Use sort data with java dispatcher

3 files changed, 39 insertions, 20 deletions

diff --git a/container-search/src/main/java/com/yahoo/fs4/DocumentInfo.java b/container-search/src/main/java/com/yahoo/fs4/DocumentInfo.java
index e5ab00fb139..dd5a4fff584 100644
--- a/container-search/src/main/java/com/yahoo/fs4/DocumentInfo.java
+++ b/container-search/src/main/java/com/yahoo/fs4/DocumentInfo.java
@@ -17,14 +17,20 @@ public class DocumentInfo implements Cloneable {
     private final double metric;
     private final int partId;
     private final int distributionKey;
+    private final byte[] sortData;
 
     DocumentInfo(ByteBuffer buffer, QueryResultPacket owner) {
+        this(buffer, owner, null);
+    }
+
+    DocumentInfo(ByteBuffer buffer, QueryResultPacket owner, byte[] sortData) {
         byte[] rawGid = new byte[GlobalId.LENGTH];
         buffer.get(rawGid);
         globalId = new GlobalId(rawGid);
         metric = decodeMetric(buffer);
         partId = owner.getMldFeature() ? buffer.getInt() : 0;
         distributionKey = owner.getMldFeature() ? buffer.getInt() : 0;
+        this.sortData = sortData;
     }
 
     public DocumentInfo(GlobalId globalId, int metric, int partId, int distributionKey) {
@@ -32,6 +38,7 @@ public class DocumentInfo implements Cloneable {
         this.metric = metric;
         this.partId = partId;
         this.distributionKey = distributionKey;
+        this.sortData = null;
     }
 
     private double decodeMetric(ByteBuffer buffer) {
@@ -49,6 +56,10 @@ public class DocumentInfo implements Cloneable {
     /** Unique key for the node this document resides on */
     public int getDistributionKey() { return distributionKey; }
 
+    public byte[] getSortData() {
+        return sortData;
+    }
+
     public String toString() {
         return "document info [globalId=" + globalId + ", metric=" + metric + "]";
     }
diff --git a/container-search/src/main/java/com/yahoo/fs4/QueryPacket.java b/container-search/src/main/java/com/yahoo/fs4/QueryPacket.java
index 60bd5314cff..d1d8b9a0f77 100644
--- a/container-search/src/main/java/com/yahoo/fs4/QueryPacket.java
+++ b/container-search/src/main/java/com/yahoo/fs4/QueryPacket.java
@@ -5,6 +5,7 @@ import com.yahoo.compress.IntegerCompressor;
 import com.yahoo.io.GrowableByteBuffer;
 import com.yahoo.prelude.query.Item;
 import com.yahoo.search.Query;
+import com.yahoo.search.dispatch.Dispatcher;
 import com.yahoo.search.grouping.vespa.GroupingExecutor;
 import com.yahoo.search.query.Ranking;
 import com.yahoo.searchlib.aggregation.Grouping;
@@ -176,18 +177,7 @@ public class QueryPacket extends Packet {
     private int getFlagInt() {
         int flags = getQueryFlags(query);
         queryPacketData.setQueryFlags(flags);
-
-        /*
-         * QFLAG_DROP_SORTDATA
-         * SORTDATA is a mangling of data from the attribute vectors
-         * which were used in the search which is byte comparable in
-         * such a way the comparing SORTDATA for two different hits
-         * will reproduce the order in which the data were returned when
-         * using sortspec.  For now we simply drop these. If they
-         * become necessary, QueryResultPacket must be
-         * updated to be able to read the sort data.
-         */
-        flags |= QFLAG_DROP_SORTDATA;
+        flags |= query.properties().getBoolean(Dispatcher.dispatchInternal, false) ? 0 : QFLAG_DROP_SORTDATA;
         return flags;
     }
 
diff --git a/container-search/src/main/java/com/yahoo/fs4/QueryResultPacket.java b/container-search/src/main/java/com/yahoo/fs4/QueryResultPacket.java
index 239101184ea..6a27beefb5e 100644
--- a/container-search/src/main/java/com/yahoo/fs4/QueryResultPacket.java
+++ b/container-search/src/main/java/com/yahoo/fs4/QueryResultPacket.java
@@ -7,7 +7,6 @@ import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
-
 /**
  * A query result packet (code 217). This packet can be decoded only.
  *
@@ -108,10 +107,9 @@ public class QueryResultPacket extends Packet {
             nodesReplied = buffer.getShort();
         }
 
-        if (sortData && documentCount > 0) { // sort data is not needed - skip
-            buffer.position(buffer.position() + (documentCount -1) * 4); // one sortIndex int per document
-            int sortDataLengthInBytes = buffer.getInt();
-            buffer.position(buffer.position() + sortDataLengthInBytes);
+        byte[][] documentSortData = null;
+        if (sortData && documentCount > 0) {
+            documentSortData = decodeSortData(buffer, documentCount);
         }
 
         if (groupDataFeature) {
@@ -125,7 +123,7 @@ public class QueryResultPacket extends Packet {
         soonActiveDocs = buffer.getLong();
         degradedReason = buffer.getInt();
 
-        decodeDocuments(buffer, documentCount);
+        decodeDocuments(buffer, documentCount, documentSortData);
         if (propsFeature) {
             int numMaps = buffer.getInt();
             propsArray = new FS4Properties[numMaps];
@@ -136,6 +134,25 @@ public class QueryResultPacket extends Packet {
         }
     }
 
+    private byte[][] decodeSortData(ByteBuffer buffer, int documentCount) {
+        int[] indexes = new int[documentCount];
+        indexes[0] = 0;
+        for (int i = 1; i < documentCount; i++) {
+            indexes[i] = buffer.getInt();
+        }
+
+        int sortDataLengthInBytes = buffer.getInt();
+        byte[][] ret = new byte[indexes.length][];
+
+        for (int i = 0; i < indexes.length; i++) {
+            int end = i + 1 >= indexes.length ? sortDataLengthInBytes : indexes[i + 1];
+            int len = end - indexes[i];
+            ret[i] = new byte[len];
+            buffer.get(ret[i], 0, len);
+        }
+        return ret;
+    }
+
     private Number decodeMaxRank(ByteBuffer buffer) {
         return Double.valueOf(buffer.getDouble());
     }
@@ -161,9 +178,10 @@ public class QueryResultPacket extends Packet {
         propsFeature     = (QRF_PROPERTIES & features) != 0;
     }
 
-    private void decodeDocuments(ByteBuffer buffer, int documentCount) {
+    private void decodeDocuments(ByteBuffer buffer, int documentCount, byte[][] documentSortData) {
         for (int i = 0; i < documentCount; i++) {
-            documents.add(new DocumentInfo(buffer, this));
+            byte[] sort = documentSortData == null ? null : documentSortData[i];
+            documents.add(new DocumentInfo(buffer, this, sort));
         }
     }