Use LongHashSet for sparse ordinal sets in DocValuesRangeIterator

lucene/CHANGES.txt

@@ -352,6 +352,9 @@ Optimizations
 * GITHUB#16370: Sort two-phase iterators by matchCost in DenseConjunctionBulkScorer.
   (Alan Woodward)
 
+* GITHUB#16297: Use LongHashSet for sparse ordinal sets in DocValuesRangeIterator so memory
+  scales with matching terms instead of segment cardinality. (Costin Leau)
+
 Bug Fixes
 ---------------------
 * GITHUB#16350: Disable bulk-scoring in monitor queries. (Alan Woodward)

lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/OrdinalSetLookupBenchmark.java

@@ -0,0 +1,114 @@
+/*
+ * 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.lucene.benchmark.jmh;
+
+import java.util.Random;
+import java.util.concurrent.TimeUnit;
+import org.apache.lucene.internal.hppc.LongHashSet;
+import org.apache.lucene.util.LongBitSet;
+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.OperationsPerInvocation;
+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.Warmup;
+
+/**
+ * Measures per-lookup cost of LongBitSet.get() vs LongHashSet.contains() for ordinal membership
+ * testing, as used by DocValuesRangeIterator.forOrdinalSet() on the per-doc hot path.
+ *
+ * <p>The ordinal space (ordCount) represents total unique values in a segment. The match count
+ * represents how many ordinals a MultiTermQuery actually matches. Lookups simulate the per-doc
+ * check where the doc's ordinal is tested against the set.
+ */
+@State(Scope.Thread)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@Warmup(iterations = 3, time = 3)
+@Measurement(iterations = 5, time = 5)
+@Fork(value = 1, warmups = 1)
+public class OrdinalSetLookupBenchmark {
+
+  private static final int LOOKUP_COUNT = 4096;
+
+  @Param({"1000", "100000", "10000000"})
+  int ordCount;
+
+  @Param({"10", "100", "1000"})
+  int matchCount;
+
+  private LongBitSet bitSet;
+  private LongHashSet hashSet;
+  private long[] lookups;
+
+  @Setup(Level.Trial)
+  public void setup() {
+    Random rng = new Random(42);
+
+    long[] matchingOrds = new long[matchCount];
+    for (int i = 0; i < matchCount; i++) {
+      matchingOrds[i] = (long) (rng.nextDouble() * ordCount);
+    }
+
+    bitSet = new LongBitSet(ordCount);
+    hashSet = new LongHashSet(matchCount);
+    for (long ord : matchingOrds) {
+      bitSet.set(ord);
+      hashSet.add(ord);
+    }
+
+    lookups = new long[LOOKUP_COUNT];
+    for (int i = 0; i < LOOKUP_COUNT; i++) {
+      if (rng.nextInt(4) == 0 && matchCount > 0) {
+        lookups[i] = matchingOrds[rng.nextInt(matchCount)];
+      } else {
+        lookups[i] = (long) (rng.nextDouble() * ordCount);
+      }
+    }
+  }
+
+  @Benchmark
+  @OperationsPerInvocation(LOOKUP_COUNT)
+  public int bitSetLookup() {
+    int hits = 0;
+    for (long ord : lookups) {
+      if (bitSet.get(ord)) {
+        hits++;
+      }
+    }
+    return hits;
+  }
+
+  @Benchmark
+  @OperationsPerInvocation(LOOKUP_COUNT)
+  public int hashSetLookup() {
+    int hits = 0;
+    for (long ord : lookups) {
+      if (hashSet.contains(ord)) {
+        hits++;
+      }
+    }
+    return hits;
+  }
+}

lucene/core/src/java/org/apache/lucene/search/DocValuesRangeIterator.java

@@ -17,12 +17,15 @@
 package org.apache.lucene.search;
 
 import java.io.IOException;
+import java.util.function.LongPredicate;
 import org.apache.lucene.index.DocValuesSkipper;
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.SortedDocValues;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.index.SortedSetDocValues;
 import org.apache.lucene.index.TermsEnum;
+import org.apache.lucene.internal.hppc.LongHashSet;
+import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.FixedBitSet;
 import org.apache.lucene.util.IOBooleanSupplier;
 import org.apache.lucene.util.LongBitSet;
@@ -135,7 +138,7 @@ public static DocValuesRangeIterator forOrdinalRange(
    *     place of the per-doc bit lookup. Computed at construction time so callers don't have to
    *     trust an undocumented invariant about the bounds of set bits in {@code ords}.
    */
-  private record OrdinalSet(long min, long max, LongBitSet ords, boolean contiguous) {
+  private record OrdinalSet(long min, long max, LongPredicate ords, boolean contiguous) {
 
     boolean disjoint(DocValuesSkipper skipper) {
       if (skipper == null) {
@@ -145,27 +148,37 @@ boolean disjoint(DocValuesSkipper skipper) {
     }
   }
 
-  private static OrdinalSet buildOrdinalSet(TermsEnum termsEnum, long ordCount) throws IOException {
+  private static OrdinalSet buildOrdinalSet(TermsEnum termsEnum) throws IOException {
     if (termsEnum.next() == null) {
       return null;
     }
-    // TODO can we be more memory efficient here? eg LongHashSet
-    LongBitSet ords = new LongBitSet(ordCount);
     long min = termsEnum.ord();
-    ords.set(min);
     long max = min;
-    // Count distinct ords via getAndSet so a TermsEnum that yields a duplicate ord doesn't fool
-    // the contiguity check below. The first set bit (min) is always new on a fresh bitset.
+    long[] collected = new long[] {min};
+    int count = 1;
+    long prev = min;
     long distinctCount = 1;
     while (termsEnum.next() != null) {
-      max = termsEnum.ord();
-      if (ords.getAndSet(max) == false) {
+      long ord = termsEnum.ord();
+      if (ord != prev) {
         distinctCount++;
+        prev = ord;
       }
+      max = ord;
+      if (count == collected.length) {
+        collected = ArrayUtil.grow(collected);
+      }
+      collected[count++] = ord;
+    }
+    boolean contiguous = distinctCount == max - min + 1;
+    if (contiguous) {
+      return new OrdinalSet(min, max, _ -> true, true);
+    }
+    LongHashSet ords = new LongHashSet(count);
+    for (int i = 0; i < count; i++) {
+      ords.add(collected[i]);
     }
-    // If every ord in [min, max] is set, the set is equivalent to forOrdinalRange and can use the
-    // cheaper range check + block-level YES short-circuit.
-    return new OrdinalSet(min, max, ords, distinctCount == max - min + 1);
+    return new OrdinalSet(min, max, ords::contains, false);
   }
 
   /**
@@ -177,14 +190,14 @@ private static OrdinalSet buildOrdinalSet(TermsEnum termsEnum, long ordCount) th
    */
   public static DocValuesRangeIterator forOrdinalSet(
       SortedDocValues values, DocValuesSkipper skipper, TermsEnum terms) throws IOException {
-    OrdinalSet ordinalSet = buildOrdinalSet(terms, values.getValueCount());
+    OrdinalSet ordinalSet = buildOrdinalSet(terms);
     if (ordinalSet == null || ordinalSet.disjoint(skipper)) {
       return new EmptyRangeIterator();
     }
     if (ordinalSet.contiguous) {
       return forOrdinalRange(values, skipper, ordinalSet.min, ordinalSet.max);
     }
-    IOBooleanSupplier check = () -> ordinalSet.ords.get(values.ordValue());
+    IOBooleanSupplier check = () -> ordinalSet.ords.test(values.ordValue());
     return skipper == null
         ? new DocValuesValueRangeIterator(values, check, 2)
         : new DocValuesBlockRangeIterator(
@@ -200,7 +213,7 @@ public static DocValuesRangeIterator forOrdinalSet(
    */
   public static DocValuesRangeIterator forOrdinalSet(
       SortedSetDocValues values, DocValuesSkipper skipper, TermsEnum terms) throws IOException {
-    OrdinalSet ordinalSet = buildOrdinalSet(terms, values.getValueCount());
+    OrdinalSet ordinalSet = buildOrdinalSet(terms);
     return forOrdinalSet(values, skipper, ordinalSet);
   }
 
@@ -221,7 +234,7 @@ public static DocValuesRangeIterator forOrdinalSet(
       LongBitSet ords) {
     // Pass contiguous=false: callers of this overload aren't required by the javadoc to keep all
     // set bits within [minOrd, maxOrd], so we can't infer contiguity from cardinality alone.
-    return forOrdinalSet(values, skipper, new OrdinalSet(minOrd, maxOrd, ords, false));
+    return forOrdinalSet(values, skipper, new OrdinalSet(minOrd, maxOrd, ords::get, false));
   }
 
   private static DocValuesRangeIterator forOrdinalSet(
@@ -239,7 +252,7 @@ private static DocValuesRangeIterator forOrdinalSet(
             if (v > ordinalSet.max) {
               return false;
             }
-            if (v >= ordinalSet.min && ordinalSet.ords.get(v)) {
+            if (v >= ordinalSet.min && ordinalSet.ords.test(v)) {
               return true;
             }
           }