You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by jp...@apache.org on 2025/09/12 14:57:34 UTC
(lucene) branch branch_10x updated: Help collectors take advantage of bulk-retrieval of doc values. (#15173)
This is an automated email from the ASF dual-hosted git repository.
jpountz pushed a commit to branch branch_10x
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/branch_10x by this push:
new 815b1cea1b6 Help collectors take advantage of bulk-retrieval of doc values. (#15173)
815b1cea1b6 is described below
commit 815b1cea1b6071e749b3636451b60519812b95b3
Author: Adrien Grand <jp...@gmail.com>
AuthorDate: Fri Sep 12 16:14:05 2025 +0200
Help collectors take advantage of bulk-retrieval of doc values. (#15173)
We added `NumericDocValues#longValues` to retrieve multiple values in a single
call and better amortize the cost of virtual function calls in #15149. But it's
not easy today to take advantage of this API in collectors, the best option is
to buffer doc IDs in `LeafCollector#collect`, but then it's not great since
this `LeafCollector#collect` call itself may be virtual, so we may still be
doing 1+ virtual call per collected doc.
This PR adds `DocIdStream#intoArray` to help collectors collect batches of doc
IDs at once and retrieve numeric doc values for them. This way, you can compute
facets/aggregations with less than one virtual call per collected document.
I benchmarked on the geonames dataset, computing the average of the elevation
field with the following collector:
```java
class AverageCollector implements Collector {
long sum;
int count;
@Override
public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
SortedNumericDocValues sortedValues = context.reader().getSortedNumericDocValues("elevation");
if (sortedValues == null) {
throw new CollectionTerminatedException();
}
NumericDocValues values = DocValues.unwrapSingleton(sortedValues);
if (values == null) { // the field is single-valued
throw new Error();
}
return new LeafCollector() {
int[] docBuffer = new int[64];
long[] valueBuffer = new long[docBuffer.length];
@Override
public void setScorer(Scorable scorer) throws IOException {
}
@Override
public void collect(DocIdStream stream) throws IOException {
for (int count = stream.intoArray(docBuffer); count != 0; count = stream.intoArray(docBuffer)) {
values.longValues(count, docBuffer, valueBuffer, Long.MIN_VALUE);
for (int i = 0; i < count; ++i) {
long v = valueBuffer[i];
sum += v;
if (v != Long.MIN_VALUE) {
AverageCollector.this.count++;
}
}
}
}
@Override
public void collect(int doc) throws IOException {
if (values.advanceExact(doc)) {
sum += values.longValue();
count += 1;
}
}
};
}
@Override
public ScoreMode scoreMode() {
return ScoreMode.COMPLETE_NO_SCORES;
}
}
```
The benchmark runs a few queries with different collectors first, so that
`LeafCollector#collect` is polymorphic. I get the following numbers:
| Query | | Latency without DocIdStream#intoArray (ms) | Latency with DocIdStream#intoArray (ms) |
| - | - | - | - |
| `MatchAllDocsQuery` | Uses `RangeDocIdStream` under the hood | 81 | 65 |
| `featureClass:(S P)` | Matches spots or cities (5M docs), uses `BitSetDocIdStream` under the hood | 54 | 43 |
Note that `NumericDocValues#advanceExact` and `NumericDocValues#longValue` are
not polymorphic in this benchmark since all segments use the same impl on the
`elevation` field. The difference would be bigger if they were.
My profiler suggests that retrieving doc values is the bottleneck when
computing the average on a `MatchAllDocsQuery`, while the bottleneck is a mix
of retrieving doc values and extracting set bits from the bit set when
computing the average on `featureClass:(S P)`.
---
lucene/CHANGES.txt | 4 +
.../backward_codecs/lucene103/BitSetUtil.java | 104 ---------------------
.../lucene103/Lucene103PostingsReader.java | 8 +-
.../backward_codecs/lucene103/TestBitSetUtil.java | 53 -----------
.../apache/lucene/codecs/lucene104/BitSetUtil.java | 104 ---------------------
.../codecs/lucene104/Lucene104PostingsReader.java | 8 +-
.../apache/lucene/search/BitSetDocIdStream.java | 14 +++
.../lucene/search/BooleanScorerSupplier.java | 5 +
.../java/org/apache/lucene/search/DocIdStream.java | 16 ++++
.../org/apache/lucene/search/LRUQueryCache.java | 29 ++++++
.../org/apache/lucene/search/RangeDocIdStream.java | 16 ++++
.../java/org/apache/lucene/util/FixedBitSet.java | 78 ++++++++++++++++
.../lucene/codecs/lucene104/TestBitSetUtil.java | 53 -----------
.../lucene/search/TestBitSetDocIdStream.java | 53 +++++++++++
.../apache/lucene/search/TestRangeDocIdStream.java | 42 +++++++++
.../org/apache/lucene/util/TestFixedBitSet.java | 44 +++++++++
.../facet/plain/histograms/HistogramCollector.java | 38 +++++++-
.../tests/search/AssertingLeafCollector.java | 58 +++++++++++-
18 files changed, 402 insertions(+), 325 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 782f9a772e8..4d41f1a28d2 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -10,6 +10,10 @@ API Changes
* GITHUB#15149: Introduce NumericDocValues#longValues to help speed up the
retrieval of many doc values at once. (Adrien Grand)
+* GITHUB#15173: Introduce DocIdStream#intoArray to help collect batches of doc
+ IDs at once, in particular by taking advantage of the new
+ NumericDocValues#longValues API. (Adrien Grand)
+
New Features
---------------------
* GITHUB#14565: Add ParentsChildrenBlockJoinQuery that supports parent and child filter in the same query
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene103/BitSetUtil.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene103/BitSetUtil.java
deleted file mode 100644
index fcbeca1ae69..00000000000
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene103/BitSetUtil.java
+++ /dev/null
@@ -1,104 +0,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.
- */
-package org.apache.lucene.backward_codecs.lucene103;
-
-import java.util.Objects;
-import org.apache.lucene.util.FixedBitSet;
-
-class BitSetUtil {
-
- /**
- * Converts set bits in the given bitset to an array of document IDs. Only processes bits from
- * index {@code from} (inclusive) to {@code to} (exclusive) and returns the number of bits set in
- * this range.
- *
- * <p>Each set bit's position is converted to a document ID by adding the {@code base} value and
- * stored in the provided {@code array}.
- *
- * <p>NOTE: Caller need to ensure the {@code array} has a length greater than or equal to {@code
- * bitSet.cardinality(from, to) + 1}.
- */
- static int denseBitsetToArray(FixedBitSet bitSet, int from, int to, int base, int[] array) {
- assert bitSet.cardinality(from, to) + 1 <= array.length
- : "Array length must be at least bitSet.cardinality(from, to) + 1";
-
- Objects.checkFromToIndex(from, to, bitSet.length());
-
- int offset = 0;
- long[] bits = bitSet.getBits();
- // First, align `from` with a word start, ie. a multiple of Long.SIZE (64)
- if ((from & 0x3F) != 0) {
- long word = bits[from >> 6] >>> from;
- int numBitsTilNextWord = -from & 0x3F;
- if (to - from < numBitsTilNextWord) {
- // All bits are in a single word
- word &= (1L << (to - from)) - 1L;
- return word2Array(word, from + base, array, offset);
- }
- offset = word2Array(word, from + base, array, offset);
- from += numBitsTilNextWord;
- assert (from & 0x3F) == 0;
- }
-
- for (int i = from >> 6, end = to >> 6; i < end; ++i) {
- long word = bits[i];
- offset = word2Array(word, base + (i << 6), array, offset);
- }
-
- // Now handle remaining bits in the last partial word
- if ((to & 0x3F) != 0) {
- long word = bits[to >> 6] & ((1L << to) - 1);
- offset = word2Array(word, base + (to & ~0x3F), array, offset);
- }
-
- return offset;
- }
-
- private static int word2Array(long word, int base, int[] docs, int offset) {
- final int bitCount = Long.bitCount(word);
-
- if (bitCount >= 32) {
- return denseWord2Array(word, base, docs, offset);
- }
-
- for (int i = 0; i < bitCount; i++) {
- int ntz = Long.numberOfTrailingZeros(word);
- docs[offset++] = base + ntz;
- word ^= 1L << ntz;
- }
-
- return offset;
- }
-
- private static int denseWord2Array(long word, int base, int[] docs, int offset) {
- final int lWord = (int) word;
- final int hWord = (int) (word >>> 32);
- final int offset32 = offset + Integer.bitCount(lWord);
- int hOffset = offset32;
-
- for (int i = 0; i < 32; i++) {
- docs[offset] = base + i;
- docs[hOffset] = base + i + 32;
- offset += (lWord >>> i) & 1;
- hOffset += (hWord >>> i) & 1;
- }
-
- docs[offset32] = base + 32 + Integer.numberOfTrailingZeros(hWord);
-
- return hOffset;
- }
-}
diff --git a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene103/Lucene103PostingsReader.java b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene103/Lucene103PostingsReader.java
index e2895fbf5fa..c153802b334 100644
--- a/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene103/Lucene103PostingsReader.java
+++ b/lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene103/Lucene103PostingsReader.java
@@ -1050,7 +1050,8 @@ public final class Lucene103PostingsReader extends PostingsReaderBase {
}
// Only return docs from the current block
- // +1 to make BitSetUtil#denseBitsetToArray work, see its java doc.
+ // +1 to make FixedBitSet#intoArray perform a bit faster, it uses a slower path if it doesn't
+ // have a free slot after the last element
buffer.growNoCopy(BLOCK_SIZE + 1);
upTo = (int) Math.min(upTo, level0LastDocID + 1L);
@@ -1066,8 +1067,9 @@ public final class Lucene103PostingsReader extends PostingsReaderBase {
break;
case UNARY:
buffer.size =
- BitSetUtil.denseBitsetToArray(
- docBitSet, doc - docBitSetBase, upTo - docBitSetBase, docBitSetBase, buffer.docs);
+ docBitSet.intoArray(
+ doc - docBitSetBase, upTo - docBitSetBase, docBitSetBase, buffer.docs);
+ assert buffer.size < buffer.docs.length; // buffer's capacity is BLOCK_SIZE+1
break;
}
diff --git a/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene103/TestBitSetUtil.java b/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene103/TestBitSetUtil.java
deleted file mode 100644
index 094d2111f3a..00000000000
--- a/lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene103/TestBitSetUtil.java
+++ /dev/null
@@ -1,53 +0,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.
- */
-package org.apache.lucene.backward_codecs.lucene103;
-
-import static org.apache.lucene.backward_codecs.lucene103.Lucene103PostingsFormat.BLOCK_SIZE;
-
-import org.apache.lucene.tests.util.LuceneTestCase;
-import org.apache.lucene.util.FixedBitSet;
-
-public class TestBitSetUtil extends LuceneTestCase {
-
- public void testDenseBitsetToArray() throws Exception {
- for (int iter = 0; iter < 1000; iter++) {
- FixedBitSet bitSet = new FixedBitSet(BLOCK_SIZE + random().nextInt(200));
- while (bitSet.cardinality() < BLOCK_SIZE) {
- bitSet.set(random().nextInt(bitSet.length()));
- }
-
- int[] array = new int[BLOCK_SIZE + 1];
- int from, to, size;
- if (random().nextBoolean()) {
- from = 0;
- to = bitSet.length();
- } else {
- from = random().nextInt(bitSet.length());
- to = from + random().nextInt(bitSet.length() - from);
- }
-
- int base = random().nextInt(1000);
- size = BitSetUtil.denseBitsetToArray(bitSet, from, to, base, array);
-
- assertEquals(bitSet.cardinality(from, to), size);
-
- int[] index = new int[] {0};
- bitSet.forEach(from, to, base, bit -> assertEquals(bit, array[index[0]++]));
- assertEquals(size, index[0]);
- }
- }
-}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/BitSetUtil.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/BitSetUtil.java
deleted file mode 100644
index d5d72d35d84..00000000000
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/BitSetUtil.java
+++ /dev/null
@@ -1,104 +0,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.
- */
-package org.apache.lucene.codecs.lucene104;
-
-import java.util.Objects;
-import org.apache.lucene.util.FixedBitSet;
-
-class BitSetUtil {
-
- /**
- * Converts set bits in the given bitset to an array of document IDs. Only processes bits from
- * index {@code from} (inclusive) to {@code to} (exclusive) and returns the number of bits set in
- * this range.
- *
- * <p>Each set bit's position is converted to a document ID by adding the {@code base} value and
- * stored in the provided {@code array}.
- *
- * <p>NOTE: Caller need to ensure the {@code array} has a length greater than or equal to {@code
- * bitSet.cardinality(from, to) + 1}.
- */
- static int denseBitsetToArray(FixedBitSet bitSet, int from, int to, int base, int[] array) {
- assert bitSet.cardinality(from, to) + 1 <= array.length
- : "Array length must be at least bitSet.cardinality(from, to) + 1";
-
- Objects.checkFromToIndex(from, to, bitSet.length());
-
- int offset = 0;
- long[] bits = bitSet.getBits();
- // First, align `from` with a word start, ie. a multiple of Long.SIZE (64)
- if ((from & 0x3F) != 0) {
- long word = bits[from >> 6] >>> from;
- int numBitsTilNextWord = -from & 0x3F;
- if (to - from < numBitsTilNextWord) {
- // All bits are in a single word
- word &= (1L << (to - from)) - 1L;
- return word2Array(word, from + base, array, offset);
- }
- offset = word2Array(word, from + base, array, offset);
- from += numBitsTilNextWord;
- assert (from & 0x3F) == 0;
- }
-
- for (int i = from >> 6, end = to >> 6; i < end; ++i) {
- long word = bits[i];
- offset = word2Array(word, base + (i << 6), array, offset);
- }
-
- // Now handle remaining bits in the last partial word
- if ((to & 0x3F) != 0) {
- long word = bits[to >> 6] & ((1L << to) - 1);
- offset = word2Array(word, base + (to & ~0x3F), array, offset);
- }
-
- return offset;
- }
-
- private static int word2Array(long word, int base, int[] docs, int offset) {
- final int bitCount = Long.bitCount(word);
-
- if (bitCount >= 32) {
- return denseWord2Array(word, base, docs, offset);
- }
-
- for (int i = 0; i < bitCount; i++) {
- int ntz = Long.numberOfTrailingZeros(word);
- docs[offset++] = base + ntz;
- word ^= 1L << ntz;
- }
-
- return offset;
- }
-
- private static int denseWord2Array(long word, int base, int[] docs, int offset) {
- final int lWord = (int) word;
- final int hWord = (int) (word >>> 32);
- final int offset32 = offset + Integer.bitCount(lWord);
- int hOffset = offset32;
-
- for (int i = 0; i < 32; i++) {
- docs[offset] = base + i;
- docs[hOffset] = base + i + 32;
- offset += (lWord >>> i) & 1;
- hOffset += (hWord >>> i) & 1;
- }
-
- docs[offset32] = base + 32 + Integer.numberOfTrailingZeros(hWord);
-
- return hOffset;
- }
-}
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104PostingsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104PostingsReader.java
index ee028c4f847..19118d64dbe 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104PostingsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104PostingsReader.java
@@ -1050,7 +1050,8 @@ public final class Lucene104PostingsReader extends PostingsReaderBase {
}
// Only return docs from the current block
- // +1 to make BitSetUtil#denseBitsetToArray work, see its java doc.
+ // +1 to make FixedBitSet#intoArray perform a bit faster, it uses a slower path if it doesn't
+ // have a free slot after the last element
buffer.growNoCopy(BLOCK_SIZE + 1);
upTo = (int) Math.min(upTo, level0LastDocID + 1L);
@@ -1066,8 +1067,9 @@ public final class Lucene104PostingsReader extends PostingsReaderBase {
break;
case UNARY:
buffer.size =
- BitSetUtil.denseBitsetToArray(
- docBitSet, doc - docBitSetBase, upTo - docBitSetBase, docBitSetBase, buffer.docs);
+ docBitSet.intoArray(
+ doc - docBitSetBase, upTo - docBitSetBase, docBitSetBase, buffer.docs);
+ assert buffer.size < buffer.docs.length; // buffer's capacity is BLOCK_SIZE+1
break;
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/BitSetDocIdStream.java b/lucene/core/src/java/org/apache/lucene/search/BitSetDocIdStream.java
index 8c2f01e38e3..ebe1ea9f737 100644
--- a/lucene/core/src/java/org/apache/lucene/search/BitSetDocIdStream.java
+++ b/lucene/core/src/java/org/apache/lucene/search/BitSetDocIdStream.java
@@ -58,4 +58,18 @@ final class BitSetDocIdStream extends DocIdStream {
return 0;
}
}
+
+ @Override
+ public int intoArray(int upTo, int[] array) {
+ if (upTo > this.upTo) {
+ upTo = Math.min(upTo, max);
+ int count = bitSet.intoArray(this.upTo - offset, upTo - offset, offset, array);
+ if (count == array.length) { // The whole range of doc IDs may not have been copied
+ upTo = array[array.length - 1] + 1;
+ }
+ this.upTo = upTo;
+ return count;
+ }
+ return 0;
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/BooleanScorerSupplier.java b/lucene/core/src/java/org/apache/lucene/search/BooleanScorerSupplier.java
index 944812b8b18..04e75cf10fa 100644
--- a/lucene/core/src/java/org/apache/lucene/search/BooleanScorerSupplier.java
+++ b/lucene/core/src/java/org/apache/lucene/search/BooleanScorerSupplier.java
@@ -266,6 +266,11 @@ final class BooleanScorerSupplier extends ScorerSupplier {
public void collect(int doc) throws IOException {
collector.collect(doc);
}
+
+ @Override
+ public void collect(DocIdStream stream) throws IOException {
+ collector.collect(stream);
+ }
};
return scorer.score(noScoreCollector, acceptDocs, min, max);
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/DocIdStream.java b/lucene/core/src/java/org/apache/lucene/search/DocIdStream.java
index b13e0ce6ec1..b22a0aa83ba 100644
--- a/lucene/core/src/java/org/apache/lucene/search/DocIdStream.java
+++ b/lucene/core/src/java/org/apache/lucene/search/DocIdStream.java
@@ -58,6 +58,22 @@ public abstract class DocIdStream {
// so would defeat the purpose of collecting hits via a DocIdStream.
public abstract int count(int upTo) throws IOException;
+ /**
+ * Copy some matching doc IDs into the provided array and return the number of copied elements. A
+ * return value of {@code 0} indicates that there are no remaining doc IDs. The given array must
+ * not be empty.
+ */
+ public int intoArray(int[] array) {
+ return intoArray(DocIdSetIterator.NO_MORE_DOCS, array);
+ }
+
+ /**
+ * Copy some matching doc IDs under {@code upTo} (exclusive) into the provided array and return
+ * the number of copied elements. A return value of {@code 0} indicates that there are no matching
+ * doc IDs under {@code upTo} anymore. The given array must not be empty.
+ */
+ public abstract int intoArray(int upTo, int[] array);
+
/**
* Return {@code true} if this stream may have remaining doc IDs. This must eventually return
* {@code false} when the stream is exhausted.
diff --git a/lucene/core/src/java/org/apache/lucene/search/LRUQueryCache.java b/lucene/core/src/java/org/apache/lucene/search/LRUQueryCache.java
index 05ef52ff327..0add1c1bc5e 100644
--- a/lucene/core/src/java/org/apache/lucene/search/LRUQueryCache.java
+++ b/lucene/core/src/java/org/apache/lucene/search/LRUQueryCache.java
@@ -539,6 +539,8 @@ public class LRUQueryCache implements QueryCache, Accountable {
scorer.score(
new LeafCollector() {
+ private int[] buffer;
+
@Override
public void setScorer(Scorable scorer) throws IOException {}
@@ -547,6 +549,19 @@ public class LRUQueryCache implements QueryCache, Accountable {
count[0]++;
bitSet.set(doc);
}
+
+ @Override
+ public void collect(DocIdStream stream) throws IOException {
+ if (buffer == null) {
+ buffer = new int[128];
+ }
+ for (int c = stream.intoArray(buffer); c != 0; c = stream.intoArray(buffer)) {
+ for (int i = 0; i < c; ++i) {
+ bitSet.set(buffer[i]);
+ }
+ count[0] += c;
+ }
+ }
},
null,
0,
@@ -560,6 +575,8 @@ public class LRUQueryCache implements QueryCache, Accountable {
scorer.score(
new LeafCollector() {
+ private int[] buffer = null;
+
@Override
public void setScorer(Scorable scorer) throws IOException {}
@@ -567,6 +584,18 @@ public class LRUQueryCache implements QueryCache, Accountable {
public void collect(int doc) throws IOException {
builder.add(doc);
}
+
+ @Override
+ public void collect(DocIdStream stream) throws IOException {
+ if (buffer == null) {
+ buffer = new int[128];
+ }
+ for (int c = stream.intoArray(buffer); c != 0; c = stream.intoArray(buffer)) {
+ for (int i = 0; i < c; ++i) {
+ builder.add(buffer[i]);
+ }
+ }
+ }
},
null,
0,
diff --git a/lucene/core/src/java/org/apache/lucene/search/RangeDocIdStream.java b/lucene/core/src/java/org/apache/lucene/search/RangeDocIdStream.java
index fc128880fc1..07aa4ab334f 100644
--- a/lucene/core/src/java/org/apache/lucene/search/RangeDocIdStream.java
+++ b/lucene/core/src/java/org/apache/lucene/search/RangeDocIdStream.java
@@ -17,6 +17,7 @@
package org.apache.lucene.search;
import java.io.IOException;
+import org.apache.lucene.util.MathUtil;
final class RangeDocIdStream extends DocIdStream {
@@ -58,4 +59,19 @@ final class RangeDocIdStream extends DocIdStream {
return 0;
}
}
+
+ @Override
+ public int intoArray(int upTo, int[] array) {
+ int start = this.upTo;
+ upTo = Math.min(upTo, max);
+ upTo = MathUtil.unsignedMin(upTo, start + array.length);
+ if (upTo > start) {
+ for (int doc = start; doc < upTo; ++doc) {
+ array[doc - start] = doc;
+ }
+ this.upTo = upTo;
+ return upTo - start;
+ }
+ return 0;
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java b/lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java
index 1d8b4d16ecd..f1d4cc9cf17 100644
--- a/lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java
+++ b/lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java
@@ -870,4 +870,82 @@ public final class FixedBitSet extends BitSet {
bits ^= 1L << ntz;
}
}
+
+ /**
+ * Converts set bits in this bitset to an array of document IDs. Only processes bits from index
+ * {@code from} (inclusive) to {@code to} (exclusive) and returns the number of bits copied.
+ *
+ * <p>Each set bit's position is converted to a document ID by adding the {@code base} value and
+ * stored in the provided {@code array}. This method stops when there are no more set bits before
+ * {@code to} or when there is no capacity left in the given {@code array}, whichever comes first.
+ */
+ public int intoArray(int from, int to, int base, int[] array) {
+ Objects.checkFromToIndex(from, to, length());
+
+ int offset = 0;
+ // First, align `from` with a word start, ie. a multiple of Long.SIZE (64)
+ if ((from & 0x3F) != 0) {
+ long word = bits[from >> 6] >>> from;
+ int numBitsTilNextWord = -from & 0x3F;
+ if (to - from < numBitsTilNextWord) {
+ // All bits are in a single word
+ word &= (1L << (to - from)) - 1L;
+ return word2Array(word, from + base, array, offset);
+ }
+ offset = word2Array(word, from + base, array, offset);
+ from += numBitsTilNextWord;
+ assert (from & 0x3F) == 0;
+ }
+
+ for (int i = from >> 6, end = to >> 6; i < end; ++i) {
+ long word = bits[i];
+ offset = word2Array(word, base + (i << 6), array, offset);
+ }
+
+ // Now handle remaining bits in the last partial word
+ if ((to & 0x3F) != 0) {
+ long word = bits[to >> 6] & ((1L << to) - 1);
+ offset = word2Array(word, base + (to & ~0x3F), array, offset);
+ }
+
+ return offset;
+ }
+
+ private static int word2Array(long word, int base, int[] docs, int offset) {
+ final int bitCount = Long.bitCount(word);
+
+ if (bitCount >= 32 && docs.length - offset > bitCount) {
+ return denseWord2Array(word, base, docs, offset);
+ }
+
+ int numBitsToCopy = Math.min(bitCount, docs.length - offset);
+
+ for (int i = 0; i < numBitsToCopy; i++) {
+ int ntz = Long.numberOfTrailingZeros(word);
+ docs[offset + i] = base + ntz;
+ word ^= 1L << ntz;
+ }
+
+ return offset + numBitsToCopy;
+ }
+
+ private static int denseWord2Array(long word, int base, int[] docs, int offset) {
+ assert docs.length - offset >= Long.bitCount(word) + 1;
+
+ final int lWord = (int) word;
+ final int hWord = (int) (word >>> 32);
+ final int offset32 = offset + Integer.bitCount(lWord);
+ int hOffset = offset32;
+
+ for (int i = 0; i < 32; i++) {
+ docs[offset] = base + i;
+ docs[hOffset] = base + i + 32;
+ offset += (lWord >>> i) & 1;
+ hOffset += (hWord >>> i) & 1;
+ }
+
+ docs[offset32] = base + 32 + Integer.numberOfTrailingZeros(hWord);
+
+ return hOffset;
+ }
}
diff --git a/lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestBitSetUtil.java b/lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestBitSetUtil.java
deleted file mode 100644
index 966c61cfe28..00000000000
--- a/lucene/core/src/test/org/apache/lucene/codecs/lucene104/TestBitSetUtil.java
+++ /dev/null
@@ -1,53 +0,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.
- */
-package org.apache.lucene.codecs.lucene104;
-
-import static org.apache.lucene.codecs.lucene104.Lucene104PostingsFormat.BLOCK_SIZE;
-
-import org.apache.lucene.tests.util.LuceneTestCase;
-import org.apache.lucene.util.FixedBitSet;
-
-public class TestBitSetUtil extends LuceneTestCase {
-
- public void testDenseBitsetToArray() throws Exception {
- for (int iter = 0; iter < 1000; iter++) {
- FixedBitSet bitSet = new FixedBitSet(BLOCK_SIZE + random().nextInt(200));
- while (bitSet.cardinality() < BLOCK_SIZE) {
- bitSet.set(random().nextInt(bitSet.length()));
- }
-
- int[] array = new int[BLOCK_SIZE + 1];
- int from, to, size;
- if (random().nextBoolean()) {
- from = 0;
- to = bitSet.length();
- } else {
- from = random().nextInt(bitSet.length());
- to = from + random().nextInt(bitSet.length() - from);
- }
-
- int base = random().nextInt(1000);
- size = BitSetUtil.denseBitsetToArray(bitSet, from, to, base, array);
-
- assertEquals(bitSet.cardinality(from, to), size);
-
- int[] index = new int[] {0};
- bitSet.forEach(from, to, base, bit -> assertEquals(bit, array[index[0]++]));
- assertEquals(size, index[0]);
- }
- }
-}
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestBitSetDocIdStream.java b/lucene/core/src/test/org/apache/lucene/search/TestBitSetDocIdStream.java
index 25fcea1a803..8ffd525d7a4 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestBitSetDocIdStream.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestBitSetDocIdStream.java
@@ -54,6 +54,29 @@ public class TestBitSetDocIdStream extends LuceneTestCase {
assertEquals(bitSet.cardinality(), stream.count());
}
+ public void testIntoArray() {
+ FixedBitSet bitSet = randomBitSet();
+ BitSetIterator bitSetIterator = new BitSetIterator(bitSet, 0L);
+ int base = 42;
+ BitSetDocIdStream stream = new BitSetDocIdStream(bitSet, base);
+ int[] array = new int[16];
+
+ int o = array.length;
+ int count = array.length;
+
+ for (int expected = bitSetIterator.nextDoc();
+ expected != DocIdSetIterator.NO_MORE_DOCS;
+ expected = bitSetIterator.nextDoc()) {
+ if (o == count) {
+ count = stream.intoArray(array);
+ o = 0;
+ }
+ assertEquals(expected + base, array[o++]);
+ }
+ assertEquals(count, o);
+ assertEquals(0, stream.intoArray(array));
+ }
+
public void testForEachUpTo() throws IOException {
FixedBitSet bitSet = randomBitSet();
int base = 42;
@@ -99,6 +122,36 @@ public class TestBitSetDocIdStream extends LuceneTestCase {
assertFalse(stream.mayHaveRemaining());
}
+ public void testIntoArrayUpTo() throws IOException {
+ FixedBitSet bitSet = randomBitSet();
+ BitSetIterator bitSetIterator = new BitSetIterator(bitSet, 0L);
+ bitSetIterator.nextDoc();
+ int base = 42;
+ BitSetDocIdStream stream = new BitSetDocIdStream(bitSet, base);
+ int[] array = new int[16];
+
+ int o = array.length;
+ int count = array.length;
+
+ for (int upTo = 0; upTo < bitSet.length(); ) {
+ int newUpTo = Math.min(upTo + random().nextInt(40), 100);
+
+ for (int expected = bitSetIterator.docID();
+ expected < newUpTo;
+ expected = bitSetIterator.nextDoc()) {
+ if (o == count) {
+ count = stream.intoArray(base + newUpTo, array);
+ o = 0;
+ }
+ assertEquals(expected + base, array[o++]);
+ }
+ assertEquals(count, o);
+ assertEquals(0, stream.intoArray(base + newUpTo, array));
+
+ upTo = newUpTo;
+ }
+ }
+
public void testMixForEachCountUpTo() throws IOException {
FixedBitSet bitSet = randomBitSet();
int base = 42;
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestRangeDocIdStream.java b/lucene/core/src/test/org/apache/lucene/search/TestRangeDocIdStream.java
index 1aa0bdb0fa8..356417bc1c7 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestRangeDocIdStream.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestRangeDocIdStream.java
@@ -36,6 +36,25 @@ public class TestRangeDocIdStream extends LuceneTestCase {
assertEquals(100 - 42, stream.count());
}
+ public void testIntoArray() throws IOException {
+ RangeDocIdStream stream = new RangeDocIdStream(42, 100);
+ int[] array = new int[16];
+ int o = array.length;
+ int count = array.length;
+
+ for (int i = 42; i < 100; ++i) {
+ if (o == count) {
+ count = stream.intoArray(array);
+ o = 0;
+ if (100 - i >= array.length) {
+ assertEquals(array.length, count);
+ }
+ }
+ assertEquals(i, array[o++]);
+ }
+ assertEquals(count, o);
+ }
+
public void testForEachUpTo() throws IOException {
RangeDocIdStream stream = new RangeDocIdStream(42, 100);
int[] expected = new int[] {42};
@@ -73,6 +92,29 @@ public class TestRangeDocIdStream extends LuceneTestCase {
assertFalse(stream.mayHaveRemaining());
}
+ public void testIntoArrayUpTo() throws IOException {
+ RangeDocIdStream stream = new RangeDocIdStream(42, 100);
+ int[] array = new int[16];
+ int o = array.length;
+ int count = array.length;
+
+ for (int upTo = 42; upTo < 100; ) {
+ int newUpTo = Math.min(upTo + random().nextInt(40), 100);
+ for (int i = upTo; i < newUpTo; ++i) {
+ if (o == count) {
+ count = stream.intoArray(newUpTo, array);
+ o = 0;
+ if (newUpTo - i >= array.length) {
+ assertEquals(array.length, count);
+ }
+ }
+ assertEquals(i, array[o++]);
+ }
+ assertEquals(count, o);
+ upTo = newUpTo;
+ }
+ }
+
public void testMixForEachCountUpTo() throws IOException {
RangeDocIdStream stream = new RangeDocIdStream(42, 100);
int[] expected = new int[] {42};
diff --git a/lucene/core/src/test/org/apache/lucene/util/TestFixedBitSet.java b/lucene/core/src/test/org/apache/lucene/util/TestFixedBitSet.java
index efdebee7158..b15776840f4 100644
--- a/lucene/core/src/test/org/apache/lucene/util/TestFixedBitSet.java
+++ b/lucene/core/src/test/org/apache/lucene/util/TestFixedBitSet.java
@@ -775,4 +775,48 @@ public class TestFixedBitSet extends BaseBitSetTestCase<FixedBitSet> {
assertEquals(expected, actual);
}
+
+ public void testIntoArray() throws Exception {
+ for (int outerIter = 0; outerIter < 100; outerIter++) {
+ int numBits = TestUtil.nextInt(random(), 10, 1_000);
+ int numSetBits = TestUtil.nextInt(random(), 0, numBits);
+
+ FixedBitSet bitSet = new FixedBitSet(numBits);
+ while (bitSet.cardinality() < numSetBits) {
+ bitSet.set(random().nextInt(bitSet.length()));
+ }
+
+ for (int innerIter = 0; innerIter < 100; ++innerIter) {}
+
+ int[] array = new int[TestUtil.nextInt(random(), 0, numSetBits + 10)];
+ int from, to;
+ if (random().nextBoolean()) {
+ from = 0;
+ to = bitSet.length();
+ } else {
+ from = random().nextInt(bitSet.length());
+ to = from + random().nextInt(bitSet.length() - from);
+ }
+
+ int base = random().nextInt(1000);
+ int size = bitSet.intoArray(from, to, base, array);
+
+ if (size < array.length) {
+ // All set bits should have been copied
+ assertEquals(bitSet.cardinality(from, to), size);
+ }
+
+ int[] index = new int[] {0};
+ bitSet.forEach(
+ from,
+ to,
+ base,
+ bit -> {
+ if (index[0] < array.length) {
+ assertEquals(bit, array[index[0]++]);
+ }
+ });
+ assertEquals(size, index[0]);
+ }
+ }
}
diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/facet/plain/histograms/HistogramCollector.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/facet/plain/histograms/HistogramCollector.java
index cca4e2b92a4..4b77cc7dd64 100644
--- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/facet/plain/histograms/HistogramCollector.java
+++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/facet/plain/histograms/HistogramCollector.java
@@ -181,6 +181,9 @@ final class HistogramCollector implements Collector {
private final int maxBuckets;
private final LongIntHashMap counts;
+ private final int[] docBuffer = new int[64];
+ private final long[] valueBuffer = new long[docBuffer.length];
+
HistogramNaiveSingleValuedLeafCollector(
NumericDocValues values, long bucketWidth, int maxBuckets, LongIntHashMap counts) {
this.values = values;
@@ -192,15 +195,42 @@ final class HistogramCollector implements Collector {
@Override
public void setScorer(Scorable scorer) throws IOException {}
+ @Override
+ public void collect(DocIdStream stream) throws IOException {
+ // Optimize bulk retrieval of doc values
+ for (int count = stream.intoArray(docBuffer);
+ count != 0;
+ count = stream.intoArray(docBuffer)) {
+ int firstDoc = docBuffer[0];
+ int lastDoc = docBuffer[count - 1];
+
+ if (values.advanceExact(firstDoc) && values.docIDRunEnd() > lastDoc) {
+ // This guarantees that all docs in docBuffer have a value.
+ values.longValues(count, docBuffer, valueBuffer, 0L);
+ for (int i = 0; i < count; ++i) {
+ collectValue(valueBuffer[i]);
+ }
+ } else {
+ // Delegate to #collect to handle docs with missing values
+ for (int i = 0; i < count; ++i) {
+ collect(docBuffer[i]);
+ }
+ }
+ }
+ }
+
@Override
public void collect(int doc) throws IOException {
if (values.advanceExact(doc)) {
- final long value = values.longValue();
- final long bucket = Math.floorDiv(value, bucketWidth);
- counts.addTo(bucket, 1);
- checkMaxBuckets(counts.size(), maxBuckets);
+ collectValue(values.longValue());
}
}
+
+ private void collectValue(long value) {
+ final long bucket = Math.floorDiv(value, bucketWidth);
+ counts.addTo(bucket, 1);
+ checkMaxBuckets(counts.size(), maxBuckets);
+ }
}
/**
diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java
index c4cc2394a78..d481daa858a 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java
@@ -16,6 +16,8 @@
*/
package org.apache.lucene.tests.search;
+import static org.apache.lucene.tests.util.LuceneTestCase.rarely;
+
import java.io.IOException;
import org.apache.lucene.search.CheckedIntConsumer;
import org.apache.lucene.search.DocIdSetIterator;
@@ -35,6 +37,8 @@ class AssertingLeafCollector extends FilterLeafCollector {
private int lastCollected = -1;
private boolean finishCalled;
+ private int[] docBuffer;
+
AssertingLeafCollector(LeafCollector collector, int min, int max) {
super(collector);
this.min = min;
@@ -48,7 +52,20 @@ class AssertingLeafCollector extends FilterLeafCollector {
@Override
public void collect(DocIdStream stream) throws IOException {
- in.collect(new AssertingDocIdStream(stream));
+ if (rarely()) {
+ if (docBuffer == null) {
+ docBuffer = new int[32];
+ }
+ for (int count = stream.intoArray(docBuffer);
+ count != 0;
+ count = stream.intoArray(docBuffer)) {
+ for (int i = 0; i < count; ++i) {
+ collect(docBuffer[i]);
+ }
+ }
+ } else {
+ in.collect(new AssertingDocIdStream(stream));
+ }
}
@Override
@@ -179,6 +196,45 @@ class AssertingLeafCollector extends FilterLeafCollector {
return count;
}
+ @Override
+ public int intoArray(int[] array) {
+ assert array.length > 0;
+ int count = stream.intoArray(array);
+ assert lastUpTo != DocIdSetIterator.NO_MORE_DOCS || count == 0;
+ if (count < array.length) {
+ lastUpTo = DocIdSetIterator.NO_MORE_DOCS;
+ assert stream.mayHaveRemaining() == false;
+ } else {
+ lastUpTo = array[array.length - 1] + 1;
+ }
+ return count;
+ }
+
+ @Override
+ public int intoArray(int upTo, int[] array) {
+ assert array.length > 0;
+ assert lastUpTo <= upTo : "upTo=" + upTo + " but previous upTo=" + lastUpTo;
+ int count = stream.intoArray(upTo, array);
+
+ assert lastUpTo != upTo || count == 0;
+
+ if (count != 0) {
+ assert array[count - 1] < upTo;
+ }
+
+ if (count < array.length) {
+ lastUpTo = upTo;
+ } else {
+ lastUpTo = array[array.length - 1] + 1;
+ }
+
+ if (upTo == DocIdSetIterator.NO_MORE_DOCS && count < array.length) {
+ assert stream.mayHaveRemaining() == false;
+ }
+
+ return count;
+ }
+
@Override
public boolean mayHaveRemaining() {
boolean mayHaveRemaining = stream.mayHaveRemaining();