Code flow cleanups

4 files changed, 30 insertions, 20 deletions

diff --git a/container-search/src/main/java/com/yahoo/prelude/fastsearch/SortDataHitSorter.java b/container-search/src/main/java/com/yahoo/prelude/fastsearch/SortDataHitSorter.java
index e87c94fa3be..ee64632d40a 100644
--- a/container-search/src/main/java/com/yahoo/prelude/fastsearch/SortDataHitSorter.java
+++ b/container-search/src/main/java/com/yahoo/prelude/fastsearch/SortDataHitSorter.java
@@ -20,13 +20,13 @@ public class SortDataHitSorter {
         Collections.sort(hits, getComparator(sorting, fallbackComparator));
     }
 
-    public static boolean isSortable(Hit hit, HitGroup hitGroup) {
-        if (hitGroup.getQuery() == null) {
+    public static boolean isSortable(Hit hit, Sorting sorting) {
+        if (sorting == null) {
             return false;
         }
         if (hit instanceof FastHit) {
             var fhit = (FastHit) hit;
-            return fhit.hasSortData(hitGroup.getQuery().getRanking().getSorting());
+            return fhit.hasSortData(sorting);
         } else {
             return false;
         }
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
index 5a8102ff5c7..033b507b4a4 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
@@ -154,14 +154,15 @@ public class ProtobufSerialization {
         return convertFromResult(searchResult).toByteArray();
     }
 
-    public static Result deserializeToSearchResult(byte[] payload, Query query, VespaBackEndSearcher searcher)
+    public static Result deserializeToSearchResult(byte[] payload, Query query, VespaBackEndSearcher searcher, int partId, int distKey)
             throws InvalidProtocolBufferException {
         var protobuf = SearchProtocol.SearchReply.parseFrom(payload);
-        var result = convertToResult(query, protobuf, searcher.getDocumentDatabase(query));
+        var result = convertToResult(query, protobuf, searcher.getDocumentDatabase(query), partId, distKey, searcher.getName());
         return result;
     }
 
-    private static Result convertToResult(Query query, SearchProtocol.SearchReply protobuf, DocumentDatabase documentDatabase) {
+    private static Result convertToResult(Query query, SearchProtocol.SearchReply protobuf, DocumentDatabase documentDatabase, int partId,
+            int distKey, String source) {
         var result = new Result(query);
 
         result.setTotalHitCount(protobuf.getTotalHitCount());
@@ -188,12 +189,20 @@ public class ProtobufSerialization {
 
             hit.setRelevance(new Relevance(replyHit.getRelevance()));
             hit.setGlobalId(new GlobalId(replyHit.getGlobalId().toByteArray()));
-            hit.setSortData(replyHit.getSortData().toByteArray(), sorting);
+            if (!replyHit.getSortData().isEmpty()) {
+                hit.setSortData(replyHit.getSortData().toByteArray(), sorting);
+            }
             hit.setFillable();
             hit.setCached(false);
+            hit.setPartId(partId);
+            hit.setDistributionKey(distKey);
+            hit.setSource(source);
 
             result.hits().add(hit);
         }
+        if(sorting != null) {
+            result.hits().setSorted(true);
+        }
 
         return result;
     }
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/RpcSearchInvoker.java b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/RpcSearchInvoker.java
index 75e9b06f445..d118d1835a8 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/RpcSearchInvoker.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/RpcSearchInvoker.java
@@ -3,7 +3,6 @@ package com.yahoo.search.dispatch.rpc;
 
 import com.yahoo.compress.CompressionType;
 import com.yahoo.compress.Compressor;
-import com.yahoo.prelude.fastsearch.FastHit;
 import com.yahoo.prelude.fastsearch.VespaBackEndSearcher;
 import com.yahoo.search.Query;
 import com.yahoo.search.Result;
@@ -86,15 +85,7 @@ public class RpcSearchInvoker extends SearchInvoker implements Client.ResponseRe
         CompressionType compression = CompressionType.valueOf(protobufResponse.compression());
         byte[] payload = resourcePool.compressor().decompress(protobufResponse.compressedPayload(), compression,
                 protobufResponse.uncompressedSize());
-        var result = ProtobufSerialization.deserializeToSearchResult(payload, query, searcher);
-        result.hits().unorderedIterator().forEachRemaining(hit -> {
-            if(hit instanceof FastHit) {
-                FastHit fhit = (FastHit) hit;
-                fhit.setPartId(node.pathIndex());
-                fhit.setDistributionKey(node.key());
-            }
-            hit.setSource(getName());
-        });
+        var result = ProtobufSerialization.deserializeToSearchResult(payload, query, searcher, node.pathIndex(), node.key());
 
         return result;
     }
diff --git a/container-search/src/main/java/com/yahoo/search/result/HitGroup.java b/container-search/src/main/java/com/yahoo/search/result/HitGroup.java
index fdebe78c20b..8755d0b0897 100644
--- a/container-search/src/main/java/com/yahoo/search/result/HitGroup.java
+++ b/container-search/src/main/java/com/yahoo/search/result/HitGroup.java
@@ -11,6 +11,7 @@ import com.yahoo.processing.response.ArrayDataList;
 import com.yahoo.processing.response.DataList;
 import com.yahoo.processing.response.DefaultIncomingData;
 import com.yahoo.processing.response.IncomingData;
+import com.yahoo.search.query.Sorting;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -687,7 +688,7 @@ public class HitGroup extends Hit implements DataList<Hit>, Cloneable, Iterable<
             hit.setAddNumber(size());
         }
 
-        if (SortDataHitSorter.isSortable(hit, this)) {
+        if (SortDataHitSorter.isSortable(hit, currentSorting())) {
             hitsWithSortData++;
         }
 
@@ -739,7 +740,7 @@ public class HitGroup extends Hit implements DataList<Hit>, Cloneable, Iterable<
             errorHit = null;
         }
 
-        if (SortDataHitSorter.isSortable(hit, this)) {
+        if (SortDataHitSorter.isSortable(hit, currentSorting())) {
             hitsWithSortData--;
         }
 
@@ -758,7 +759,7 @@ public class HitGroup extends Hit implements DataList<Hit>, Cloneable, Iterable<
         if (!hit.isCached())
             notCachedCount++;
 
-        if (SortDataHitSorter.isSortable(hit, this)) {
+        if (SortDataHitSorter.isSortable(hit, currentSorting())) {
             hitsWithSortData++;
         }
     }
@@ -927,6 +928,15 @@ public class HitGroup extends Hit implements DataList<Hit>, Cloneable, Iterable<
         return Iterables.filter(hits, isFillable);
     }
 
+    private Sorting currentSorting() {
+        var query = getQuery();
+        if(query == null) {
+            return null;
+        } else {
+            return query.getRanking().getSorting();
+        }
+    }
+
     /** Returns the incoming hit buffer to which new hits can be added to this asynchronous, if supported by the instance */
     @Override
     public IncomingData<Hit> incoming() { return incomingHits; }