This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 3a36f139597 Optimize bitmap unions in query aggregation, star-tree
build, and JSON indexes (#19587)
3a36f139597 is described below
commit 3a36f139597d55a3951e9fb6d6814d1bb2ab7cfc
Author: Xiang Fu <[email protected]>
AuthorDate: Wed Oct 7 17:04:20 2026 -0700
Optimize bitmap unions in query aggregation, star-tree build, and JSON
indexes (#19587)
* Union serialized bitmaps lazily in DISTINCT_COUNT_BITMAP aggregation
The serialized-bitmap (BYTES) aggregation paths - the ones hit when unioning
pre-aggregated bitmaps such as star-tree DISTINCT_COUNT_BITMAP columns -
fold
every matching document's bitmap into a per-group accumulator with
RoaringBitmap.or(), which recomputes container cardinality and re-normalizes
container types on every union. On a production-like workload dominated by
DISTINCT_COUNT_BITMAP over star-tree pre-aggregated bitmaps, JFR attributed
~85% of server execution samples to these unions.
Switch the fold to RoaringBitmap's lazy union (the primitive FastAggregation
uses internally): container cardinality maintenance is skipped during
accumulation, array containers promote to bitmap containers past 1024
combined
cardinality so subsequent unions of hot containers become word-wise ORs, and
the accumulator is repaired once at extraction. The new
RoaringBitmapLazyUnion
shim exposes the protected lazyor()/repairAfterLazy() primitives;
extractAggregationResult/extractGroupByResult repair the accumulator before
it
escapes the function, so results, intermediate types and wire format are
unchanged.
Fold benchmarks (RoaringBitmap 1.6.23): 2.2-2.7x faster with ~flat
allocation
for dense accumulators (small value universe), 4.7x faster for sparse
accumulators (at the cost of transient young-gen allocation from
pre-promotion
array copies), 5.4x faster for singleton-heavy input streams.
* Extend lazy bitmap unions to star-tree build, json index, and inverted
index count
- Move the RoaringBitmapLazyUnion shim from pinot-core to pinot-common and
add a
buffer-bitmap counterpart (MutableRoaringBitmapLazyUnion) so
pinot-segment-local
can use the same primitives.
- DistinctCountBitmapValueAggregator (star-tree build / segment
aggregation):
union serialized-bitmap raw values and aggregated values lazily,
repairing at
serializeAggregatedValue(). Inputs to applyAggregatedValue() are repaired
first
because on-heap star-tree records may themselves be lazy accumulators. The
per-row _maxByteSize tracking stays valid on a lazy accumulator:
serialized
container sizes do not depend on the deferred cardinality values, and a
lazy
bitmap container measures at its fixed 8KiB size, which upper-bounds the
size
after repair. Fold benchmarks: 2.2-6.2x faster.
- JSON index (ImmutableJsonIndexReader and MutableJsonIndexImpl): fold the
IN / NOT_IN / REGEXP_LIKE / RANGE posting-list unions through a lazy-union
accumulator that borrows single inputs without copying and repairs once
at the
end. Fold benchmarks: 5-11x faster.
- InvertedIndexFilterOperator: lazy fold for the >2-dictId union-cardinality
case (still materializes the union, but skips per-union cardinality
maintenance).
* Optimize sparse bitmap intermediate-result merges
* Use RoaringBitmapUnion accumulators for bitmap unions
Replace the three classes this branch shipped in the org.roaringbitmap
packages with RoaringBitmapUnion and MutableRoaringBitmapUnion in
org.apache.pinot.common.utils. They carry the names and methods of the
accumulators proposed to RoaringBitmap, so moving to the library classes
later is an import change plus one line in RoaringBitmapUtils. They reach
the protected lazy union through private subclasses, which removes the
split package and the use of package-private and deprecated API, and they
build against the released RoaringBitmap 1.6.23.
The released lazy union only pays off for dense containers and inserts
new container keys one at a time, so the interim classes union eagerly
while the accumulator is sparse or when an input brings many new keys.
Star-tree build: the bitmap aggregator reports the largest value it
serialized instead of measuring the accumulator on every apply. The
on-heap builder serializes BYTES metrics before it sizes the forward
indexes so that maximum is complete when it is read.
The sparse merge optimization is dropped from Pinot; it is part of the
RoaringBitmap change and merge() goes back to a plain or().
* Keep repeated sparse bitmap unions eager
* Avoid sparse union copies and unrelated StarTree serialization
* Trim unused bitmap union APIs and unrelated merge tests
---------
Co-authored-by: intentlab-ai[bot]
<intentlab-ai[bot]@users.noreply.github.com>
---
.../common/utils/MutableRoaringBitmapUnion.java | 182 ++++++++++
.../pinot/common/utils/RoaringBitmapUnion.java | 240 +++++++++++++
.../pinot/common/utils/RoaringBitmapUtils.java | 8 +
.../utils/MutableRoaringBitmapUnionTest.java | 299 +++++++++++++++
.../pinot/common/utils/RoaringBitmapUnionTest.java | 399 +++++++++++++++++++++
.../filter/InvertedIndexFilterOperator.java | 9 +-
.../DistinctCountBitmapAggregationFunction.java | 78 ++--
.../function/DistinctCountBitmapLazyUnionTest.java | 221 ++++++++++++
.../pinot/core/startree/v2/BaseStarTreeV2Test.java | 11 +-
.../v2/DistinctCountBitmapStarTreeV2Test.java | 18 +-
.../DistinctCountBitmapWithMVStarTreeV2Test.java | 17 +-
...a => DistinctCountHLLOnHeapStarTreeV2Test.java} | 32 +-
.../BenchmarkDistinctCountBitmapAggregation.java | 121 +++++++
.../DistinctCountBitmapValueAggregator.java | 60 ++--
.../segment/local/aggregator/ValueAggregator.java | 5 +
.../realtime/impl/json/MutableJsonIndexImpl.java | 67 ++--
.../readers/json/ImmutableJsonIndexReader.java | 59 ++-
.../startree/v2/builder/BaseSingleTreeBuilder.java | 17 +
.../v2/builder/OnHeapSingleTreeBuilder.java | 24 ++
.../DistinctCountBitmapValueAggregatorTest.java | 161 +++++++++
.../v2/builder/NullAwareStarTreeBuilderTest.java | 99 ++++-
21 files changed, 1993 insertions(+), 134 deletions(-)
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/utils/MutableRoaringBitmapUnion.java
b/pinot-common/src/main/java/org/apache/pinot/common/utils/MutableRoaringBitmapUnion.java
new file mode 100644
index 00000000000..8e41a931abb
--- /dev/null
+++
b/pinot-common/src/main/java/org/apache/pinot/common/utils/MutableRoaringBitmapUnion.java
@@ -0,0 +1,182 @@
+/**
+ * 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.pinot.common.utils;
+
+import java.util.Objects;
+import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+import org.roaringbitmap.buffer.MappeableContainerPointer;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
+
+
+/// Interim stand-in for `org.roaringbitmap.buffer.MutableRoaringBitmapUnion`,
the buffer counterpart of
+/// [RoaringBitmapUnion]. It exposes only the proposed methods used by Pinot,
so that the swap is an import change
+/// plus deleting this class. It is internal to Pinot and goes away with the
swap. See [RoaringBitmapUnion] for the
+/// contract and the swap steps.
+///
+/// Inputs may be read-only bitmaps over memory-mapped buffers; they are only
read.
+///
+/// Differences from the library class:
+///
+/// - The accumulated state is a private subclass of [MutableRoaringBitmap]
that reaches the protected lazy primitives
+/// through inheritance; [#get()] and [#take()] hand out instances of it.
+/// - An input is unioned eagerly while the accumulator holds few values per
container, when overlapping arrays would
+/// stay small, or when it would insert many container keys that are new to
the accumulator, because the released
+/// lazy union only pays off for dense containers and inserts new keys one
at a time.
+///
+/// Instances are not thread-safe.
+public final class MutableRoaringBitmapUnion {
+ // See RoaringBitmapUnion
+ private static final int MIN_VALUES_PER_CONTAINER_FOR_LAZY_UNION = 1024;
+ private static final int MAX_NEW_KEYS_FOR_LAZY_UNION = 4;
+
+ // The accumulated bitmap. Never null.
+ private LazyBitmap _bitmap;
+ // Whether the bitmap holds lazy state that must be repaired before it is
read
+ private boolean _dirty;
+ // Whether an alias returned by get() is outstanding and must not be mutated
+ private boolean _published;
+ // Number of values added so far, counting duplicates: an upper bound of the
cardinality that is known without
+ // repairing, used to estimate how dense the containers are
+ private long _numValuesAdded;
+ // A cardinality observed on repaired state. Unions never remove values, so
this remains a lower bound used to
+ // check average density without repairing every lazy union. Overlapping
arrays are checked separately.
+ private long _lastKnownCardinality;
+
+ /// Creates an empty union.
+ public MutableRoaringBitmapUnion() {
+ _bitmap = new LazyBitmap();
+ }
+
+ /// Unions the input into the accumulated state. The input is neither
modified nor retained. Adding this union's own
+ /// [#get()] result, or an empty bitmap, is a no-op.
+ public void add(ImmutableRoaringBitmap input) {
+ Objects.requireNonNull(input, "input");
+ if (input == _bitmap || input.isEmpty()) {
+ return;
+ }
+ beforeMutation();
+ boolean dense = isDenseForLazyUnion();
+ _numValuesAdded += input.getLongCardinality();
+ if (dense && !shouldUnionEagerly(input)) {
+ _dirty = true;
+ _bitmap.lazyOr(input);
+ } else {
+ repairIfDirty();
+ _bitmap.or(input);
+ }
+ }
+
+ /// Returns the accumulated bitmap with all pending lazy state repaired. The
returned instance is the union's
+ /// current state, not a copy; the union copies it before its next mutation,
so the returned bitmap stays valid.
+ /// Repeated calls without an intervening mutation return the same instance.
+ public MutableRoaringBitmap get() {
+ repairIfDirty();
+ _published = true;
+ return _bitmap;
+ }
+
+ /// Returns the accumulated bitmap, repaired as by [#get()], and transfers
its ownership to the caller. The union is
+ /// empty afterwards and can be reused.
+ public MutableRoaringBitmap take() {
+ repairIfDirty();
+ MutableRoaringBitmap result = _bitmap;
+ _bitmap = new LazyBitmap();
+ _dirty = false;
+ _published = false;
+ _numValuesAdded = 0;
+ _lastKnownCardinality = 0;
+ return result;
+ }
+
+ /// Returns whether the lazy union would copy small overlapping arrays or
insert more than a few new keys between
+ /// existing ones. Keys past the accumulator's last key do not count: they
are appended in one step.
+ private boolean shouldUnionEagerly(ImmutableRoaringBitmap input) {
+ // Called after the average density check: an empty or sole dense
container has no small-array overlaps.
+ // A small input cannot insert too many keys. Avoid allocating scan
pointers for this common case.
+ if (_bitmap.getContainerCount() <= 1 && input.getContainerCount() <=
MAX_NEW_KEYS_FOR_LAZY_UNION) {
+ return false;
+ }
+ MappeableContainerPointer own = _bitmap.getContainerPointer();
+ MappeableContainerPointer other = input.getContainerPointer();
+ int numNewKeys = 0;
+ while (other.hasContainer()) {
+ char key = other.key();
+ while (own.hasContainer() && own.key() < key) {
+ own.advance();
+ }
+ if (!own.hasContainer()) {
+ return false;
+ }
+ if (own.key() == key) {
+ if (!own.isBitmapContainer() && !own.isRunContainer() &&
!other.isBitmapContainer() && !other.isRunContainer()
+ && own.getCardinality() + other.getCardinality() <=
MIN_VALUES_PER_CONTAINER_FOR_LAZY_UNION) {
+ return true;
+ }
+ } else if (++numNewKeys > MAX_NEW_KEYS_FOR_LAZY_UNION) {
+ return true;
+ }
+ other.advance();
+ }
+ return false;
+ }
+
+ private boolean isDenseForLazyUnion() {
+ long minCardinality = (long) MIN_VALUES_PER_CONTAINER_FOR_LAZY_UNION *
_bitmap.getContainerCount();
+ if (_lastKnownCardinality >= minCardinality) {
+ return true;
+ }
+ if (_numValuesAdded < minCardinality) {
+ return false;
+ }
+ // The estimate counts duplicates. Confirm actual density before switching
a sparse accumulator to lazy unions.
+ repairIfDirty();
+ _lastKnownCardinality = _bitmap.getLongCardinality();
+ _numValuesAdded = _lastKnownCardinality;
+ return _lastKnownCardinality >= minCardinality;
+ }
+
+ private void repairIfDirty() {
+ if (_dirty) {
+ _bitmap.repair();
+ _dirty = false;
+ }
+ }
+
+ private void beforeMutation() {
+ if (_published) {
+ _bitmap = (LazyBitmap) _bitmap.clone();
+ _published = false;
+ }
+ }
+
+ /// Reaches the protected lazy primitives of [MutableRoaringBitmap] through
inheritance.
+ private static final class LazyBitmap extends MutableRoaringBitmap {
+ // Public so that an instance that escaped through get() or take() stays
deserializable like a plain bitmap
+ public LazyBitmap() {
+ }
+
+ void lazyOr(ImmutableRoaringBitmap other) {
+ lazyor(other);
+ }
+
+ void repair() {
+ repairAfterLazy();
+ }
+ }
+}
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/utils/RoaringBitmapUnion.java
b/pinot-common/src/main/java/org/apache/pinot/common/utils/RoaringBitmapUnion.java
new file mode 100644
index 00000000000..2c7a05986ad
--- /dev/null
+++
b/pinot-common/src/main/java/org/apache/pinot/common/utils/RoaringBitmapUnion.java
@@ -0,0 +1,240 @@
+/**
+ * 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.pinot.common.utils;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Objects;
+import org.roaringbitmap.ContainerPointer;
+import org.roaringbitmap.RoaringBitmap;
+
+
+/// Interim stand-in for `org.roaringbitmap.RoaringBitmapUnion`, which is
proposed to RoaringBitmap for
+/// apache/pinot#19587 and not released yet. It exposes only the proposed
methods used by Pinot, so that once a
+/// RoaringBitmap release ships the class the swap is mechanical:
+///
+/// - change the import at the call sites to
`org.roaringbitmap.RoaringBitmapUnion`,
+/// - make [RoaringBitmapUtils#deserializeToUnion] call
`RoaringBitmapUnion.takeOwnership(deserialize(bytes))`,
+/// - delete this class, [MutableRoaringBitmapUnion] and their two tests (the
library tests its own classes).
+///
+/// `RoaringBitmapUnionTest` fails as soon as the library class is on the
classpath, as a reminder. This class is
+/// internal to Pinot and goes away with the swap; code outside this
repository should not depend on it.
+///
+/// An incremental union accumulator: it folds bitmaps that arrive over time
into one bitmap with RoaringBitmap's lazy
+/// union, which skips cardinality maintenance and keeps overlapping
containers in a form that makes further unions
+/// cheap, and it repairs the result only when it is read. Callers only ever
observe valid bitmaps:
+///
+/// - [#add(RoaringBitmap)] never modifies or retains its argument.
+/// - [#get()] returns the accumulated bitmap without copying. The union never
modifies that instance again: its next
+/// mutating call first copies it, so the returned bitmap stays valid.
Callers must treat it as read-only while they
+/// still intend to use the union.
+/// - [#take()] transfers ownership of the accumulated bitmap and leaves the
union empty.
+///
+/// Differences from the library class:
+///
+/// - The lazy primitives of [RoaringBitmap] are `protected`, so the
accumulated state is a private subclass that
+/// reaches them through inheritance. The bitmaps handed out by [#get()] and
[#take()] are instances of that
+/// subclass; it adds no state and overrides nothing.
+/// - Bitmaps are deserialized straight into a union with
[RoaringBitmapUtils#deserializeToUnion], and copies are
+/// made by adding to an empty union.
+/// - [#add(int)] repairs pending lazy state before inserting, because
inserting into a lazy container needs the
+/// library's internals, and it does not re-encode a run container that
received the value. A given accumulator
+/// receives either bitmaps or single values in Pinot, never both.
+/// - The released lazy union only pays off once containers are dense: while
they are small arrays it allocates a new
+/// container per union where the eager union merges in place, and it
inserts container keys that are new to the
+/// accumulator one at a time, which is quadratic when an input brings many
of them. So an input is unioned eagerly
+/// while the accumulator holds few values per container (hashed values
spread over many containers stay there for
+/// good), when overlapping arrays would stay small, or when the input would
insert many new keys, and lazily
+/// otherwise. The library class does not need this: its lazy union merges
small arrays in place and new keys in
+/// one pass.
+///
+/// Instances are not thread-safe.
+public final class RoaringBitmapUnion {
+ // The lazy union starts to pay off when containers hold more values than
the library's lazy array bound: past it
+ // they are kept as bitmaps whose unions skip cardinality maintenance, below
it they are arrays that the lazy union
+ // copies on every union
+ private static final int MIN_VALUES_PER_CONTAINER_FOR_LAZY_UNION = 1024;
+ // An input that would insert more new keys than this between existing ones
is unioned eagerly. Each insert shifts
+ // the accumulator's key and container arrays, and the eager union's single
merge pass costs about as much as a
+ // few of those shifts.
+ private static final int MAX_NEW_KEYS_FOR_LAZY_UNION = 4;
+
+ // The accumulated bitmap. Never null.
+ private LazyBitmap _bitmap;
+ // Whether the bitmap holds lazy state that must be repaired before it is
read
+ private boolean _dirty;
+ // Whether an alias returned by get() is outstanding and must not be mutated
+ private boolean _published;
+ // Number of values added so far, counting duplicates: an upper bound of the
cardinality that is known without
+ // repairing, used to estimate how dense the containers are
+ private long _numValuesAdded;
+ // A cardinality observed on repaired state. Unions never remove values, so
this remains a lower bound used to
+ // check average density without repairing every lazy union. Overlapping
arrays are checked separately.
+ private long _lastKnownCardinality;
+
+ /// Creates an empty union.
+ public RoaringBitmapUnion() {
+ _bitmap = new LazyBitmap();
+ }
+
+ private RoaringBitmapUnion(LazyBitmap adopted) {
+ _bitmap = adopted;
+ // Adopted state is normalized on the first read, like the library class
does
+ _dirty = true;
+ _numValuesAdded = adopted.getLongCardinality();
+ _lastKnownCardinality = _numValuesAdded;
+ }
+
+ /// Deserializes a bitmap straight into a union that owns it. Not part of
the library class: callers go through
+ /// [RoaringBitmapUtils#deserializeToUnion].
+ static RoaringBitmapUnion deserialize(ByteBuffer byteBuffer) {
+ LazyBitmap bitmap = new LazyBitmap();
+ try {
+ bitmap.deserialize(byteBuffer);
+ } catch (IOException e) {
+ throw new RuntimeException("Caught exception while deserializing
RoaringBitmap", e);
+ }
+ return new RoaringBitmapUnion(bitmap);
+ }
+
+ /// Unions the input into the accumulated state. The input is neither
modified nor retained, so it may be reused or
+ /// mutated afterwards. Adding this union's own [#get()] result, or an empty
bitmap, is a no-op.
+ public void add(RoaringBitmap input) {
+ Objects.requireNonNull(input, "input");
+ if (input == _bitmap || input.isEmpty()) {
+ return;
+ }
+ beforeMutation();
+ boolean dense = isDenseForLazyUnion();
+ _numValuesAdded += input.getLongCardinality();
+ if (dense && !shouldUnionEagerly(input)) {
+ _dirty = true;
+ _bitmap.lazyOr(input);
+ } else {
+ repairIfDirty();
+ _bitmap.or(input);
+ }
+ }
+
+ /// Unions a single value, treated as unsigned, into the accumulated state.
+ public void add(int value) {
+ beforeMutation();
+ repairIfDirty();
+ _bitmap.add(value);
+ _numValuesAdded++;
+ }
+
+ /// Returns the accumulated bitmap with all pending lazy state repaired. The
returned instance is the union's
+ /// current state, not a copy; the union copies it before its next mutation,
so the returned bitmap stays valid.
+ /// Repeated calls without an intervening mutation return the same instance.
+ public RoaringBitmap get() {
+ repairIfDirty();
+ _published = true;
+ return _bitmap;
+ }
+
+ /// Returns the accumulated bitmap, repaired as by [#get()], and transfers
its ownership to the caller. The union is
+ /// empty afterwards and can be reused.
+ public RoaringBitmap take() {
+ repairIfDirty();
+ RoaringBitmap result = _bitmap;
+ _bitmap = new LazyBitmap();
+ _dirty = false;
+ _published = false;
+ _numValuesAdded = 0;
+ _lastKnownCardinality = 0;
+ return result;
+ }
+
+ /// Returns whether the lazy union would copy small overlapping arrays or
insert more than a few new keys between
+ /// existing ones. Keys past the accumulator's last key do not count: they
are appended in one step.
+ private boolean shouldUnionEagerly(RoaringBitmap input) {
+ // Called after the average density check: an empty or sole dense
container has no small-array overlaps.
+ // A small input cannot insert too many keys. Avoid allocating scan
pointers for this common case.
+ if (_bitmap.getContainerCount() <= 1 && input.getContainerCount() <=
MAX_NEW_KEYS_FOR_LAZY_UNION) {
+ return false;
+ }
+ ContainerPointer own = _bitmap.getContainerPointer();
+ ContainerPointer other = input.getContainerPointer();
+ int numNewKeys = 0;
+ while (other.getContainer() != null) {
+ char key = other.key();
+ while (own.getContainer() != null && own.key() < key) {
+ own.advance();
+ }
+ if (own.getContainer() == null) {
+ return false;
+ }
+ if (own.key() == key) {
+ if (!own.isBitmapContainer() && !own.isRunContainer() &&
!other.isBitmapContainer() && !other.isRunContainer()
+ && own.getCardinality() + other.getCardinality() <=
MIN_VALUES_PER_CONTAINER_FOR_LAZY_UNION) {
+ return true;
+ }
+ } else if (++numNewKeys > MAX_NEW_KEYS_FOR_LAZY_UNION) {
+ return true;
+ }
+ other.advance();
+ }
+ return false;
+ }
+
+ private boolean isDenseForLazyUnion() {
+ long minCardinality = (long) MIN_VALUES_PER_CONTAINER_FOR_LAZY_UNION *
_bitmap.getContainerCount();
+ if (_lastKnownCardinality >= minCardinality) {
+ return true;
+ }
+ if (_numValuesAdded < minCardinality) {
+ return false;
+ }
+ // The estimate counts duplicates. Confirm actual density before switching
a sparse accumulator to lazy unions.
+ repairIfDirty();
+ _lastKnownCardinality = _bitmap.getLongCardinality();
+ _numValuesAdded = _lastKnownCardinality;
+ return _lastKnownCardinality >= minCardinality;
+ }
+
+ private void repairIfDirty() {
+ if (_dirty) {
+ _bitmap.repair();
+ _dirty = false;
+ }
+ }
+
+ private void beforeMutation() {
+ if (_published) {
+ _bitmap = (LazyBitmap) _bitmap.clone();
+ _published = false;
+ }
+ }
+
+ /// Reaches the protected lazy primitives of [RoaringBitmap] through
inheritance.
+ private static final class LazyBitmap extends RoaringBitmap {
+ // Public so that an instance that escaped through get() or take() stays
deserializable like a plain bitmap
+ public LazyBitmap() {
+ }
+
+ void lazyOr(RoaringBitmap other) {
+ lazyor(other);
+ }
+
+ void repair() {
+ repairAfterLazy();
+ }
+ }
+}
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/utils/RoaringBitmapUtils.java
b/pinot-common/src/main/java/org/apache/pinot/common/utils/RoaringBitmapUtils.java
index 59b63760103..e37814d5d3a 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/utils/RoaringBitmapUtils.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/utils/RoaringBitmapUtils.java
@@ -54,6 +54,14 @@ public class RoaringBitmapUtils {
return bitmap;
}
+ /// Deserializes a bitmap into a [RoaringBitmapUnion] that owns it, without
copying it.
+ ///
+ /// Once RoaringBitmap ships its own `RoaringBitmapUnion`, this becomes
+ /// `RoaringBitmapUnion.takeOwnership(deserialize(bytes))`.
+ public static RoaringBitmapUnion deserializeToUnion(byte[] bytes) {
+ return RoaringBitmapUnion.deserialize(ByteBuffer.wrap(bytes));
+ }
+
/// Iterates over the ranges of unset bits and calls the consumer for each
range. This is more performant to
/// alternatives like calling [RoaringBitmap#contains(int)] in a loop or
cloning and flipping the bitmap before
/// iterating, especially for sparse bitmaps.
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/utils/MutableRoaringBitmapUnionTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/utils/MutableRoaringBitmapUnionTest.java
new file mode 100644
index 00000000000..7f6f036fe3c
--- /dev/null
+++
b/pinot-common/src/test/java/org/apache/pinot/common/utils/MutableRoaringBitmapUnionTest.java
@@ -0,0 +1,299 @@
+/**
+ * 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.pinot.common.utils;
+
+import java.lang.reflect.Field;
+import java.nio.ByteBuffer;
+import java.util.Random;
+import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+import org.roaringbitmap.buffer.MappeableContainer;
+import org.roaringbitmap.buffer.MappeableContainerPointer;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+
+
+/// Covers the interim [MutableRoaringBitmapUnion] with the inputs Pinot feeds
it: read-only bitmaps over heap and
+/// direct buffers, as the index readers hand out, next to plain mutable
bitmaps.
+public class MutableRoaringBitmapUnionTest {
+ private static final long SEED = 20260929L;
+
+ private static MutableRoaringBitmap[] inputs(Random random, int numInputs) {
+ MutableRoaringBitmap[] inputs = new MutableRoaringBitmap[numInputs];
+ for (int i = 0; i < numInputs; i++) {
+ MutableRoaringBitmap bitmap = new MutableRoaringBitmap();
+ int base = random.nextInt(8) << 16;
+ switch (i % 3) {
+ case 0:
+ for (int j = 0; j < 300; j++) {
+ bitmap.add(base | random.nextInt(1 << 16));
+ }
+ break;
+ case 1:
+ bitmap.add((long) base + random.nextInt(1000), (long) base + 10000);
+ break;
+ default:
+ bitmap.add((long) base + random.nextInt(1000), (long) base + 6000);
+ bitmap.runOptimize();
+ break;
+ }
+ inputs[i] = bitmap;
+ }
+ return inputs;
+ }
+
+ /// The input as the readers present it: every third one over a read-only
heap buffer, every third one over a
+ /// read-only direct buffer, the rest as is.
+ private static ImmutableRoaringBitmap present(MutableRoaringBitmap bitmap,
int index) {
+ if (index % 3 == 0) {
+ return bitmap;
+ }
+ int size = bitmap.serializedSizeInBytes();
+ ByteBuffer buffer = index % 3 == 1 ? ByteBuffer.allocate(size) :
ByteBuffer.allocateDirect(size);
+ bitmap.serialize(buffer);
+ buffer.flip();
+ return new ImmutableRoaringBitmap(buffer.asReadOnlyBuffer());
+ }
+
+ private static MutableRoaringBitmap eagerUnion(MutableRoaringBitmap[]
inputs, int numInputs) {
+ MutableRoaringBitmap expected = new MutableRoaringBitmap();
+ for (int i = 0; i < numInputs; i++) {
+ expected.or(inputs[i]);
+ }
+ return expected;
+ }
+
+ private static void assertValid(MutableRoaringBitmap actual,
MutableRoaringBitmap expected) {
+ // A bitmap left in lazy state reports wrong cardinalities and fails
validation
+ assertTrue(actual.validate());
+ assertTrue(actual.equals(expected));
+ assertEquals(actual.getLongCardinality(), expected.getLongCardinality());
+ ByteBuffer buffer = ByteBuffer.allocate(actual.serializedSizeInBytes());
+ actual.serialize(buffer);
+ buffer.flip();
+ assertTrue(new ImmutableRoaringBitmap(buffer).equals(expected));
+ }
+
+ @Test
+ public void testFoldOverMappedInputsMatchesEagerUnion() {
+ Random random = new Random(SEED);
+ MutableRoaringBitmap[] inputs = inputs(random, 60);
+ MutableRoaringBitmapUnion union = new MutableRoaringBitmapUnion();
+ for (int i = 0; i < inputs.length; i++) {
+ MutableRoaringBitmap copy = inputs[i].clone();
+ ImmutableRoaringBitmap presented = present(inputs[i], i);
+ union.add(presented);
+ if (i % 7 == 0) {
+ union.add(presented);
+ union.add(new MutableRoaringBitmap());
+ union.add(union.get());
+ }
+ assertTrue(presented.equals(copy));
+ if (i % 10 == 9) {
+ assertValid(union.get(), eagerUnion(inputs, i + 1));
+ }
+ }
+ assertValid(union.take(), eagerUnion(inputs, inputs.length));
+ }
+
+ @Test
+ public void testGetIsStableAndTakeResets() {
+ Random random = new Random(SEED + 1);
+ MutableRoaringBitmap[] inputs = inputs(random, 20);
+ MutableRoaringBitmapUnion union = new MutableRoaringBitmapUnion();
+ for (int i = 0; i < 10; i++) {
+ union.add(present(inputs[i], i));
+ }
+ MutableRoaringBitmap published = union.get();
+ assertSame(union.get(), published);
+ MutableRoaringBitmap snapshot = published.clone();
+ for (int i = 10; i < 20; i++) {
+ union.add(present(inputs[i], i));
+ }
+ assertValid(published, snapshot);
+
+ MutableRoaringBitmap expected = eagerUnion(inputs, 20);
+ MutableRoaringBitmap taken = union.take();
+ assertNotSame(taken, published);
+ assertValid(taken, expected);
+ assertTrue(union.take().isEmpty());
+ }
+
+ @Test
+ public void testMappedInputsWithManyNewKeys() {
+ Random random = new Random(SEED + 3);
+ MutableRoaringBitmap[] dense = inputs(random, 30);
+ MutableRoaringBitmapUnion union = new MutableRoaringBitmapUnion();
+ MutableRoaringBitmap expected = new MutableRoaringBitmap();
+ for (int i = 0; i < dense.length; i++) {
+ // Hashed values: every input brings many container keys the accumulator
does not have yet
+ MutableRoaringBitmap hashed = new MutableRoaringBitmap();
+ for (int j = 0; j < 500; j++) {
+ hashed.add(random.nextInt());
+ }
+ union.add(present(dense[i], i));
+ union.add(present(hashed, i + 1));
+ expected.or(dense[i]);
+ expected.or(hashed);
+ if (i % 8 == 7) {
+ assertValid(union.get(), expected);
+ }
+ }
+ assertValid(union.take(), expected);
+ }
+
+ @Test
+ public void testSparseMappedInputsInterleavedWithDenseInputs() {
+ Random random = new Random(SEED + 4);
+ MutableRoaringBitmap[] dense = inputs(random, 80);
+ MutableRoaringBitmapUnion union = new MutableRoaringBitmapUnion();
+ MutableRoaringBitmap expected = new MutableRoaringBitmap();
+ for (int i = 0; i < dense.length; i++) {
+ union.add(present(dense[i], i));
+ expected.or(dense[i]);
+ if (i % 8 == 7) {
+ // A few values spread over new keys, added to an accumulator that
mostly holds dense containers
+ MutableRoaringBitmap sparse = new MutableRoaringBitmap();
+ for (int j = 0; j < 24; j++) {
+ sparse.add(random.nextInt());
+ }
+ union.add(present(sparse, i));
+ expected.or(sparse);
+ }
+ if (i % 16 == 15) {
+ assertValid(union.get(), expected);
+ }
+ }
+ assertValid(union.take(), expected);
+ }
+
+ // Inspect container identity without get(), whose publication would
legitimately trigger a copy on the next add.
+ private static MutableRoaringBitmap
accumulatedBitmap(MutableRoaringBitmapUnion union)
+ throws ReflectiveOperationException {
+ Field field = MutableRoaringBitmapUnion.class.getDeclaredField("_bitmap");
+ field.setAccessible(true);
+ return (MutableRoaringBitmap) field.get(union);
+ }
+
+ @Test
+ public void testRepeatedSparseMappedInputsStayInPlace()
+ throws ReflectiveOperationException {
+ MutableRoaringBitmap input = new MutableRoaringBitmap();
+ for (int key = 0; key < 24; key++) {
+ input.add(key << 16);
+ }
+ for (int presentation = 1; presentation <= 2; presentation++) {
+ ImmutableRoaringBitmap mapped = present(input, presentation);
+ MutableRoaringBitmapUnion union = new MutableRoaringBitmapUnion();
+ union.add(mapped);
+ MutableRoaringBitmap owned = accumulatedBitmap(union);
+ MappeableContainer[] containers = new
MappeableContainer[owned.getContainerCount()];
+ MappeableContainerPointer pointer = owned.getContainerPointer();
+ for (int i = 0; i < containers.length; i++) {
+ containers[i] = pointer.getContainer();
+ pointer.advance();
+ }
+ // Exercise read-only heap and direct buffers without publishing
intermediate accumulator state.
+ for (int i = 0; i < 2048; i++) {
+ union.add(mapped);
+ }
+ MutableRoaringBitmap result = union.take();
+ assertValid(result, input);
+ assertTrue(mapped.equals(input));
+ pointer = result.getContainerPointer();
+ for (MappeableContainer container : containers) {
+ assertSame(pointer.getContainer(), container);
+ pointer.advance();
+ }
+ }
+ }
+
+ @Test
+ public void testSparseMappedContainersStayInPlaceWithSkewedDensity()
+ throws ReflectiveOperationException {
+ for (int numSparseKeys : new int[]{1, 50}) {
+ MutableRoaringBitmap initial = new MutableRoaringBitmap();
+ initial.add(0L, 1L << 16);
+ for (int key = 1; key <= numSparseKeys; key++) {
+ initial.add(key << 16);
+ }
+ for (int presentation = 1; presentation <= 2; presentation++) {
+ MutableRoaringBitmapUnion union = new MutableRoaringBitmapUnion();
+ union.add(present(initial, presentation));
+ MutableRoaringBitmap owned = accumulatedBitmap(union);
+ MappeableContainerPointer pointer = owned.getContainerPointer();
+ pointer.advance();
+ MappeableContainer[] sparseContainers = new
MappeableContainer[numSparseKeys];
+ for (int i = 0; i < numSparseKeys; i++) {
+ sparseContainers[i] = pointer.getContainer();
+ pointer.advance();
+ }
+ MutableRoaringBitmap expected = initial.clone();
+ // The full first container masks the sparse arrays in the average
density.
+ for (int i = 1; i <= 500; i++) {
+ MutableRoaringBitmap input = new MutableRoaringBitmap();
+ for (int key = 1; key <= numSparseKeys; key++) {
+ input.add((key << 16) | (i * 2));
+ }
+ ImmutableRoaringBitmap mapped = present(input, presentation);
+ union.add(mapped);
+ expected.or(input);
+ assertTrue(mapped.equals(input));
+ }
+ MutableRoaringBitmap result = union.take();
+ assertValid(result, expected);
+ pointer = result.getContainerPointer();
+ pointer.advance();
+ for (MappeableContainer container : sparseContainers) {
+ assertSame(pointer.getContainer(), container);
+ pointer.advance();
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testDenseMappedInputsStillUseLazyUnion()
+ throws ReflectiveOperationException {
+ MutableRoaringBitmap input = new MutableRoaringBitmap();
+ for (int value = 0; value < 1024; value++) {
+ input.add(value);
+ }
+ MutableRoaringBitmapUnion union = new MutableRoaringBitmapUnion();
+ union.add(input);
+ MutableRoaringBitmap owned = accumulatedBitmap(union);
+ MappeableContainer container = owned.getContainerPointer().getContainer();
+ MutableRoaringBitmap duplicate = new MutableRoaringBitmap();
+ duplicate.add(0);
+ union.add(present(duplicate, 2));
+ MutableRoaringBitmap result = union.take();
+ assertValid(result, input);
+ assertNotSame(result.getContainerPointer().getContainer(), container);
+ }
+
+ @Test
+ public void testNullArguments() {
+ assertThrows(NullPointerException.class, () -> new
MutableRoaringBitmapUnion().add(null));
+ }
+}
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/utils/RoaringBitmapUnionTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/utils/RoaringBitmapUnionTest.java
new file mode 100644
index 00000000000..95f9643a8cf
--- /dev/null
+++
b/pinot-common/src/test/java/org/apache/pinot/common/utils/RoaringBitmapUnionTest.java
@@ -0,0 +1,399 @@
+/**
+ * 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.pinot.common.utils;
+
+import java.lang.reflect.Field;
+import java.util.Random;
+import org.roaringbitmap.Container;
+import org.roaringbitmap.ContainerPointer;
+import org.roaringbitmap.RoaringBitmap;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
+
+
+/// Covers the interim [RoaringBitmapUnion]: folds are compared with eager
unions, and the ownership rules of
+/// `get()` and `take()` are pinned so that the swap to the library class
cannot change behavior unnoticed.
+public class RoaringBitmapUnionTest {
+ private static final long SEED = 20260929L;
+
+ /// Inputs mixing sparse values, dense ranges and run-optimized ranges, so
that array, bitmap and run containers
+ /// all take part, including keys in the upper (unsigned) half of the range.
+ private static RoaringBitmap[] inputs(Random random, int numInputs) {
+ RoaringBitmap[] inputs = new RoaringBitmap[numInputs];
+ for (int i = 0; i < numInputs; i++) {
+ RoaringBitmap bitmap = new RoaringBitmap();
+ int base = random.nextInt(8) << 16;
+ if (random.nextBoolean()) {
+ base |= 0x80000000;
+ }
+ switch (i % 3) {
+ case 0:
+ for (int j = 0; j < 300; j++) {
+ bitmap.add(base | random.nextInt(1 << 16));
+ }
+ break;
+ case 1:
+ long start = (base & 0xFFFFFFFFL) + random.nextInt(1000);
+ bitmap.add(start, start + 9000);
+ break;
+ default:
+ long runStart = (base & 0xFFFFFFFFL) + random.nextInt(1000);
+ bitmap.add(runStart, runStart + 5000);
+ bitmap.runOptimize();
+ break;
+ }
+ inputs[i] = bitmap;
+ }
+ return inputs;
+ }
+
+ private static RoaringBitmap eagerUnion(RoaringBitmap[] inputs, int
numInputs) {
+ RoaringBitmap expected = new RoaringBitmap();
+ for (int i = 0; i < numInputs; i++) {
+ expected.or(inputs[i]);
+ }
+ return expected;
+ }
+
+ private static void assertValid(RoaringBitmap actual, RoaringBitmap
expected) {
+ // A bitmap left in lazy state reports wrong cardinalities and fails
validation
+ assertTrue(actual.validate());
+ assertTrue(actual.equals(expected));
+ assertEquals(actual.getLongCardinality(), expected.getLongCardinality());
+
assertTrue(RoaringBitmapUtils.deserialize(RoaringBitmapUtils.serialize(actual)).equals(expected));
+ }
+
+ /// Bitmaps of hashed values: a few hundred values spread over the whole
unsigned range, so that every input brings
+ /// many container keys the accumulator does not have yet.
+ private static RoaringBitmap[] hashedInputs(Random random, int numInputs,
int valuesPerInput) {
+ RoaringBitmap[] inputs = new RoaringBitmap[numInputs];
+ for (int i = 0; i < numInputs; i++) {
+ RoaringBitmap bitmap = new RoaringBitmap();
+ for (int j = 0; j < valuesPerInput; j++) {
+ bitmap.add(random.nextInt());
+ }
+ inputs[i] = bitmap;
+ }
+ return inputs;
+ }
+
+ @Test
+ public void testFoldMatchesEagerUnionAndLeavesInputsUntouched() {
+ Random random = new Random(SEED);
+ RoaringBitmap[] inputs = inputs(random, 60);
+ RoaringBitmap[] copies = new RoaringBitmap[inputs.length];
+ RoaringBitmapUnion union = new RoaringBitmapUnion();
+ for (int i = 0; i < inputs.length; i++) {
+ copies[i] = inputs[i].clone();
+ union.add(inputs[i]);
+ // The same input twice, an empty input and the union's own result never
change the outcome
+ if (i % 7 == 0) {
+ union.add(inputs[i]);
+ union.add(new RoaringBitmap());
+ union.add(union.get());
+ }
+ if (i % 10 == 9) {
+ assertValid(union.get(), eagerUnion(inputs, i + 1));
+ }
+ }
+ assertValid(union.take(), eagerUnion(inputs, inputs.length));
+ for (int i = 0; i < inputs.length; i++) {
+ assertTrue(inputs[i].equals(copies[i]));
+ }
+ }
+
+ @Test
+ public void testGetIsStableAndNeverInvalidatedByLaterAdds() {
+ Random random = new Random(SEED + 1);
+ RoaringBitmap[] inputs = inputs(random, 20);
+ RoaringBitmapUnion union = new RoaringBitmapUnion();
+ for (int i = 0; i < 10; i++) {
+ union.add(inputs[i]);
+ }
+ RoaringBitmap published = union.get();
+ assertSame(union.get(), published);
+ RoaringBitmap snapshot = published.clone();
+
+ for (int i = 10; i < 20; i++) {
+ union.add(inputs[i]);
+ }
+ union.add(12345);
+ // The published bitmap was copied before the union changed again
+ assertValid(published, snapshot);
+ RoaringBitmap expected = eagerUnion(inputs, 20);
+ expected.add(12345);
+ RoaringBitmap latest = union.get();
+ assertNotSame(latest, published);
+ assertValid(latest, expected);
+ }
+
+ @Test
+ public void testTakeTransfersOwnershipAndResets() {
+ Random random = new Random(SEED + 2);
+ RoaringBitmap[] inputs = inputs(random, 10);
+ RoaringBitmapUnion union = new RoaringBitmapUnion();
+ for (RoaringBitmap input : inputs) {
+ union.add(input);
+ }
+ RoaringBitmap published = union.get();
+ RoaringBitmap taken = union.take();
+ assertSame(taken, published);
+ RoaringBitmap snapshot = taken.clone();
+
+ assertTrue(union.get().isEmpty());
+ union.add(inputs[0]);
+ union.add(7);
+ assertValid(taken, snapshot);
+ RoaringBitmap expected = inputs[0].clone();
+ expected.add(7);
+ assertValid(union.take(), expected);
+ assertTrue(union.take().isEmpty());
+ }
+
+ @Test
+ public void testDeserializeToUnionKeepsAccumulating() {
+ Random random = new Random(SEED + 3);
+ RoaringBitmap[] inputs = inputs(random, 12);
+ RoaringBitmapUnion union =
RoaringBitmapUtils.deserializeToUnion(RoaringBitmapUtils.serialize(inputs[0]));
+ assertValid(union.get(), inputs[0]);
+ for (int i = 1; i < inputs.length; i++) {
+ union.add(inputs[i]);
+ }
+ RoaringBitmap result = union.take();
+ assertValid(result, eagerUnion(inputs, inputs.length));
+ }
+
+ @Test
+ public void testInputsWithManyNewKeysInterleavedWithDenseInputs() {
+ Random random = new Random(SEED + 5);
+ RoaringBitmap[] hashed = hashedInputs(random, 40, 500);
+ RoaringBitmap[] dense = inputs(random, 40);
+ RoaringBitmapUnion union = new RoaringBitmapUnion();
+ RoaringBitmap expected = new RoaringBitmap();
+ for (int i = 0; i < hashed.length; i++) {
+ RoaringBitmap hashedCopy = hashed[i].clone();
+ // Inputs that bring many new keys and inputs that bring few alternate,
so pending lazy state has to be
+ // repaired before the accumulator is unioned eagerly
+ union.add(dense[i]);
+ union.add(hashed[i]);
+ expected.or(dense[i]);
+ expected.or(hashed[i]);
+ assertTrue(hashed[i].equals(hashedCopy));
+ if (i % 8 == 7) {
+ assertValid(union.get(), expected);
+ }
+ }
+ // Once the accumulator has most keys, the same inputs bring few new ones
+ for (RoaringBitmap input : hashed) {
+ union.add(input);
+ }
+ assertValid(union.take(), expected);
+ }
+
+ @Test
+ public void testSparseInputsInterleavedWithDenseInputs() {
+ Random random = new Random(SEED + 6);
+ RoaringBitmap[] dense = inputs(random, 80);
+ RoaringBitmap[] sparse = hashedInputs(random, 10, 24);
+ RoaringBitmapUnion union = new RoaringBitmapUnion();
+ RoaringBitmap expected = new RoaringBitmap();
+ for (int i = 0; i < dense.length; i++) {
+ union.add(dense[i]);
+ expected.or(dense[i]);
+ if (i % 8 == 7) {
+ // A few values spread over new keys, added to an accumulator that
mostly holds dense containers
+ union.add(sparse[i / 8]);
+ expected.or(sparse[i / 8]);
+ }
+ if (i % 16 == 15) {
+ assertValid(union.get(), expected);
+ }
+ }
+ // Enough hashed values to leave the accumulator with mostly sparse
containers, then dense inputs again
+ for (RoaringBitmap input : hashedInputs(random, 40, 500)) {
+ union.add(input);
+ expected.or(input);
+ }
+ assertValid(union.get(), expected);
+ for (RoaringBitmap input : inputs(random, 20)) {
+ union.add(input);
+ expected.or(input);
+ }
+ assertValid(union.take(), expected);
+ }
+
+ @Test
+ public void testSingleValuesMixedWithBitmaps() {
+ Random random = new Random(SEED + 4);
+ RoaringBitmap[] inputs = inputs(random, 30);
+ RoaringBitmapUnion union = new RoaringBitmapUnion();
+ RoaringBitmap expected = new RoaringBitmap();
+ for (int i = 0; i < inputs.length; i++) {
+ union.add(inputs[i]);
+ expected.or(inputs[i]);
+ // Single values landing in containers the previous lazy union just
touched, and in new ones
+ for (int j = 0; j < 5; j++) {
+ int value = random.nextBoolean() ? inputs[i].first() + j :
random.nextInt();
+ union.add(value);
+ expected.add(value);
+ }
+ }
+ union.add(-1);
+ expected.add(-1);
+ assertValid(union.take(), expected);
+ }
+
+ @Test
+ public void testRawValuesOnlyCrossTheArrayToBitmapThreshold() {
+ RoaringBitmapUnion union = new RoaringBitmapUnion();
+ RoaringBitmap expected = new RoaringBitmap();
+ for (int i = 0; i < 5000; i++) {
+ union.add(3 * i);
+ expected.add(3 * i);
+ }
+ assertValid(union.get(), expected);
+ }
+
+ // Inspect container identity without get(), whose publication would
legitimately trigger a copy on the next add.
+ private static RoaringBitmap accumulatedBitmap(RoaringBitmapUnion union)
+ throws ReflectiveOperationException {
+ Field field = RoaringBitmapUnion.class.getDeclaredField("_bitmap");
+ field.setAccessible(true);
+ return (RoaringBitmap) field.get(union);
+ }
+
+ @Test
+ public void testRepeatedSparseInputsStayInPlace()
+ throws ReflectiveOperationException {
+ RoaringBitmap input = new RoaringBitmap();
+ for (int key = 0; key < 24; key++) {
+ input.add(key << 16);
+ }
+ RoaringBitmapUnion union = new RoaringBitmapUnion();
+ union.add(input);
+ RoaringBitmap owned = accumulatedBitmap(union);
+ Container[] containers = new Container[owned.getContainerCount()];
+ ContainerPointer pointer = owned.getContainerPointer();
+ for (int i = 0; i < containers.length; i++) {
+ containers[i] = pointer.getContainer();
+ pointer.advance();
+ }
+ // Duplicate counts exceed the lazy threshold, but the result remains one
value per container.
+ for (int i = 0; i < 2048; i++) {
+ union.add(input);
+ }
+ RoaringBitmap result = union.take();
+ assertValid(result, input);
+ pointer = result.getContainerPointer();
+ for (Container container : containers) {
+ // Eager unions reuse these sparse arrays; lazy unions replace them even
when no values change.
+ assertSame(pointer.getContainer(), container);
+ pointer.advance();
+ }
+ }
+
+ @Test
+ public void testSparseContainersStayInPlaceWithSkewedDensity()
+ throws ReflectiveOperationException {
+ for (int numSparseKeys : new int[]{1, 50}) {
+ RoaringBitmap initial = new RoaringBitmap();
+ initial.add(0L, 1L << 16);
+ for (int key = 1; key <= numSparseKeys; key++) {
+ initial.add(key << 16);
+ }
+ RoaringBitmapUnion union = new RoaringBitmapUnion();
+ union.add(initial);
+ RoaringBitmap owned = accumulatedBitmap(union);
+ ContainerPointer pointer = owned.getContainerPointer();
+ pointer.advance();
+ Container[] sparseContainers = new Container[numSparseKeys];
+ for (int i = 0; i < numSparseKeys; i++) {
+ sparseContainers[i] = pointer.getContainer();
+ pointer.advance();
+ }
+ RoaringBitmap expected = initial.clone();
+ // The full first container masks the sparse arrays in the average
density, even with 50 sparse keys.
+ for (int i = 1; i <= 500; i++) {
+ RoaringBitmap input = new RoaringBitmap();
+ for (int key = 1; key <= numSparseKeys; key++) {
+ input.add((key << 16) | (i * 2));
+ }
+ union.add(input);
+ expected.or(input);
+ }
+ RoaringBitmap result = union.take();
+ assertValid(result, expected);
+ pointer = result.getContainerPointer();
+ pointer.advance();
+ for (Container container : sparseContainers) {
+ // Eager unions grow these arrays in place; lazy unions allocate a
replacement on every input.
+ assertSame(pointer.getContainer(), container);
+ pointer.advance();
+ }
+ }
+ }
+
+ @Test
+ public void testDenseInputsStillUseLazyUnion()
+ throws ReflectiveOperationException {
+ RoaringBitmap input = new RoaringBitmap();
+ for (int value = 0; value < 1024; value++) {
+ input.add(value);
+ }
+ RoaringBitmapUnion union = new RoaringBitmapUnion();
+ union.add(input);
+ RoaringBitmap owned = accumulatedBitmap(union);
+ Container container = owned.getContainerPointer().getContainer();
+ union.add(RoaringBitmap.bitmapOf(0));
+ RoaringBitmap result = union.take();
+ assertValid(result, input);
+ // Lazy union promotes the dense array and repair replaces it; eager union
would keep the array.
+ assertNotSame(result.getContainerPointer().getContainer(), container);
+ }
+
+ @Test
+ public void testNullArguments() {
+ assertThrows(NullPointerException.class, () -> new
RoaringBitmapUnion().add((RoaringBitmap) null));
+ }
+
+ /// Fails once the RoaringBitmap library on the classpath ships the classes
these interim ones stand in for, with
+ /// the steps of the swap in the failure message.
+ @Test
+ public void testLibraryDoesNotShipTheUnionClassesYet() {
+ for (String libraryClass : new String[]{
+ "org.roaringbitmap.RoaringBitmapUnion",
"org.roaringbitmap.buffer.MutableRoaringBitmapUnion"
+ }) {
+ try {
+ Class.forName(libraryClass);
+ } catch (ClassNotFoundException e) {
+ continue;
+ }
+ fail("RoaringBitmap now ships " + libraryClass + ". Replace Pinot's
interim copies: import the union classes "
+ + "from org.roaringbitmap and org.roaringbitmap.buffer at the call
sites, make "
+ + "RoaringBitmapUtils.deserializeToUnion call
RoaringBitmapUnion.takeOwnership(deserialize(bytes)), and "
+ + "delete org.apache.pinot.common.utils.RoaringBitmapUnion,
MutableRoaringBitmapUnion and their two tests");
+ }
+ }
+}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/InvertedIndexFilterOperator.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/InvertedIndexFilterOperator.java
index f4555a0b408..2a83176e2eb 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/InvertedIndexFilterOperator.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/InvertedIndexFilterOperator.java
@@ -21,6 +21,7 @@ package org.apache.pinot.core.operator.filter;
import com.google.common.base.CaseFormat;
import java.util.List;
import org.apache.pinot.common.request.context.predicate.Predicate;
+import org.apache.pinot.common.utils.MutableRoaringBitmapUnion;
import org.apache.pinot.core.common.BlockDocIdSet;
import org.apache.pinot.core.common.Operator;
import org.apache.pinot.core.operator.ExplainAttributeBuilder;
@@ -145,11 +146,13 @@ public class InvertedIndexFilterOperator extends
BaseColumnFilterOperator {
_invertedIndexReader.getDocIds(dictIds[1]));
break;
default:
- MutableRoaringBitmap bitmap = new MutableRoaringBitmap();
+ // Union lazily and finalize once at the end: the fold still
materializes the union, but skips per-union
+ // cardinality maintenance
+ MutableRoaringBitmapUnion union = new MutableRoaringBitmapUnion();
for (int dictId : dictIds) {
- bitmap.or(_invertedIndexReader.getDocIds(dictId));
+ union.add(_invertedIndexReader.getDocIds(dictId));
}
- count = bitmap.getCardinality();
+ count = union.get().getCardinality();
break;
}
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java
index 25fb0ea5029..160fd83d5e9 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java
@@ -25,6 +25,7 @@ import javax.annotation.Nullable;
import org.apache.pinot.common.CustomObject;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.common.utils.RoaringBitmapUnion;
import org.apache.pinot.common.utils.RoaringBitmapUtils;
import org.apache.pinot.core.common.BlockValSet;
import org.apache.pinot.core.common.ObjectSerDeUtils;
@@ -78,19 +79,14 @@ public class DistinctCountBitmapAggregationFunction extends
BaseSingleInputAggre
// Logical BYTES is a serialized RoaringBitmap and always uses the
single-value representation.
byte[][] bytesValues = blockValSet.getBytesValuesSV();
forEachNotNull(length, blockValSet, (from, to) -> {
- int i = from;
- RoaringBitmap valueBitmap = aggregationResultHolder.getResult();
- if (valueBitmap == null) {
- if (i == to) {
- return;
- }
- // The first bitmap read becomes the accumulator instead of being
merged into a fresh one
- valueBitmap = RoaringBitmapUtils.deserialize(bytesValues[i++]);
- aggregationResultHolder.setValue(valueBitmap);
- }
- for (; i < to; i++) {
- valueBitmap.or(RoaringBitmapUtils.deserialize(bytesValues[i]));
+ Object result = aggregationResultHolder.getResult();
+ for (int i = from; i < to; i++) {
+ // The first bitmap read is deserialized straight into the
accumulator instead of being merged into a fresh
+ // one; later ones are unioned lazily, and
extractAggregationResult() finalizes the accumulator exactly once
+ result = result == null ?
RoaringBitmapUtils.deserializeToUnion(bytesValues[i])
+ : addBitmap(result,
RoaringBitmapUtils.deserialize(bytesValues[i]));
}
+ aggregationResultHolder.setValue(result);
});
return;
}
@@ -271,14 +267,11 @@ public class DistinctCountBitmapAggregationFunction
extends BaseSingleInputAggre
byte[][] bytesValues = blockValSet.getBytesValuesSV();
forEachNotNull(length, blockValSet, (from, to) -> {
for (int i = from; i < to; i++) {
- RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]);
int groupKey = groupKeyArray[i];
- RoaringBitmap valueBitmap = groupByResultHolder.getResult(groupKey);
- if (valueBitmap != null) {
- valueBitmap.or(value);
- } else {
- groupByResultHolder.setValueForKey(groupKey, value);
- }
+ Object result = groupByResultHolder.getResult(groupKey);
+ groupByResultHolder.setValueForKey(groupKey,
+ result == null ?
RoaringBitmapUtils.deserializeToUnion(bytesValues[i])
+ : addBitmap(result,
RoaringBitmapUtils.deserialize(bytesValues[i])));
}
});
return;
@@ -461,12 +454,12 @@ public class DistinctCountBitmapAggregationFunction
extends BaseSingleInputAggre
for (int i = from; i < to; i++) {
RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]);
for (int groupKey : groupKeysArray[i]) {
- RoaringBitmap bitmap = groupByResultHolder.getResult(groupKey);
- if (bitmap != null) {
- bitmap.or(value);
+ Object result = groupByResultHolder.getResult(groupKey);
+ if (result != null) {
+ groupByResultHolder.setValueForKey(groupKey, addBitmap(result,
value));
} else {
- // Clone a bitmap for the group
- groupByResultHolder.setValueForKey(groupKey, value.clone());
+ // Each group owns an independent accumulator: the same input is
added to every group of the row
+ groupByResultHolder.setValueForKey(groupKey, addBitmap(new
RoaringBitmapUnion(), value));
}
}
}
@@ -661,10 +654,16 @@ public class DistinctCountBitmapAggregationFunction
extends BaseSingleInputAggre
if (result instanceof DictIdsWrapper) {
// For dictionary-encoded expression, convert dictionary ids to hash
code of the values
return convertToValueBitmap((DictIdsWrapper) result);
- } else {
- // For serialized RoaringBitmap and non-dictionary-encoded expression,
directly return the value bitmap
- return (RoaringBitmap) result;
}
+ if (result instanceof RoaringBitmapUnion) {
+ // Serialized RoaringBitmap expression: finalize the lazy accumulator
once and publish the plain bitmap, so
+ // repeated extraction returns the same instance and the union no longer
references it
+ RoaringBitmap bitmap = ((RoaringBitmapUnion) result).take();
+ aggregationResultHolder.setValue(bitmap);
+ return bitmap;
+ }
+ // For non-dictionary-encoded expression (or an already published bitmap),
directly return the value bitmap
+ return (RoaringBitmap) result;
}
@Override
@@ -677,10 +676,31 @@ public class DistinctCountBitmapAggregationFunction
extends BaseSingleInputAggre
if (result instanceof DictIdsWrapper) {
// For dictionary-encoded expression, convert dictionary ids to hash
code of the values
return convertToValueBitmap((DictIdsWrapper) result);
+ }
+ if (result instanceof RoaringBitmapUnion) {
+ // See extractAggregationResult()
+ RoaringBitmap bitmap = ((RoaringBitmapUnion) result).take();
+ groupByResultHolder.setValueForKey(groupKey, bitmap);
+ return bitmap;
+ }
+ // For non-dictionary-encoded expression (or an already published bitmap),
directly return the value bitmap
+ return (RoaringBitmap) result;
+ }
+
+ /// Unions a deserialized bitmap into the accumulator of a serialized-bitmap
aggregation. The accumulator is a
+ /// [RoaringBitmapUnion] while aggregating; a plain [RoaringBitmap] means
the result was already published by
+ /// extraction (no operator aggregates again after extracting today, so this
path only keeps a published bitmap
+ /// immutable), in which case a fresh accumulator is started from a copy of
it.
+ private static RoaringBitmapUnion addBitmap(Object result, RoaringBitmap
input) {
+ RoaringBitmapUnion union;
+ if (result instanceof RoaringBitmapUnion) {
+ union = (RoaringBitmapUnion) result;
} else {
- // For serialized RoaringBitmap and non-dictionary-encoded expression,
directly return the value bitmap
- return (RoaringBitmap) result;
+ union = new RoaringBitmapUnion();
+ union.add((RoaringBitmap) result);
}
+ union.add(input);
+ return union;
}
@Override
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapLazyUnionTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapLazyUnionTest.java
new file mode 100644
index 00000000000..84d846edba1
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapLazyUnionTest.java
@@ -0,0 +1,221 @@
+/**
+ * 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.pinot.core.query.aggregation.function;
+
+import java.util.Map;
+import java.util.Random;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.utils.RoaringBitmapUtils;
+import org.apache.pinot.core.common.BlockValSet;
+import org.apache.pinot.core.query.aggregation.AggregationResultHolder;
+import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.mockito.Mockito;
+import org.roaringbitmap.RoaringBitmap;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertSame;
+
+
+/// Covers the serialized-bitmap (`BYTES`) aggregation paths of
[DistinctCountBitmapAggregationFunction], which union
+/// input bitmaps lazily and repair the accumulator at extraction. The input
cardinalities are chosen to cross the
+/// container thresholds of the lazy union (array containers promote to bitmap
containers past 1024 combined
+/// cardinality; repair converts back to array containers at up to 4096), so
both promoted and non-promoted
+/// accumulator states are verified against an eagerly unioned reference.
+public class DistinctCountBitmapLazyUnionTest {
+ private static final ExpressionContext EXPRESSION =
ExpressionContext.forIdentifier("bitmapCol");
+ private static final Random RANDOM = new Random(42);
+
+ private static byte[][] serializedBitmaps(int numBitmaps, int
valuesPerBitmap, int maxValue) {
+ byte[][] serialized = new byte[numBitmaps][];
+ for (int i = 0; i < numBitmaps; i++) {
+ RoaringBitmap bitmap = new RoaringBitmap();
+ for (int j = 0; j < valuesPerBitmap; j++) {
+ bitmap.add(RANDOM.nextInt(maxValue));
+ }
+ serialized[i] = RoaringBitmapUtils.serialize(bitmap);
+ }
+ return serialized;
+ }
+
+ private static RoaringBitmap eagerUnion(byte[][] serialized, int from, int
to) {
+ RoaringBitmap expected = new RoaringBitmap();
+ for (int i = from; i < to; i++) {
+ expected.or(RoaringBitmapUtils.deserialize(serialized[i]));
+ }
+ return expected;
+ }
+
+ private static Map<ExpressionContext, BlockValSet>
mockBlockValSetMap(byte[][] serialized) {
+ return mockBlockValSetMap(serialized, null);
+ }
+
+ private static Map<ExpressionContext, BlockValSet>
mockBlockValSetMap(byte[][] serialized,
+ @Nullable RoaringBitmap nullBitmap) {
+ BlockValSet blockValSet = Mockito.mock(BlockValSet.class);
+ Mockito.when(blockValSet.getValueType()).thenReturn(DataType.BYTES);
+ Mockito.when(blockValSet.isSingleValue()).thenReturn(true);
+ Mockito.when(blockValSet.getBytesValuesSV()).thenReturn(serialized);
+ Mockito.when(blockValSet.getNullBitmap()).thenReturn(nullBitmap);
+ return Map.of(EXPRESSION, blockValSet);
+ }
+
+ @Test
+ public void testAggregateSkipsNullRowsAcrossRanges() {
+ // With null handling enabled the block is consumed as several non-null
ranges; the accumulator is re-read from
+ // the holder at the start of each range and written back at its end, and
null rows never reach the union
+ byte[][] serialized = serializedBitmaps(120, 50, 100_000);
+ RoaringBitmap nullBitmap = RoaringBitmap.bitmapOf(0, 1, 17, 40, 41, 42,
99, 119);
+ DistinctCountBitmapAggregationFunction function =
+ new DistinctCountBitmapAggregationFunction(EXPRESSION, true);
+ AggregationResultHolder holder = function.createAggregationResultHolder();
+ int blockSize = 60;
+ for (int from = 0; from < serialized.length; from += blockSize) {
+ byte[][] block = new byte[blockSize][];
+ System.arraycopy(serialized, from, block, 0, blockSize);
+ RoaringBitmap blockNulls = new RoaringBitmap();
+ for (int i = 0; i < blockSize; i++) {
+ if (nullBitmap.contains(from + i)) {
+ blockNulls.add(i);
+ }
+ }
+ function.aggregate(blockSize, holder, mockBlockValSetMap(block,
blockNulls));
+ }
+
+ RoaringBitmap expected = new RoaringBitmap();
+ for (int i = 0; i < serialized.length; i++) {
+ if (!nullBitmap.contains(i)) {
+ expected.or(RoaringBitmapUtils.deserialize(serialized[i]));
+ }
+ }
+ RoaringBitmap result = function.extractAggregationResult(holder);
+ assertEquals(result, expected);
+ assertEquals(result.getCardinality(), expected.getCardinality());
+ assertFalse(result.isEmpty());
+ }
+
+ @Test
+ public void testAggregateAcrossBlocks() {
+ // Enough overlap-heavy inputs to promote accumulator containers to (lazy)
bitmap containers, split into
+ // multiple aggregate() calls to verify the accumulator stays valid across
blocks until extraction. An early
+ // extraction after the first block guards the published-bitmap safety
net: no operator extracts and then keeps
+ // aggregating today, but if one did, the published bitmap must stay
immutable and later extraction must still
+ // be complete.
+ byte[][] serialized = serializedBitmaps(200, 100, 200_000);
+ DistinctCountBitmapAggregationFunction function =
+ new DistinctCountBitmapAggregationFunction(EXPRESSION, false);
+ AggregationResultHolder holder = function.createAggregationResultHolder();
+ int blockSize = 50;
+ RoaringBitmap published = null;
+ RoaringBitmap publishedCopy = null;
+ for (int from = 0; from < serialized.length; from += blockSize) {
+ byte[][] block = new byte[blockSize][];
+ System.arraycopy(serialized, from, block, 0, blockSize);
+ function.aggregate(blockSize, holder, mockBlockValSetMap(block));
+ if (published != null) {
+ assertEquals(published, publishedCopy);
+ assertEquals(published.getCardinality(),
publishedCopy.getCardinality());
+ }
+ if (from == 0) {
+ published = function.extractAggregationResult(holder);
+ assertEquals(published, eagerUnion(serialized, 0, blockSize));
+ // Repeated extraction returns the same published instance
+ assertSame(function.extractAggregationResult(holder), published);
+ publishedCopy = published.clone();
+ }
+ }
+
+ RoaringBitmap result = function.extractAggregationResult(holder);
+ RoaringBitmap expected = eagerUnion(serialized, 0, serialized.length);
+ assertEquals(result, expected);
+ assertEquals(result.getCardinality(), expected.getCardinality());
+ // The repaired accumulator must serialize into a form that round-trips
+
assertEquals(RoaringBitmapUtils.deserialize(RoaringBitmapUtils.serialize(result)),
expected);
+ }
+
+ @Test
+ public void testAggregateGroupBySV() {
+ int numGroups = 4;
+ byte[][] serialized = serializedBitmaps(400, 50, 100_000);
+ int[] groupKeys = new int[serialized.length];
+ for (int i = 0; i < groupKeys.length; i++) {
+ groupKeys[i] = i % numGroups;
+ }
+
+ DistinctCountBitmapAggregationFunction function =
+ new DistinctCountBitmapAggregationFunction(EXPRESSION, false);
+ GroupByResultHolder holder = function.createGroupByResultHolder(numGroups,
numGroups);
+ function.aggregateGroupBySV(serialized.length, groupKeys, holder,
mockBlockValSetMap(serialized));
+
+ for (int groupKey = 0; groupKey < numGroups; groupKey++) {
+ RoaringBitmap expected = new RoaringBitmap();
+ for (int i = groupKey; i < serialized.length; i += numGroups) {
+ expected.or(RoaringBitmapUtils.deserialize(serialized[i]));
+ }
+ RoaringBitmap result = function.extractGroupByResult(holder, groupKey);
+ assertEquals(result, expected, "group " + groupKey);
+ assertEquals(result.getCardinality(), expected.getCardinality(), "group
" + groupKey);
+ }
+ }
+
+ @Test
+ public void testAggregateGroupByMV() {
+ // Every row belongs to multiple groups, so the same deserialized input
bitmap is unioned into one group's
+ // accumulator and cloned as another group's initial accumulator —
verifies no aliasing between accumulators
+ int numGroups = 3;
+ byte[][] serialized = serializedBitmaps(150, 80, 150_000);
+ int[][] groupKeys = new int[serialized.length][];
+ for (int i = 0; i < groupKeys.length; i++) {
+ groupKeys[i] = new int[]{i % numGroups, (i + 1) % numGroups};
+ }
+
+ DistinctCountBitmapAggregationFunction function =
+ new DistinctCountBitmapAggregationFunction(EXPRESSION, false);
+ GroupByResultHolder holder = function.createGroupByResultHolder(numGroups,
numGroups);
+ function.aggregateGroupByMV(serialized.length, groupKeys, holder,
mockBlockValSetMap(serialized));
+
+ for (int groupKey = 0; groupKey < numGroups; groupKey++) {
+ RoaringBitmap expected = new RoaringBitmap();
+ for (int i = 0; i < serialized.length; i++) {
+ if (i % numGroups == groupKey || (i + 1) % numGroups == groupKey) {
+ expected.or(RoaringBitmapUtils.deserialize(serialized[i]));
+ }
+ }
+ RoaringBitmap result = function.extractGroupByResult(holder, groupKey);
+ assertEquals(result, expected, "group " + groupKey);
+ assertEquals(result.getCardinality(), expected.getCardinality(), "group
" + groupKey);
+ }
+ }
+
+ @Test
+ public void testSmallBitmapsStayCorrectBelowPromotionThreshold() {
+ // All containers stay small array containers: the lazy path must degrade
to plain array unions
+ byte[][] serialized = serializedBitmaps(50, 5, 1_000);
+ DistinctCountBitmapAggregationFunction function =
+ new DistinctCountBitmapAggregationFunction(EXPRESSION, false);
+ AggregationResultHolder holder = function.createAggregationResultHolder();
+ function.aggregate(serialized.length, holder,
mockBlockValSetMap(serialized));
+
+ RoaringBitmap expected = eagerUnion(serialized, 0, serialized.length);
+ assertEquals(function.extractAggregationResult(holder), expected);
+ }
+}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/BaseStarTreeV2Test.java
b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/BaseStarTreeV2Test.java
index 79ac5a7dd29..24dbabe15b0 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/BaseStarTreeV2Test.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/BaseStarTreeV2Test.java
@@ -179,11 +179,8 @@ abstract class BaseStarTreeV2Test<R, A> {
_valueAggregator.getAggregationType().getName(), null,
getCompressionCodec(),
true, getIndexVersion(), null, null)), MAX_LEAF_RECORDS);
File indexDir = new File(TEMP_DIR, SEGMENT_NAME);
- // Randomly build star-tree using on-heap or off-heap mode
- MultipleTreesBuilder.BuildMode buildMode =
- RANDOM.nextBoolean() ? MultipleTreesBuilder.BuildMode.ON_HEAP :
MultipleTreesBuilder.BuildMode.OFF_HEAP;
try (MultipleTreesBuilder builder = new
MultipleTreesBuilder(List.of(starTreeIndexConfig), false,
- indexDir, buildMode)) {
+ indexDir, getBuildMode())) {
builder.build();
}
@@ -494,6 +491,12 @@ abstract class BaseStarTreeV2Test<R, A> {
return compressionCodec.isApplicableToRawIndex() ? compressionCodec : null;
}
+ /// Star-tree build mode. Randomized by default so both builders are
exercised over time; tests that need a
+ /// specific builder override this.
+ MultipleTreesBuilder.BuildMode getBuildMode() {
+ return RANDOM.nextBoolean() ? MultipleTreesBuilder.BuildMode.ON_HEAP :
MultipleTreesBuilder.BuildMode.OFF_HEAP;
+ }
+
/// Can be overridden to force the index version.
@Nullable
Integer getIndexVersion() {
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountBitmapStarTreeV2Test.java
b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountBitmapStarTreeV2Test.java
index 044a9143837..ebc38c0c728 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountBitmapStarTreeV2Test.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountBitmapStarTreeV2Test.java
@@ -19,18 +19,19 @@
package org.apache.pinot.core.startree.v2;
import java.util.Random;
+import org.apache.pinot.common.utils.RoaringBitmapUnion;
import
org.apache.pinot.segment.local.aggregator.DistinctCountBitmapValueAggregator;
import org.apache.pinot.segment.local.aggregator.ValueAggregator;
+import org.apache.pinot.segment.local.startree.v2.builder.MultipleTreesBuilder;
import org.apache.pinot.spi.data.FieldSpec.DataType;
-import org.roaringbitmap.RoaringBitmap;
import static org.testng.Assert.assertEquals;
-public class DistinctCountBitmapStarTreeV2Test extends
BaseStarTreeV2Test<Object, RoaringBitmap> {
+public class DistinctCountBitmapStarTreeV2Test extends
BaseStarTreeV2Test<Object, RoaringBitmapUnion> {
@Override
- ValueAggregator<Object, RoaringBitmap> getValueAggregator() {
+ ValueAggregator<Object, RoaringBitmapUnion> getValueAggregator() {
return new DistinctCountBitmapValueAggregator();
}
@@ -45,7 +46,14 @@ public class DistinctCountBitmapStarTreeV2Test extends
BaseStarTreeV2Test<Object
}
@Override
- void assertAggregatedValue(RoaringBitmap starTreeResult, RoaringBitmap
nonStarTreeResult) {
- assertEquals(starTreeResult, nonStarTreeResult);
+ MultipleTreesBuilder.BuildMode getBuildMode() {
+ // Pin the on-heap builder so its path is always covered: it
pre-serializes the variable-length bitmap metric
+ // before sizing the forward index
+ return MultipleTreesBuilder.BuildMode.ON_HEAP;
+ }
+
+ @Override
+ void assertAggregatedValue(RoaringBitmapUnion starTreeResult,
RoaringBitmapUnion nonStarTreeResult) {
+ assertEquals(starTreeResult.get(), nonStarTreeResult.get());
}
}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountBitmapWithMVStarTreeV2Test.java
b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountBitmapWithMVStarTreeV2Test.java
index b58a9d31cc6..194526adf0a 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountBitmapWithMVStarTreeV2Test.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountBitmapWithMVStarTreeV2Test.java
@@ -19,18 +19,19 @@
package org.apache.pinot.core.startree.v2;
import java.util.Random;
+import org.apache.pinot.common.utils.RoaringBitmapUnion;
import
org.apache.pinot.segment.local.aggregator.DistinctCountBitmapValueAggregator;
import org.apache.pinot.segment.local.aggregator.ValueAggregator;
+import org.apache.pinot.segment.local.startree.v2.builder.MultipleTreesBuilder;
import org.apache.pinot.spi.data.FieldSpec.DataType;
-import org.roaringbitmap.RoaringBitmap;
import static org.testng.Assert.assertEquals;
-public class DistinctCountBitmapWithMVStarTreeV2Test extends
BaseStarTreeV2Test<Object, RoaringBitmap> {
+public class DistinctCountBitmapWithMVStarTreeV2Test extends
BaseStarTreeV2Test<Object, RoaringBitmapUnion> {
@Override
- ValueAggregator<Object, RoaringBitmap> getValueAggregator() {
+ ValueAggregator<Object, RoaringBitmapUnion> getValueAggregator() {
return new DistinctCountBitmapValueAggregator();
}
@@ -45,8 +46,14 @@ public class DistinctCountBitmapWithMVStarTreeV2Test extends
BaseStarTreeV2Test<
}
@Override
- void assertAggregatedValue(RoaringBitmap starTreeResult, RoaringBitmap
nonStarTreeResult) {
- assertEquals(starTreeResult, nonStarTreeResult);
+ MultipleTreesBuilder.BuildMode getBuildMode() {
+ // Pin the off-heap builder so its path is always covered: it serializes
every record when it is appended
+ return MultipleTreesBuilder.BuildMode.OFF_HEAP;
+ }
+
+ @Override
+ void assertAggregatedValue(RoaringBitmapUnion starTreeResult,
RoaringBitmapUnion nonStarTreeResult) {
+ assertEquals(starTreeResult.get(), nonStarTreeResult.get());
}
@Override
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountBitmapStarTreeV2Test.java
b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountHLLOnHeapStarTreeV2Test.java
similarity index 50%
copy from
pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountBitmapStarTreeV2Test.java
copy to
pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountHLLOnHeapStarTreeV2Test.java
index 044a9143837..2f456424e8b 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountBitmapStarTreeV2Test.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/DistinctCountHLLOnHeapStarTreeV2Test.java
@@ -18,34 +18,16 @@
*/
package org.apache.pinot.core.startree.v2;
-import java.util.Random;
-import
org.apache.pinot.segment.local.aggregator.DistinctCountBitmapValueAggregator;
-import org.apache.pinot.segment.local.aggregator.ValueAggregator;
-import org.apache.pinot.spi.data.FieldSpec.DataType;
-import org.roaringbitmap.RoaringBitmap;
+import org.apache.pinot.segment.local.startree.v2.builder.MultipleTreesBuilder;
-import static org.testng.Assert.assertEquals;
-
-public class DistinctCountBitmapStarTreeV2Test extends
BaseStarTreeV2Test<Object, RoaringBitmap> {
-
- @Override
- ValueAggregator<Object, RoaringBitmap> getValueAggregator() {
- return new DistinctCountBitmapValueAggregator();
- }
-
- @Override
- DataType getRawValueType() {
- return DataType.INT;
- }
-
- @Override
- Object getRandomRawValue(Random random) {
- return random.nextInt(100);
- }
+/// Runs [DistinctCountHLLStarTreeV2Test] with the on-heap builder pinned, so
that its pass that serializes
+/// variable-length metrics before the forward indexes are sized is always
covered for an aggregator other than the
+/// bitmap one. The parent class keeps drawing the build mode at random.
+public class DistinctCountHLLOnHeapStarTreeV2Test extends
DistinctCountHLLStarTreeV2Test {
@Override
- void assertAggregatedValue(RoaringBitmap starTreeResult, RoaringBitmap
nonStarTreeResult) {
- assertEquals(starTreeResult, nonStarTreeResult);
+ MultipleTreesBuilder.BuildMode getBuildMode() {
+ return MultipleTreesBuilder.BuildMode.ON_HEAP;
}
}
diff --git
a/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkDistinctCountBitmapAggregation.java
b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkDistinctCountBitmapAggregation.java
new file mode 100644
index 00000000000..78f9604166b
--- /dev/null
+++
b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkDistinctCountBitmapAggregation.java
@@ -0,0 +1,121 @@
+/**
+ * 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.pinot.perf.aggregation;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.concurrent.TimeUnit;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.utils.RoaringBitmapUtils;
+import org.apache.pinot.core.common.BlockValSet;
+import org.apache.pinot.core.common.SyntheticBlockValSets;
+import org.apache.pinot.core.plan.DocIdSetPlanNode;
+import org.apache.pinot.core.query.aggregation.AggregationResultHolder;
+import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
+import
org.apache.pinot.core.query.aggregation.function.DistinctCountBitmapAggregationFunction;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+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.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.RunnerException;
+import org.openjdk.jmh.runner.options.Options;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+import org.roaringbitmap.RoaringBitmap;
+
+
+/// Benchmarks the serialized-bitmap (`BYTES`) aggregation path of
DISTINCT_COUNT_BITMAP, i.e. the path used when
+/// unioning pre-aggregated bitmaps such as star-tree DISTINCT_COUNT_BITMAP
columns. `_valuesPerBitmap` controls the
+/// input bitmap size and `_maxValue` the value universe: a small universe
produces dense accumulator containers
+/// (bitmap containers), a large one sparse accumulators (array containers).
+@Fork(1)
+@BenchmarkMode(Mode.Throughput)
+@Warmup(iterations = 30, time = 100, timeUnit = TimeUnit.MILLISECONDS)
+@Measurement(iterations = 30, time = 100, timeUnit = TimeUnit.MILLISECONDS)
+@OutputTimeUnit(TimeUnit.SECONDS)
+@State(Scope.Benchmark)
+public class BenchmarkDistinctCountBitmapAggregation extends
AbstractAggregationFunctionBenchmark.Stable {
+ private static final ExpressionContext EXPR =
ExpressionContext.forIdentifier("col");
+
+ @Param({"200", "2000"})
+ private int _valuesPerBitmap;
+
+ // The last universe spreads values over 32768 container keys, like hash
codes of a non-numeric column do
+ @Param({"2000000", "100000000", "2147483647"})
+ private int _maxValue;
+
+ public static void main(String[] args) throws RunnerException {
+ Options opt =
+ new
OptionsBuilder().include(BenchmarkDistinctCountBitmapAggregation.class.getSimpleName()).build();
+ new Runner(opt).run();
+ }
+
+ @Override
+ protected AggregationFunction<?, ?> createAggregationFunction() {
+ return new DistinctCountBitmapAggregationFunction(List.of(EXPR), false);
+ }
+
+ @Override
+ protected AggregationResultHolder createResultHolder() {
+ return getAggregationFunction().createAggregationResultHolder();
+ }
+
+ @Override
+ protected Map<ExpressionContext, BlockValSet> createBlockValSetMap() {
+ Random random = new Random(420);
+ int numDocs = DocIdSetPlanNode.MAX_DOC_PER_CALL;
+ BlockValSet block = SyntheticBlockValSets.Bytes.create(numDocs, null, ()
-> {
+ RoaringBitmap bitmap = new RoaringBitmap();
+ for (int i = 0; i < _valuesPerBitmap; i++) {
+ bitmap.add(random.nextInt(_maxValue));
+ }
+ return RoaringBitmapUtils.serialize(bitmap);
+ });
+ return Map.of(EXPR, block);
+ }
+
+ @Override
+ protected Object createExpectedResult(Map<ExpressionContext, BlockValSet>
map) {
+ RoaringBitmap expected = new RoaringBitmap();
+ for (byte[] serialized : map.get(EXPR).getBytesValuesSV()) {
+ expected.or(RoaringBitmapUtils.deserialize(serialized));
+ }
+ return expected.getCardinality();
+ }
+
+ @Override
+ protected Comparable extractFinalResult(AggregationResultHolder
resultHolder) {
+ // Route through extractAggregationResult(): the holder's accumulator is
not directly consumable (it may be in
+ // the lazy union state that extraction repairs), matching how the query
engine reads aggregation results
+ DistinctCountBitmapAggregationFunction function =
+ (DistinctCountBitmapAggregationFunction) getAggregationFunction();
+ return
function.extractFinalResult(function.extractAggregationResult(resultHolder));
+ }
+
+ @Override
+ protected void resetResultHolder(AggregationResultHolder resultHolder) {
+ resultHolder.setValue(null);
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountBitmapValueAggregator.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountBitmapValueAggregator.java
index 74fbb95f2d8..827dec88fad 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountBitmapValueAggregator.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/DistinctCountBitmapValueAggregator.java
@@ -18,13 +18,21 @@
*/
package org.apache.pinot.segment.local.aggregator;
+import org.apache.pinot.common.utils.RoaringBitmapUnion;
import org.apache.pinot.common.utils.RoaringBitmapUtils;
import org.apache.pinot.segment.spi.AggregationFunctionType;
import org.apache.pinot.spi.data.FieldSpec.DataType;
-import org.roaringbitmap.RoaringBitmap;
-public class DistinctCountBitmapValueAggregator implements
ValueAggregator<Object, RoaringBitmap> {
+/// Pre-aggregates distinct values into a [RoaringBitmapUnion], which unions
serialized-bitmap (`byte[]`) inputs
+/// lazily and finalizes the bitmap only when it is read or serialized. A
given metric column always provides either
+/// serialized bitmaps or plain raw values, never both.
+///
+/// [#getMaxAggregatedValueByteSize] reports the largest value produced by
[#serializeAggregatedValue] so far. This
+/// aggregator is only used by the star-tree builders (it has no fixed size,
so ingestion-time aggregation rejects
+/// it), and both builders serialize every record before that size is
consumed: the off-heap builder when it appends
+/// a record, the on-heap builder in its pre-serialization pass before the
forward indexes are sized.
+public class DistinctCountBitmapValueAggregator implements
ValueAggregator<Object, RoaringBitmapUnion> {
public static final DataType AGGREGATED_VALUE_TYPE = DataType.BYTES;
private int _maxByteSize;
@@ -40,57 +48,53 @@ public class DistinctCountBitmapValueAggregator implements
ValueAggregator<Objec
}
@Override
- public RoaringBitmap getInitialAggregatedValue(Object rawValue) {
+ public RoaringBitmapUnion getInitialAggregatedValue(Object rawValue) {
// NOTE: rawValue cannot be null because this aggregator can only be used
for star-tree index, and the builder
// never passes a null raw value: a null-aware star-tree leaves the
aggregated value null until the group sees
// its first non-null input.
assert rawValue != null;
- RoaringBitmap initialValue;
if (rawValue instanceof byte[]) {
- byte[] bytes = (byte[]) rawValue;
- initialValue = deserializeAggregatedValue(bytes);
- _maxByteSize = Math.max(_maxByteSize, bytes.length);
- } else {
- initialValue = new RoaringBitmap();
- addToValue(initialValue, rawValue);
- _maxByteSize = Math.max(_maxByteSize,
initialValue.serializedSizeInBytes());
+ return deserializeAggregatedValue((byte[]) rawValue);
}
+ RoaringBitmapUnion initialValue = new RoaringBitmapUnion();
+ addToValue(initialValue, rawValue);
return initialValue;
}
@Override
- public RoaringBitmap applyRawValue(RoaringBitmap value, Object rawValue) {
+ public RoaringBitmapUnion applyRawValue(RoaringBitmapUnion value, Object
rawValue) {
if (rawValue instanceof byte[]) {
- value.or(deserializeAggregatedValue((byte[]) rawValue));
+ value.add(RoaringBitmapUtils.deserialize((byte[]) rawValue));
} else {
addToValue(value, rawValue);
}
- _maxByteSize = Math.max(_maxByteSize, value.serializedSizeInBytes());
return value;
}
- /// Adds a raw value (single value or multi-value array) to the
RoaringBitmap.
- protected void addToValue(RoaringBitmap bitmap, Object rawValue) {
+ /// Adds a raw value (single value or multi-value array) to the union.
+ protected void addToValue(RoaringBitmapUnion union, Object rawValue) {
if (rawValue instanceof Object[]) {
Object[] values = (Object[]) rawValue;
for (Object value : values) {
- bitmap.add(value.hashCode());
+ union.add(value.hashCode());
}
} else {
- bitmap.add(rawValue.hashCode());
+ union.add(rawValue.hashCode());
}
}
@Override
- public RoaringBitmap applyAggregatedValue(RoaringBitmap value, RoaringBitmap
aggregatedValue) {
- value.or(aggregatedValue);
- _maxByteSize = Math.max(_maxByteSize, value.serializedSizeInBytes());
+ public RoaringBitmapUnion applyAggregatedValue(RoaringBitmapUnion value,
RoaringBitmapUnion aggregatedValue) {
+ // get() finalizes the other accumulator without consuming it; it stays
usable (and is serialized later itself)
+ value.add(aggregatedValue.get());
return value;
}
@Override
- public RoaringBitmap cloneAggregatedValue(RoaringBitmap value) {
- return value.clone();
+ public RoaringBitmapUnion cloneAggregatedValue(RoaringBitmapUnion value) {
+ RoaringBitmapUnion clone = new RoaringBitmapUnion();
+ clone.add(value.get());
+ return clone;
}
@Override
@@ -104,12 +108,14 @@ public class DistinctCountBitmapValueAggregator
implements ValueAggregator<Objec
}
@Override
- public byte[] serializeAggregatedValue(RoaringBitmap value) {
- return RoaringBitmapUtils.serialize(value);
+ public byte[] serializeAggregatedValue(RoaringBitmapUnion value) {
+ byte[] bytes = RoaringBitmapUtils.serialize(value.get());
+ _maxByteSize = Math.max(_maxByteSize, bytes.length);
+ return bytes;
}
@Override
- public RoaringBitmap deserializeAggregatedValue(byte[] bytes) {
- return RoaringBitmapUtils.deserialize(bytes);
+ public RoaringBitmapUnion deserializeAggregatedValue(byte[] bytes) {
+ return RoaringBitmapUtils.deserializeToUnion(bytes);
}
}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregator.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregator.java
index 350535d76d9..6ee71277681 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregator.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregator.java
@@ -81,6 +81,11 @@ public interface ValueAggregator<R, A> {
boolean isAggregatedValueFixedSize();
/// Returns the maximum size in bytes of the aggregated values seen so far.
+ ///
+ /// For variable-size (`BYTES`) aggregated values this is only complete once
every value has been serialized
+ /// through [#serializeAggregatedValue]: an aggregator may derive the size
from the values it serialized rather than
+ /// measure every intermediate state. Builders must therefore serialize all
values before consuming the size, as the
+ /// star-tree builders do.
int getMaxAggregatedValueByteSize();
/// Serializes an aggregated value into a byte array.
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/json/MutableJsonIndexImpl.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/json/MutableJsonIndexImpl.java
index 4b9a3c99c55..8010ab5d20d 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/json/MutableJsonIndexImpl.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/json/MutableJsonIndexImpl.java
@@ -48,6 +48,7 @@ import
org.apache.pinot.common.request.context.predicate.NotInPredicate;
import org.apache.pinot.common.request.context.predicate.Predicate;
import org.apache.pinot.common.request.context.predicate.RangePredicate;
import org.apache.pinot.common.request.context.predicate.RegexpLikePredicate;
+import org.apache.pinot.common.utils.RoaringBitmapUnion;
import org.apache.pinot.common.utils.SegmentUtils;
import org.apache.pinot.common.utils.regex.Matcher;
import org.apache.pinot.common.utils.regex.Pattern;
@@ -200,6 +201,40 @@ public class MutableJsonIndexImpl implements
MutableJsonIndex {
return predicateType == Predicate.Type.IS_NULL;
}
+ /// Folds many posting lists into one [LazyBitmap] through a
[RoaringBitmapUnion], which unions lazily and
+ /// finalizes the bitmap once at [#get]. A single input is borrowed: it is
wrapped as an immutable [LazyBitmap],
+ /// never adopted or modified.
+ private static class UnionAccumulator {
+ @Nullable
+ private RoaringBitmap _first;
+ @Nullable
+ private RoaringBitmapUnion _union;
+
+ void add(RoaringBitmap docIds) {
+ if (docIds.isEmpty()) {
+ return;
+ }
+ if (_union != null) {
+ _union.add(docIds);
+ } else if (_first == null) {
+ _first = docIds;
+ } else {
+ // The borrowed first posting list is copied into the union, never
adopted
+ _union = new RoaringBitmapUnion();
+ _union.add(_first);
+ _union.add(docIds);
+ _first = null;
+ }
+ }
+
+ LazyBitmap get() {
+ if (_union != null) {
+ return LazyBitmap.createMutable(_union.take());
+ }
+ return _first != null ? LazyBitmap.createImmutable(_first) :
LazyBitmap.EMPTY_BITMAP;
+ }
+ }
+
/// This class allows delaying of cloning posting list bitmap for as long as
possible
/// It stores either a bitmap from posting list that must be cloned before
mutating (readOnly=true)
/// or an already cloned bitmap.
@@ -280,20 +315,6 @@ public class MutableJsonIndexImpl implements
MutableJsonIndex {
return createMutable(RoaringBitmap.or(_value, other._value));
}
- LazyBitmap or(RoaringBitmap bitmap) {
- if (isEmpty()) {
- return createImmutable(bitmap);
- }
- if (bitmap.isEmpty()) {
- return this;
- }
- if (isMutable()) {
- _value.or(bitmap);
- return this;
- }
- return createMutable(RoaringBitmap.or(_value, bitmap));
- }
-
LazyBitmap andNot(LazyBitmap other) {
if (isEmpty()) {
return EMPTY_BITMAP;
@@ -438,17 +459,17 @@ public class MutableJsonIndexImpl implements
MutableJsonIndex {
StringBuilder buffer = new StringBuilder(key);
buffer.append(JsonIndexCreator.KEY_VALUE_SEPARATOR);
int pos = buffer.length();
- LazyBitmap result = LazyBitmap.EMPTY_BITMAP;
+ UnionAccumulator result = new UnionAccumulator();
List<String> values = ((InPredicate) predicate).getValues();
for (String value : values) {
buffer.setLength(pos);
buffer.append(value);
RoaringBitmap docIds = _postingListMap.get(buffer.toString());
if (docIds != null) {
- result = result.or(docIds);
+ result.add(docIds);
}
}
- return result;
+ return result.get();
}
case NOT_IN: {
@@ -488,7 +509,7 @@ public class MutableJsonIndexImpl implements
MutableJsonIndex {
}
Pattern pattern = ((RegexpLikePredicate) predicate).getPattern();
Matcher matcher = pattern.matcher("");
- LazyBitmap result = LazyBitmap.EMPTY_BITMAP;
+ UnionAccumulator result = new UnionAccumulator();
StringBuilder value = new StringBuilder();
int valueStart = key.length() + 1;
for (Map.Entry<String, RoaringBitmap> entry : subMap.entrySet()) {
@@ -496,10 +517,10 @@ public class MutableJsonIndexImpl implements
MutableJsonIndex {
value.setLength(0);
value.append(keyValue, valueStart, keyValue.length());
if (matcher.reset(value).matches()) {
- result = result.or(entry.getValue());
+ result.add(entry.getValue());
}
}
- return result;
+ return result.get();
}
case RANGE: {
@@ -521,7 +542,7 @@ public class MutableJsonIndexImpl implements
MutableJsonIndex {
boolean upperInclusive = upperUnbounded ||
rangePredicate.isUpperInclusive();
Object lowerBound = lowerUnbounded ? null :
rangeDataType.convert(rangePredicate.getLowerBound());
Object upperBound = upperUnbounded ? null :
rangeDataType.convert(rangePredicate.getUpperBound());
- LazyBitmap result = LazyBitmap.EMPTY_BITMAP;
+ UnionAccumulator result = new UnionAccumulator();
int valueStart = key.length() + 1;
for (Map.Entry<String, RoaringBitmap> entry : subMap.entrySet()) {
Object valueObj =
rangeDataType.convert(entry.getKey().substring(valueStart));
@@ -532,10 +553,10 @@ public class MutableJsonIndexImpl implements
MutableJsonIndex {
upperUnbounded || (upperInclusive ?
rangeDataType.compare(valueObj, upperBound) <= 0
: rangeDataType.compare(valueObj, upperBound) < 0);
if (lowerCompareResult && upperCompareResult) {
- result = result.or(entry.getValue());
+ result.add(entry.getValue());
}
}
- return result;
+ return result.get();
}
default:
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/json/ImmutableJsonIndexReader.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/json/ImmutableJsonIndexReader.java
index 242185c2be3..8e759ea57f6 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/json/ImmutableJsonIndexReader.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/json/ImmutableJsonIndexReader.java
@@ -46,6 +46,7 @@ import
org.apache.pinot.common.request.context.predicate.NotInPredicate;
import org.apache.pinot.common.request.context.predicate.Predicate;
import org.apache.pinot.common.request.context.predicate.RangePredicate;
import org.apache.pinot.common.request.context.predicate.RegexpLikePredicate;
+import org.apache.pinot.common.utils.MutableRoaringBitmapUnion;
import org.apache.pinot.common.utils.regex.Matcher;
import org.apache.pinot.common.utils.regex.Pattern;
import
org.apache.pinot.segment.local.segment.creator.impl.inv.json.BaseJsonIndexCreator;
@@ -176,6 +177,40 @@ public class ImmutableJsonIndexReader implements
JsonIndexReader {
return ImmutableRoaringBitmap.and(target, other);
}
+ /// Folds many posting lists into one result through a
[MutableRoaringBitmapUnion], which unions lazily and
+ /// finalizes the bitmap once at [#get]. The zero- and single-input cases
return the borrowed posting list without
+ /// copying, matching the ownership behavior of [#or].
+ private static class UnionAccumulator {
+ @Nullable
+ private ImmutableRoaringBitmap _first;
+ @Nullable
+ private MutableRoaringBitmapUnion _union;
+
+ void add(ImmutableRoaringBitmap docIds) {
+ if (docIds.isEmpty()) {
+ return;
+ }
+ if (_union != null) {
+ _union.add(docIds);
+ } else if (_first == null) {
+ _first = docIds;
+ } else {
+ // The borrowed first posting list is copied into the union, never
adopted
+ _union = new MutableRoaringBitmapUnion();
+ _union.add(_first);
+ _union.add(docIds);
+ _first = null;
+ }
+ }
+
+ ImmutableRoaringBitmap get() {
+ if (_union != null) {
+ return _union.take();
+ }
+ return _first != null ? _first : EMPTY_BITMAP;
+ }
+ }
+
private static ImmutableRoaringBitmap or(ImmutableRoaringBitmap target,
ImmutableRoaringBitmap other) {
if (target.isEmpty()) {
return other;
@@ -320,17 +355,17 @@ public class ImmutableJsonIndexReader implements
JsonIndexReader {
StringBuilder buffer = new StringBuilder(key);
buffer.append(JsonIndexCreator.KEY_VALUE_SEPARATOR);
int pos = buffer.length();
- ImmutableRoaringBitmap result = EMPTY_BITMAP;
+ UnionAccumulator result = new UnionAccumulator();
List<String> values = ((InPredicate) predicate).getValues();
for (String value : values) {
buffer.setLength(pos);
buffer.append(value);
int dictId = _dictionary.indexOf(buffer.toString());
if (dictId >= 0) {
- result = or(result, _invertedIndex.getDocIds(dictId));
+ result.add(_invertedIndex.getDocIds(dictId));
}
}
- return result;
+ return result.get();
}
case NOT_IN: {
@@ -376,13 +411,13 @@ public class ImmutableJsonIndexReader implements
JsonIndexReader {
notInDictIds.add(dictId);
}
}
- ImmutableRoaringBitmap result = EMPTY_BITMAP;
+ UnionAccumulator result = new UnionAccumulator();
for (int dictId = dictIdRange[0]; dictId < dictIdRange[1]; dictId++)
{
if (!notInDictIds.contains(dictId)) {
- result = or(result, _invertedIndex.getDocIds(dictId));
+ result.add(_invertedIndex.getDocIds(dictId));
}
}
- return result;
+ return result.get();
}
}
@@ -404,7 +439,7 @@ public class ImmutableJsonIndexReader implements
JsonIndexReader {
}
Pattern pattern = ((RegexpLikePredicate) predicate).getPattern();
Matcher matcher = pattern.matcher("");
- ImmutableRoaringBitmap result = EMPTY_BITMAP;
+ UnionAccumulator result = new UnionAccumulator();
byte[] dictBuffer = _dictionary.getBuffer();
StringBuilder value = new StringBuilder();
int valueStart = key.length() + 1;
@@ -413,10 +448,10 @@ public class ImmutableJsonIndexReader implements
JsonIndexReader {
value.setLength(0);
value.append(keyValue, valueStart, keyValue.length());
if (matcher.reset(value).matches()) {
- result = or(result, _invertedIndex.getDocIds(dictId));
+ result.add(_invertedIndex.getDocIds(dictId));
}
}
- return result;
+ return result.get();
}
case RANGE: {
@@ -439,7 +474,7 @@ public class ImmutableJsonIndexReader implements
JsonIndexReader {
boolean upperInclusive = upperUnbounded ||
rangePredicate.isUpperInclusive();
Object lowerBound = lowerUnbounded ? null :
rangeDataType.convert(rangePredicate.getLowerBound());
Object upperBound = upperUnbounded ? null :
rangeDataType.convert(rangePredicate.getUpperBound());
- ImmutableRoaringBitmap result = EMPTY_BITMAP;
+ UnionAccumulator result = new UnionAccumulator();
byte[] dictBuffer = _dictionary.getBuffer();
int valueStart = key.length() + 1;
for (int dictId = dictIds[0]; dictId < dictIds[1]; dictId++) {
@@ -452,10 +487,10 @@ public class ImmutableJsonIndexReader implements
JsonIndexReader {
upperUnbounded || (upperInclusive ?
rangeDataType.compare(valueObj, upperBound) <= 0
: rangeDataType.compare(valueObj, upperBound) < 0);
if (lowerCompareResult && upperCompareResult) {
- result = or(result, _invertedIndex.getDocIds(dictId));
+ result.add(_invertedIndex.getDocIds(dictId));
}
}
- return result;
+ return result.get();
}
default:
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/BaseSingleTreeBuilder.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/BaseSingleTreeBuilder.java
index 5ef46fc14ed..c0aa0da4954 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/BaseSingleTreeBuilder.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/BaseSingleTreeBuilder.java
@@ -212,6 +212,17 @@ abstract class BaseSingleTreeBuilder implements
SingleTreeBuilder {
abstract Record getStarTreeRecord(int docId)
throws IOException;
+ /// Serializes metrics whose [ValueAggregator#getMaxAggregatedValueByteSize]
is derived during serialization before
+ /// the forward indexes are sized. Returns, per metric, whether
[#getStarTreeRecord] now holds the serialized
+ /// `byte[]` in place of the aggregated value. A builder that already
serialized every record when it was appended
+ /// (the off-heap builder) has nothing to do and returns all `false`.
+ ///
+ /// @return per-metric flag telling whether the star-tree records hold
pre-serialized values
+ boolean[] preSerializeMetrics()
+ throws IOException {
+ return new boolean[_numMetrics];
+ }
+
/// Returns the dimension value of the given document and dimension Id in
the star-tree.
///
/// @param docId Document Id
@@ -541,6 +552,9 @@ abstract class BaseSingleTreeBuilder implements
SingleTreeBuilder {
private void createForwardIndexes()
throws Exception {
+ // Finalize sizes for aggregators that track their maximum during
serialization before opening any creator.
+ boolean[] preSerializedMetrics = preSerializeMetrics();
+
SingleValueUnsortedForwardIndexCreator[] dimensionIndexCreators =
new SingleValueUnsortedForwardIndexCreator[_numDimensions];
for (int i = 0; i < _numDimensions; i++) {
@@ -613,6 +627,9 @@ abstract class BaseSingleTreeBuilder implements
SingleTreeBuilder {
if (value == null) {
metricNullValueVectorCreators[i].setNull(docId);
}
+ } else if (preSerializedMetrics[i]) {
+ metricIndexCreators[i].putBytes((byte[]) value);
+ continue;
}
putMetricValue(metricIndexCreators[i], valueAggregator, value);
}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/OnHeapSingleTreeBuilder.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/OnHeapSingleTreeBuilder.java
index 60c60aa67ce..88ccd3a2b3f 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/OnHeapSingleTreeBuilder.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/builder/OnHeapSingleTreeBuilder.java
@@ -24,6 +24,8 @@ import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import org.apache.commons.configuration2.Configuration;
+import
org.apache.pinot.segment.local.aggregator.DistinctCountBitmapValueAggregator;
+import org.apache.pinot.segment.local.aggregator.ValueAggregator;
import org.apache.pinot.segment.spi.ImmutableSegment;
import org.apache.pinot.segment.spi.index.startree.StarTreeV2Constants;
@@ -53,6 +55,28 @@ public class OnHeapSingleTreeBuilder extends
BaseSingleTreeBuilder {
return _records.get(docId);
}
+ @Override
+ @SuppressWarnings("unchecked")
+ boolean[] preSerializeMetrics() {
+ // Bitmap distinct count derives its maximum size during serialization.
Other aggregators already know their
+ // maximum size and can serialize during the index write without retaining
a byte array for every record.
+ boolean[] preSerializedMetrics = new boolean[_numMetrics];
+ for (int i = 0; i < _numMetrics; i++) {
+ ValueAggregator valueAggregator = _valueAggregators[i];
+ if (!(valueAggregator instanceof DistinctCountBitmapValueAggregator)) {
+ continue;
+ }
+ for (Record record : _records) {
+ Object value = record._metrics[i];
+ if (value != null) {
+ record._metrics[i] = valueAggregator.serializeAggregatedValue(value);
+ }
+ }
+ preSerializedMetrics[i] = true;
+ }
+ return preSerializedMetrics;
+ }
+
@Override
int getDimensionValue(int docId, int dimensionId) {
return _records.get(docId)._dimensions[dimensionId];
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/aggregator/DistinctCountBitmapValueAggregatorTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/aggregator/DistinctCountBitmapValueAggregatorTest.java
new file mode 100644
index 00000000000..92b5f11735c
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/aggregator/DistinctCountBitmapValueAggregatorTest.java
@@ -0,0 +1,161 @@
+/**
+ * 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.pinot.segment.local.aggregator;
+
+import java.util.Random;
+import org.apache.pinot.common.utils.RoaringBitmapUnion;
+import org.apache.pinot.common.utils.RoaringBitmapUtils;
+import org.roaringbitmap.RoaringBitmap;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotSame;
+
+
+/// Covers the serialized-bitmap fold of [DistinctCountBitmapValueAggregator],
which unions inputs lazily through a
+/// [RoaringBitmapUnion] and finalizes the bitmap when it is read or
serialized. Input cardinalities cross the lazy
+/// container-promotion threshold so both promoted and non-promoted
accumulator states are verified against an eagerly
+/// unioned reference.
+public class DistinctCountBitmapValueAggregatorTest {
+ private static final Random RANDOM = new Random(42);
+
+ private static byte[][] serializedBitmaps(int numBitmaps, int
valuesPerBitmap, int maxValue) {
+ byte[][] serialized = new byte[numBitmaps][];
+ for (int i = 0; i < numBitmaps; i++) {
+ RoaringBitmap bitmap = new RoaringBitmap();
+ for (int j = 0; j < valuesPerBitmap; j++) {
+ bitmap.add(RANDOM.nextInt(maxValue));
+ }
+ serialized[i] = RoaringBitmapUtils.serialize(bitmap);
+ }
+ return serialized;
+ }
+
+ private static RoaringBitmap eagerUnion(byte[][] serialized) {
+ RoaringBitmap expected = new RoaringBitmap();
+ for (byte[] bytes : serialized) {
+ expected.or(RoaringBitmapUtils.deserialize(bytes));
+ }
+ return expected;
+ }
+
+ private static RoaringBitmapUnion fold(DistinctCountBitmapValueAggregator
aggregator, byte[][] serialized) {
+ RoaringBitmapUnion value =
aggregator.getInitialAggregatedValue(serialized[0]);
+ for (int i = 1; i < serialized.length; i++) {
+ value = aggregator.applyRawValue(value, serialized[i]);
+ }
+ return value;
+ }
+
+ @Test
+ public void testApplyRawSerializedBitmaps() {
+ byte[][] serialized = serializedBitmaps(200, 100, 200_000);
+ DistinctCountBitmapValueAggregator aggregator = new
DistinctCountBitmapValueAggregator();
+ assertFalse(aggregator.isAggregatedValueFixedSize());
+ // Nothing has been serialized yet, so no size is known: the star-tree
builders serialize every record before
+ // they consume the size
+ assertEquals(aggregator.getMaxAggregatedValueByteSize(), 0);
+
+ RoaringBitmapUnion value = fold(aggregator, serialized);
+ byte[] result = aggregator.serializeAggregatedValue(value);
+
+ RoaringBitmap expected = eagerUnion(serialized);
+ assertEquals(RoaringBitmapUtils.deserialize(result), expected);
+ assertEquals(value.get().getCardinality(), expected.getCardinality());
+ assertEquals(aggregator.getMaxAggregatedValueByteSize(), result.length);
+
+ // Serializing is non-consuming and repeatable
+ assertEquals(aggregator.serializeAggregatedValue(value), result);
+ assertEquals(aggregator.getMaxAggregatedValueByteSize(), result.length);
+ }
+
+ @Test
+ public void testMaxByteSizeTracksLargestSerializedValue() {
+ DistinctCountBitmapValueAggregator aggregator = new
DistinctCountBitmapValueAggregator();
+ byte[] small = aggregator.serializeAggregatedValue(fold(aggregator,
serializedBitmaps(2, 10, 1_000)));
+ byte[] large = aggregator.serializeAggregatedValue(fold(aggregator,
serializedBitmaps(50, 100, 200_000)));
+ byte[] medium = aggregator.serializeAggregatedValue(fold(aggregator,
serializedBitmaps(5, 100, 200_000)));
+ assertEquals(aggregator.getMaxAggregatedValueByteSize(),
Math.max(small.length, Math.max(large.length,
+ medium.length)));
+ }
+
+ @Test
+ public void testApplyAggregatedValueWithPendingInput() {
+ // Mirrors the on-heap star-tree builder: aggregated records are built
from other records' accumulators via
+ // cloneAggregatedValue + applyAggregatedValue, and every record is only
serialized at the end
+ byte[][] serializedA = serializedBitmaps(100, 100, 150_000);
+ byte[][] serializedB = serializedBitmaps(100, 100, 150_000);
+ DistinctCountBitmapValueAggregator aggregator = new
DistinctCountBitmapValueAggregator();
+
+ RoaringBitmapUnion recordA = fold(aggregator, serializedA);
+ RoaringBitmapUnion recordB = fold(aggregator, serializedB);
+
+ // recordA and recordB both have pending lazy unions here
+ RoaringBitmapUnion aggregated = aggregator.cloneAggregatedValue(recordA);
+ assertNotSame(aggregated, recordA);
+ aggregated = aggregator.applyAggregatedValue(aggregated, recordB);
+
+ RoaringBitmap expected = eagerUnion(serializedA);
+ expected.or(eagerUnion(serializedB));
+
assertEquals(RoaringBitmapUtils.deserialize(aggregator.serializeAggregatedValue(aggregated)),
expected);
+
+ // Serializing is non-consuming: the same accumulator remains usable for
further aggregation, and the value
+ // published by the earlier serialization is not affected
+ RoaringBitmap published = aggregated.get();
+ RoaringBitmap publishedCopy = published.clone();
+ byte[] extra = serializedBitmaps(1, 100, 150_000)[0];
+ aggregated = aggregator.applyRawValue(aggregated, extra);
+ assertEquals(published, publishedCopy);
+ expected.or(RoaringBitmapUtils.deserialize(extra));
+
assertEquals(RoaringBitmapUtils.deserialize(aggregator.serializeAggregatedValue(aggregated)),
expected);
+
+ // The source records must remain intact and serializable
+
assertEquals(RoaringBitmapUtils.deserialize(aggregator.serializeAggregatedValue(recordA)),
+ eagerUnion(serializedA));
+
assertEquals(RoaringBitmapUtils.deserialize(aggregator.serializeAggregatedValue(recordB)),
+ eagerUnion(serializedB));
+ // The clone does not share state with its source
+ recordA = aggregator.applyRawValue(recordA, serializedB[0]);
+
assertEquals(RoaringBitmapUtils.deserialize(aggregator.serializeAggregatedValue(aggregated)),
expected);
+ }
+
+ @Test
+ public void testApplyRawValuesOnly() {
+ DistinctCountBitmapValueAggregator aggregator = new
DistinctCountBitmapValueAggregator();
+ RoaringBitmapUnion value = aggregator.getInitialAggregatedValue(1);
+ for (int i = 2; i <= 1000; i++) {
+ value = aggregator.applyRawValue(value, i);
+ }
+ value = aggregator.applyRawValue(value, new Object[]{1001, 1002, 1002});
+ byte[] result = aggregator.serializeAggregatedValue(value);
+ assertEquals(RoaringBitmapUtils.deserialize(result).getCardinality(),
1002);
+ assertEquals(aggregator.getMaxAggregatedValueByteSize(), result.length);
+ }
+
+ @Test
+ public void testDeserializeRoundTrip() {
+ DistinctCountBitmapValueAggregator aggregator = new
DistinctCountBitmapValueAggregator();
+ byte[][] serialized = serializedBitmaps(20, 100, 50_000);
+ byte[] bytes = aggregator.serializeAggregatedValue(fold(aggregator,
serialized));
+ RoaringBitmapUnion restored = aggregator.deserializeAggregatedValue(bytes);
+ assertEquals(restored.get(), eagerUnion(serialized));
+ assertEquals(aggregator.serializeAggregatedValue(restored), bytes);
+ }
+}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/NullAwareStarTreeBuilderTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/NullAwareStarTreeBuilderTest.java
index 11dd454ad0f..ac018cff381 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/NullAwareStarTreeBuilderTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/startree/v2/builder/NullAwareStarTreeBuilderTest.java
@@ -20,14 +20,20 @@ package org.apache.pinot.segment.local.startree.v2.builder;
import java.io.File;
import java.io.IOException;
+import java.math.BigDecimal;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.List;
+import java.util.Map;
+import org.apache.commons.configuration2.PropertiesConfiguration;
import org.apache.commons.io.FileUtils;
+import org.apache.pinot.common.utils.RoaringBitmapUtils;
import
org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
import
org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
import
org.apache.pinot.segment.local.startree.v2.builder.MultipleTreesBuilder.BuildMode;
import org.apache.pinot.segment.spi.AggregationFunctionType;
+import org.apache.pinot.segment.spi.Constants;
import org.apache.pinot.segment.spi.ImmutableSegment;
import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader;
@@ -35,6 +41,7 @@ import
org.apache.pinot.segment.spi.index.reader.ForwardIndexReaderContext;
import org.apache.pinot.segment.spi.index.reader.NullValueVectorReader;
import
org.apache.pinot.segment.spi.index.startree.AggregationFunctionColumnPair;
import org.apache.pinot.segment.spi.index.startree.StarTreeV2;
+import org.apache.pinot.spi.config.table.StarTreeAggregationConfig;
import org.apache.pinot.spi.config.table.StarTreeIndexConfig;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
@@ -43,6 +50,7 @@ import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.data.readers.GenericRow;
import org.apache.pinot.spi.utils.ReadMode;
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.roaringbitmap.RoaringBitmap;
import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.DataProvider;
@@ -77,6 +85,8 @@ public class NullAwareStarTreeBuilderTest {
private static final String MIN_COLUMN =
new AggregationFunctionColumnPair(AggregationFunctionType.MIN,
METRIC).toColumnName();
+ private static final String BITMAP_COLUMN =
+ new
AggregationFunctionColumnPair(AggregationFunctionType.DISTINCTCOUNTBITMAP,
METRIC).toColumnName();
@AfterMethod
public void cleanUp()
@@ -118,6 +128,88 @@ public class NullAwareStarTreeBuilderTest {
}
}
+ /// A bitmap metric is serialized before its forward index is sized, which
the on-heap builder does in a
+ /// pass of its own. An all-null group has nothing to serialize: it must end
up in the null vector while the other
+ /// groups keep their values.
+ @Test(dataProvider = "buildModes")
+ public void
anAllNullGroupSurvivesTheBuildWithAVariableLengthMetric(BuildMode buildMode)
+ throws Exception {
+ File indexDir = createSegment();
+ buildStarTrees(indexDir, buildMode, starTreeConfig(true,
AggregationFunctionType.DISTINCTCOUNTBITMAP));
+
+ ImmutableSegment segment = ImmutableSegmentLoader.load(indexDir,
ReadMode.mmap);
+ try {
+ StarTreeV2 starTree = segment.getStarTrees().get(0);
+ ImmutableRoaringBitmap nullBitmap = nullBitmap(starTree, BITMAP_COLUMN);
+ assertNotNull(nullBitmap, "A null-aware star-tree must write a null
vector for a metric with an all-null group");
+ assertFalse(nullBitmap.isEmpty());
+
+ // Every group that did aggregate something holds some of the two
non-null values and nothing else
+ RoaringBitmap nonNullValues = RoaringBitmap.bitmapOf(10, 20);
+ ForwardIndexReader reader =
starTree.getDataSource(BITMAP_COLUMN).getForwardIndex();
+ assertNotNull(reader);
+ int numAggregatedGroups = 0;
+ try (ForwardIndexReaderContext context = reader.createContext()) {
+ for (int docId = 0; docId < starTree.getMetadata().getNumDocs();
docId++) {
+ if (!nullBitmap.contains(docId)) {
+ RoaringBitmap bitmap =
RoaringBitmapUtils.deserialize(reader.getBytes(docId, context));
+ assertFalse(bitmap.isEmpty());
+ assertTrue(nonNullValues.contains(bitmap));
+ numAggregatedGroups++;
+ }
+ }
+ }
+ assertTrue(numAggregatedGroups > 0, "Some group must have aggregated a
value, or the loop proves nothing");
+ } finally {
+ segment.destroy();
+ }
+ }
+
+ @Test
+ public void onlyBitmapMetricsAreRetainedAsSerializedBytes()
+ throws Exception {
+ File indexDir = createSegment();
+ StarTreeIndexConfig config = new StarTreeIndexConfig(List.of(DIMENSION),
null, List.of(BITMAP_COLUMN), List.of(
+ new StarTreeAggregationConfig(DIMENSION, "SUMPRECISION",
+ Map.of(Constants.SUMPRECISION_PRECISION_KEY, 1000), null, null,
null, null, null),
+ new StarTreeAggregationConfig(METRIC, "SUMPRECISION")),
MAX_LEAF_RECORDS, true);
+ String fixedSumColumn =
+ new
AggregationFunctionColumnPair(AggregationFunctionType.SUMPRECISION,
DIMENSION).toColumnName();
+ String variableSumColumn =
+ new
AggregationFunctionColumnPair(AggregationFunctionType.SUMPRECISION,
METRIC).toColumnName();
+ File outputDir = new File(TEMP_DIR, "mixedMetrics");
+ FileUtils.forceMkdir(outputDir);
+ ImmutableSegment segment = ImmutableSegmentLoader.load(indexDir,
ReadMode.mmap);
+ try (OnHeapSingleTreeBuilder builder = new OnHeapSingleTreeBuilder(
+ StarTreeV2BuilderConfig.fromIndexConfig(config), outputDir, segment,
new PropertiesConfiguration())) {
+ // A complete build also exercises writing pre-serialized bitmaps
alongside ordinary metric values.
+ builder.build();
+ List<String> metrics = Arrays.asList(builder._metrics);
+ int bitmapId = metrics.indexOf(BITMAP_COLUMN);
+ int fixedSumId = metrics.indexOf(fixedSumColumn);
+ int variableSumId = metrics.indexOf(variableSumColumn);
+ int maxBitmapSize = 0;
+ for (int docId = 0; docId < builder._numDocs; docId++) {
+ Object[] values = builder.getStarTreeRecord(docId)._metrics;
+ if (values[bitmapId] != null) {
+ byte[] bytes = (byte[]) values[bitmapId];
+ maxBitmapSize = Math.max(maxBitmapSize, bytes.length);
+ assertEquals(RoaringBitmapUtils.deserialize(bytes),
RoaringBitmap.bitmapOf(10, 20));
+ }
+ assertTrue(values[fixedSumId] instanceof BigDecimal);
+ if (values[variableSumId] != null) {
+ assertEquals(values[variableSumId], BigDecimal.valueOf(30));
+ }
+ }
+ assertTrue(maxBitmapSize > 0);
+
assertEquals(builder._valueAggregators[bitmapId].getMaxAggregatedValueByteSize(),
maxBitmapSize);
+ // A shared zero must not be expanded to a precision-1000 byte array on
every retained record.
+ assertSame(builder.getStarTreeRecord(0)._metrics[fixedSumId],
BigDecimal.ZERO);
+ } finally {
+ segment.destroy();
+ }
+ }
+
/// The two variants answer differently, so a segment has to be able to hold
both: the builder config's identity
/// includes the flag, which is what stops one being reused for the other.
@Test(dataProvider = "buildModes")
@@ -185,8 +277,13 @@ public class NullAwareStarTreeBuilderTest {
}
private static StarTreeIndexConfig starTreeConfig(boolean
nullHandlingEnabled) {
+ return starTreeConfig(nullHandlingEnabled, AggregationFunctionType.MIN);
+ }
+
+ private static StarTreeIndexConfig starTreeConfig(boolean
nullHandlingEnabled,
+ AggregationFunctionType aggregationType) {
return new StarTreeIndexConfig(List.of(DIMENSION), null,
- List.of(new AggregationFunctionColumnPair(AggregationFunctionType.MIN,
METRIC).toColumnName()), null,
+ List.of(new AggregationFunctionColumnPair(aggregationType,
METRIC).toColumnName()), null,
MAX_LEAF_RECORDS, nullHandlingEnabled);
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
