SOLR-5818: distrib search with custom comparator does not quite work correctly

git-svn-id: https://svn.apache.org/repos/asf/lucene/dev/trunk@1575344 13f79535-47bb-0310-9956-ffa450edef68

solr/CHANGES.txt

@@ -135,6 +135,9 @@ Bug Fixes
 * SOLR-5796: Increase how long we are willing to wait for a core to see the ZK
   advertised leader in it's local state. (Timothy Potter, Mark Miller)  
 
+* SOLR-5818: distrib search with custom comparator does not quite work correctly
+  (Ryan Ernst)
+
 Optimizations
 ----------------------
 * SOLR-1880: Distributed Search skips GET_FIELDS stage if EXECUTE_QUERY

solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java

@@ -25,12 +25,14 @@ import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.FieldComparator;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.Sort;
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.grouping.GroupDocs;
 import org.apache.lucene.search.grouping.SearchGroup;
 import org.apache.lucene.search.grouping.TopGroups;
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.InPlaceMergeSorter;
 import org.apache.solr.client.solrj.SolrServerException;
 import org.apache.solr.common.SolrDocument;
 import org.apache.solr.common.SolrDocumentList;
@@ -500,12 +502,32 @@ public class QueryComponent extends SearchComponent
 
       // sort ids from lowest to highest so we can access them in order
       int nDocs = docList.size();
-      long[] sortedIds = new long[nDocs];
+      final long[] sortedIds = new long[nDocs];
-      DocIterator it = rb.getResults().docList.iterator();
+      final float[] scores = new float[nDocs]; // doc scores, parallel to sortedIds
+      DocList docs = rb.getResults().docList;
+      DocIterator it = docs.iterator();
       for (int i=0; i<nDocs; i++) {
         sortedIds[i] = (((long)it.nextDoc()) << 32) | i;
+        scores[i] = docs.hasScores() ? it.score() : Float.NaN;
       }
-      Arrays.sort(sortedIds);
+
+      // sort ids and scores together
+      new InPlaceMergeSorter() {
+        @Override
+        protected void swap(int i, int j) {
+          long tmpId = sortedIds[i];
+          float tmpScore = scores[i];
+          sortedIds[i] = sortedIds[j];
+          scores[i] = scores[j];
+          sortedIds[j] = tmpId;
+          scores[j] = tmpScore;
+        }
+
+        @Override
+        protected int compare(int i, int j) {
+          return Long.compare(sortedIds[i], sortedIds[j]);
+        }
+      }.sort(0, sortedIds.length);
 
       SortSpec sortSpec = rb.getSortSpec();
       Sort sort = searcher.weightSort(sortSpec.getSort());
@@ -527,7 +549,9 @@ public class QueryComponent extends SearchComponent
         int lastIdx = -1;
         int idx = 0;
 
-        for (long idAndPos : sortedIds) {
+        for (int i = 0; i < sortedIds.length; ++i) {
+          long idAndPos = sortedIds[i];
+          float score = scores[i];
           int doc = (int)(idAndPos >>> 32);
           int position = (int)idAndPos;
 
@@ -546,6 +570,7 @@ public class QueryComponent extends SearchComponent
           }
 
           doc -= currentLeaf.docBase;  // adjust for what segment this is in
+          comparator.setScorer(new FakeScorer(doc, score));
           comparator.copy(0, doc);
           Object val = comparator.value(0);
           if (null != ft) val = ft.marshalSortValue(val); 
@@ -1157,4 +1182,50 @@ public class QueryComponent extends SearchComponent
   public URL[] getDocs() {
     return null;
   }
+
+  /**
+   * Fake scorer for a single document
+   *
+   * TODO: when SOLR-5595 is fixed, this wont be needed, as we dont need to recompute sort values here from the comparator
+   */
+  private static class FakeScorer extends Scorer {
+    final int docid;
+    final float score;
+
+    FakeScorer(int docid, float score) {
+      super(null);
+      this.docid = docid;
+      this.score = score;
+    }
+
+    @Override
+    public int docID() {
+      return docid;
+    }
+
+    @Override
+    public float score() throws IOException {
+      return score;
+    }
+
+    @Override
+    public int freq() throws IOException {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public int nextDoc() throws IOException {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public int advance(int target) throws IOException {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public long cost() {
+      return 1;
+    }
+  }
 }

solr/core/src/test-files/solr/collection1/conf/schema-field-sort-values.xml

@@ -0,0 +1,41 @@
+<?xml version="1.0" ?>
+<!--
+ Licensed to the Apache Software Foundation (ASF) under one or more
+ contributor license agreements.  See the NOTICE file distributed with
+ this work for additional information regarding copyright ownership.
+ The ASF licenses this file to You under the Apache License, Version 2.0
+ (the "License"); you may not use this file except in compliance with
+ the License.  You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+-->
+
+<schema name="test-custom-comparator" version="1.5">
+  <types>
+    <fieldType name="int" class="solr.TrieIntField" precisionStep="0" omitNorms="true" positionIncrementGap="0"/>
+    <fieldType name="long" class="solr.TrieLongField" precisionStep="0" omitNorms="true" positionIncrementGap="0"/>
+    <fieldtype name="string" class="solr.StrField" sortMissingLast="true"/>
+    <fieldtype name="text" class="solr.TextField">
+      <analyzer>
+        <tokenizer class="solr.StandardTokenizerFactory"/>
+      </analyzer>
+    </fieldtype>
+    <fieldType class="org.apache.solr.schema.WrappedIntField" name="wrapped_int"/>
+  </types>
+  <fields>
+    <field name="id" type="int" indexed="true" stored="true" multiValued="false" required="false"/>
+    <field name="_version_" type="long" indexed="true" stored="true" multiValued="false"/>
+    <field name="text" type="text" indexed="true" stored="false"/>
+    <field name="payload" type="wrapped_int" indexed="false"
+           stored="true"  multiValued="false"    docValues="true" required="true"/>
+
+  </fields>
+  <defaultSearchField>text</defaultSearchField>
+  <uniqueKey>id</uniqueKey>
+</schema>

solr/core/src/test/org/apache/solr/schema/WrappedIntField.java

@@ -0,0 +1,46 @@
+package org.apache.solr.schema;
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import org.apache.lucene.expressions.Expression;
+import org.apache.lucene.expressions.SimpleBindings;
+import org.apache.lucene.expressions.js.JavascriptCompiler;
+import org.apache.lucene.search.SortField;
+
+/**
+ * Custom field wrapping an int, to test sorting via a custom comparator.
+ */
+public class WrappedIntField extends TrieIntField {
+  Expression expr;
+
+  public WrappedIntField() {
+    try {
+      expr = JavascriptCompiler.compile("payload % 3");
+    } catch (Exception e) {
+      throw new RuntimeException("impossible?", e);
+    }
+  }
+
+  @Override
+  public SortField getSortField(final SchemaField field, final boolean reverse) {
+    field.checkSortability();
+    SimpleBindings bindings = new SimpleBindings();
+    bindings.add(super.getSortField(field, reverse));
+    return expr.getSortField(bindings, reverse);
+  }
+}

solr/core/src/test/org/apache/solr/search/TestFieldSortValues.java

@@ -0,0 +1,53 @@
+package org.apache.solr.search;
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import org.apache.lucene.util.LuceneTestCase.SuppressCodecs;
+import org.apache.solr.SolrTestCaseJ4;
+import org.junit.BeforeClass;
+
+
+/**
+ * Test QueryComponent.doFieldSortValues
+ */
+@SuppressCodecs({"Lucene3x"})
+public class TestFieldSortValues extends SolrTestCaseJ4 {
+
+  @BeforeClass
+  public static void beforeClass() throws Exception {
+    initCore("solrconfig-minimal.xml", "schema-field-sort-values.xml");
+  }
+  
+  public void testCustomComparator() throws Exception {
+    clearIndex();
+    assertU(adoc(sdoc("id", "1", "payload", "2")));
+    assertU(adoc(sdoc("id", "2", "payload", "3")));
+    assertU(adoc(sdoc("id", "3", "payload", "1")));
+    assertU(adoc(sdoc("id", "4", "payload", "5")));
+    assertU(adoc(sdoc("id", "5", "payload", "4")));
+    assertU(commit());
+
+    // payload is backed by a custom sort field which returns the payload value mod 3
+    assertQ(req("q", "*:*", "fl", "id", "sort", "payload asc, id asc", "fsv", "true")
+        , "//result/doc[int='2'  and position()=1]"
+        , "//result/doc[int='3'  and position()=2]"
+        , "//result/doc[int='5'  and position()=3]"
+        , "//result/doc[int='1'  and position()=4]"
+        , "//result/doc[int='4'  and position()=5]");
+  }
+}