SOLR-18304: Optimize collapse performance for String fields in Solr 9.x and later

changelog/unreleased/SOLR-18304-fix-collapse-on-string-performance.yml

@@ -0,0 +1,7 @@
+title: Optimize collapse performance for String fields in Solr 9.x and later
+type: changed
+authors:
+  - name: Bartosz Fidrysiak
+links:
+  - name: SOLR-18304
+    url: https://issues.apache.org/jira/browse/SOLR-18304

solr/benchmark/src/java/org/apache/solr/bench/search/CollapsingSearch.java

@@ -0,0 +1,214 @@
+/*
+ * 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.
+ */
+package org.apache.solr.bench.search;
+
+import java.io.IOException;
+import java.time.Instant;
+import java.time.format.DateTimeFormatter;
+import java.util.Locale;
+import java.util.Random;
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+import org.apache.solr.bench.BaseBenchState;
+import org.apache.solr.bench.SolrBenchState;
+import org.apache.solr.client.solrj.SolrServerException;
+import org.apache.solr.client.solrj.request.CollectionAdminRequest;
+import org.apache.solr.client.solrj.request.QueryRequest;
+import org.apache.solr.client.solrj.request.SolrQuery;
+import org.apache.solr.client.solrj.request.UpdateRequest;
+import org.apache.solr.common.SolrInputDocument;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+
+@Fork(value = 1)
+@Warmup(time = 5, iterations = 5)
+@Measurement(time = 15, iterations = 5)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@Threads(value = 1)
+public class CollapsingSearch {
+
+  static final String COLLECTION = "collapse_bench";
+
+  @State(Scope.Benchmark)
+  public static class BenchState {
+
+    @Param({"2000000"})
+    int numDocs;
+
+    @Param({"100000"})
+    int numGroups;
+
+    @Param({"1", "10"})
+    int numSegments;
+
+    final QueryRequest qSimple =
+        new QueryRequest(new SolrQuery("q", "*:*", "rows", "10", "cache", "false"));
+
+    final QueryRequest qCollapseWithoutSort =
+        new QueryRequest(
+            new SolrQuery(
+                "q", "*:*",
+                "fq", "{!collapse field=group_s_dv cache=false}",
+                "rows", "10",
+                "cache", "false"));
+
+    final QueryRequest qCollapseByStr =
+        new QueryRequest(
+            new SolrQuery(
+                "q", "*:*",
+                "fq", "{!collapse field=group_s_dv sort='str_s_dv asc' cache=false}",
+                "rows", "10",
+                "cache", "false"));
+
+    final QueryRequest qCollapseByDate =
+        new QueryRequest(
+            new SolrQuery(
+                "q", "*:*",
+                "fq", "{!collapse field=group_s_dv sort='date_dt_dv asc' cache=false}",
+                "rows", "10",
+                "cache", "false"));
+
+    final QueryRequest qCollapseByLong =
+        new QueryRequest(
+            new SolrQuery(
+                "q", "*:*",
+                "fq", "{!collapse field=group_s_dv sort='long_l_dv asc' cache=false}",
+                "rows", "10",
+                "cache", "false"));
+
+    final QueryRequest qCollapseByDateAndStr =
+        new QueryRequest(
+            new SolrQuery(
+                "q", "*:*",
+                "fq",
+                    "{!collapse field=group_s_dv sort='date_dt_dv asc, str_s_dv asc' cache=false}",
+                "rows", "10",
+                "cache", "false"));
+
+    @Setup(Level.Trial)
+    public void setupTrial(SolrBenchState miniClusterState) throws Exception {
+      System.setProperty("commitwithin.softcommit", "false");
+      miniClusterState.startSolr(1);
+      miniClusterState.createCollection(COLLECTION, 1, 1);
+
+      indexDocs(miniClusterState);
+      miniClusterState.forceMerge(COLLECTION, numSegments);
+
+      BaseBenchState.log(
+          "CollapsingSearch ready: numDocs="
+              + numDocs
+              + " numGroups="
+              + numGroups
+              + " numSegments(requested)="
+              + numSegments);
+    }
+
+    @Setup(Level.Iteration)
+    public void setupIteration(SolrBenchState miniClusterState)
+        throws SolrServerException, IOException {
+      CollectionAdminRequest.Reload reload = CollectionAdminRequest.reloadCollection(COLLECTION);
+      miniClusterState.client.request(reload);
+    }
+
+    @TearDown(Level.Trial)
+    public void teardown(SolrBenchState miniClusterState) throws Exception {
+      CollectionAdminRequest.deleteCollection(COLLECTION).process(miniClusterState.client);
+    }
+
+    private int pickGroup(int docId) {
+      return docId % numGroups;
+    }
+
+    private void indexDocs(SolrBenchState state) throws Exception {
+      int docsPerSegment = numDocs / numSegments;
+      String[] groupIds = new String[numGroups];
+      for (int i = 0; i < numGroups; i++) {
+        groupIds[i] = String.format(Locale.ROOT, "group_%06d", i);
+      }
+      Random rng = new Random(42);
+      Instant baseDate = Instant.parse("2020-01-01T00:00:00Z");
+
+      int docId = 0;
+      for (int seg = 0; seg < numSegments; seg++) {
+        int segDocs = (seg < numSegments - 1) ? docsPerSegment : (numDocs - docId);
+        UpdateRequest req = new UpdateRequest();
+        for (int i = 0; i < segDocs; i++) {
+          SolrInputDocument doc = new SolrInputDocument();
+          doc.addField("id", String.valueOf(docId));
+          doc.addField("group_s_dv", groupIds[pickGroup(docId)]);
+          docId++;
+          doc.addField("str_s_dv", UUID.randomUUID().toString());
+          doc.addField(
+              "date_dt_dv",
+              DateTimeFormatter.ISO_INSTANT.format(baseDate.plusSeconds(rng.nextInt(10_000_000))));
+          doc.addField("long_l_dv", rng.nextInt(10_000_000));
+          req.add(doc);
+        }
+        req.commit(state.client, COLLECTION);
+      }
+    }
+  }
+
+  @Benchmark
+  public Object simple(BenchState b, SolrBenchState cluster)
+      throws SolrServerException, IOException {
+    return cluster.client.request(b.qSimple, COLLECTION);
+  }
+
+  @Benchmark
+  public Object collapseWithoutSort(BenchState b, SolrBenchState cluster)
+      throws SolrServerException, IOException {
+    return cluster.client.request(b.qCollapseWithoutSort, COLLECTION);
+  }
+
+  @Benchmark
+  public Object collapseByStr(BenchState b, SolrBenchState cluster)
+      throws SolrServerException, IOException {
+    return cluster.client.request(b.qCollapseByStr, COLLECTION);
+  }
+
+  @Benchmark
+  public Object collapseByDate(BenchState b, SolrBenchState cluster)
+      throws SolrServerException, IOException {
+    return cluster.client.request(b.qCollapseByDate, COLLECTION);
+  }
+
+  @Benchmark
+  public Object collapseByLong(BenchState b, SolrBenchState cluster)
+      throws SolrServerException, IOException {
+    return cluster.client.request(b.qCollapseByLong, COLLECTION);
+  }
+
+  @Benchmark
+  public Object collapseByDateAndStr(BenchState b, SolrBenchState cluster)
+      throws SolrServerException, IOException {
+    return cluster.client.request(b.qCollapseByDateAndStr, COLLECTION);
+  }
+}

solr/benchmark/src/resources/configs/cloud-minimal/conf/schema.xml

@@ -56,6 +56,7 @@
     <dynamicField name="*_d_dv" type="double" indexed="true" docValues="true" stored="false"/>
     <dynamicField name="*_dt" type="date" indexed="true" stored="false"/>
     <dynamicField name="*_dt_dv" type="date" indexed="true" docValues="true" stored="false"/>
+    <dynamicField name="*_s_dv" type="string" indexed="true" docValues="true" stored="false"/>
 
     <uniqueKey>id</uniqueKey>
 </schema>

solr/core/src/java/org/apache/solr/search/CollapsingQParserPlugin.java

@@ -3565,6 +3565,9 @@ private static class SortFieldsCompare {
     private Object[][] groupHeadValues; // growable
     private final Object[] nullGroupValues;
 
+    private final SortedDocValues[] stringSortDVs;
+    private final int[] stringMissingOrd;
+
     /**
      * Constructs an instance based on the (raw, un-rewritten) SortFields to be used, and an initial
      * number of expected groups (will grow as needed).
@@ -3576,6 +3579,7 @@ public SortFieldsCompare(SortField[] sorts, int initNumGroups) {
       fieldComparators = new FieldComparator[numClauses];
       leafFieldComparators = new LeafFieldComparator[numClauses];
       reverseMul = new int[numClauses];
+      stringMissingOrd = new int[numClauses];
       for (int clause = 0; clause < numClauses; clause++) {
         SortField sf = sorts[clause];
         // we only need one slot for every comparator
@@ -3587,14 +3591,28 @@ public SortFieldsCompare(SortField[] sorts, int initNumGroups) {
                     : Pruning.NONE);
 
         reverseMul[clause] = sf.getReverse() ? -1 : 1;
+        if (sf.getType() == SortField.Type.STRING) {
+          stringMissingOrd[clause] =
+              (sf.getMissingValue() == SortField.STRING_LAST) ? Integer.MAX_VALUE : -1;
+        }
       }
       groupHeadValues = new Object[initNumGroups][];
       nullGroupValues = new Object[numClauses];
+      stringSortDVs = new SortedDocValues[numClauses]; // populated in setNextReader
     }
 
     public void setNextReader(LeafReaderContext context) throws IOException {
       for (int clause = 0; clause < numClauses; clause++) {
         leafFieldComparators[clause] = fieldComparators[clause].getLeafComparator(context);
+        if (sorts[clause].getType() == SortField.Type.STRING) {
+          String field = sorts[clause].getField();
+          FieldInfo fi = context.reader().getFieldInfos().fieldInfo(field);
+          if (fi != null && fi.getDocValuesType() == DocValuesType.SORTED) {
+            stringSortDVs[clause] = DocValues.getSorted(context.reader(), field);
+          } else {
+            stringSortDVs[clause] = null;
+          }
+        }
       }
     }
 
@@ -3652,12 +3670,20 @@ public void setNullGroupValues(int contextDoc) throws IOException {
 
     /**
      * Records the SortField values for the specified contextDoc into the values array provided by
-     * the caller.
+     * the caller. STRING clauses with SORTED DocValues are stored as {@link LazyStringValue} to
+     * defer {@code lookupOrd()} until a second document actually competes for the group.
      */
     private void setGroupValues(Object[] values, int contextDoc) throws IOException {
       for (int clause = 0; clause < numClauses; clause++) {
-        leafFieldComparators[clause].copy(0, contextDoc);
-        values[clause] = cloneIfBytesRef(fieldComparators[clause].value(0));
+        if (stringSortDVs[clause] != null) {
+          SortedDocValues dv = stringSortDVs[clause];
+          int missingOrd = stringMissingOrd[clause];
+          int ord = dv.advanceExact(contextDoc) ? dv.ordValue() : missingOrd;
+          values[clause] = new LazyStringValue(dv, ord, missingOrd);
+        } else {
+          leafFieldComparators[clause].copy(0, contextDoc);
+          values[clause] = cloneIfBytesRef(fieldComparators[clause].value(0));
+        }
       }
     }
 
@@ -3696,6 +3722,20 @@ private boolean testAndSetGroupValues(Object[] values, int contextDoc) throws IO
       int testClause = 0;
       for (
       /* testClause */ ; testClause < numClauses; testClause++) {
+        if (values[testClause] instanceof LazyStringValue headVal) {
+          SortedDocValues segDV = stringSortDVs[testClause];
+          if (segDV != null && segDV == headVal.dv) {
+            int missingOrd = headVal.missingOrd;
+            int candidateOrd = segDV.advanceExact(contextDoc) ? segDV.ordValue() : missingOrd;
+            lastCompare = reverseMul[testClause] * Integer.compare(candidateOrd, headVal.ord);
+            stash[testClause] = new LazyStringValue(segDV, candidateOrd, missingOrd);
+            if (0 != lastCompare) break;
+            continue;
+          }
+          var materializedHeadVal = headVal.materialize();
+          values[testClause] =
+              materializedHeadVal != null ? BytesRef.deepCopyOf(materializedHeadVal) : null;
+        }
         leafFieldComparators[testClause].copy(0, contextDoc);
         FieldComparator fcomp = fieldComparators[testClause];
         stash[testClause] = cloneIfBytesRef(fcomp.value(0));
@@ -3719,6 +3759,13 @@ private boolean testAndSetGroupValues(Object[] values, int contextDoc) throws IO
       System.arraycopy(stash, 0, values, 0, testClause);
       // read the remaining values we didn't need to test
       for (int copyClause = testClause; copyClause < numClauses; copyClause++) {
+        SortedDocValues segDV = stringSortDVs[copyClause];
+        if (segDV != null) {
+          int missingOrd = stringMissingOrd[copyClause];
+          int candidateOrd = segDV.advanceExact(contextDoc) ? segDV.ordValue() : missingOrd;
+          values[copyClause] = new LazyStringValue(segDV, candidateOrd, missingOrd);
+          continue;
+        }
         leafFieldComparators[copyClause].copy(0, contextDoc);
         values[copyClause] = cloneIfBytesRef(fieldComparators[copyClause].value(0));
       }

solr/core/src/java/org/apache/solr/search/LazyStringValue.java

@@ -0,0 +1,41 @@
+/*
+ * 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.
+ */
+package org.apache.solr.search;
+
+import java.io.IOException;
+import org.apache.lucene.index.SortedDocValues;
+import org.apache.lucene.util.BytesRef;
+
+final class LazyStringValue {
+
+  final SortedDocValues dv;
+  final int ord;
+  final int missingOrd;
+
+  LazyStringValue(SortedDocValues dv, int ord, int missingOrd) {
+    this.dv = dv;
+    this.ord = ord;
+    this.missingOrd = missingOrd;
+  }
+
+  BytesRef materialize() throws IOException {
+    if (ord == missingOrd || ord < 0) {
+      return null;
+    }
+    return dv.lookupOrd(ord);
+  }
+}