This is an automated email from the ASF dual-hosted git repository.
gortiz 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 c15ceadb2d1 Merge and serialize large intermediate results without
full-size copies (#19806)
c15ceadb2d1 is described below
commit c15ceadb2d1e56712686e8124178f7831cb83ac4
Author: Gonzalo Ortiz Jaureguizar <[email protected]>
AuthorDate: Thu Oct 15 15:56:17 2026 +0200
Merge and serialize large intermediate results without full-size copies
(#19806)
* Merge and serialize large intermediate results without full-size copies
Exact PERCENTILE keeps every value in a DoubleArrayList. For large groups
the
list, its serialized byte array and its deserialized copy are each G1
humongous
objects. Two of those copies exist only to move bytes between buffers.
Merge: add AggregationFunction#mergeSerializedIntermediateResult and
AggregationFunctionUtils#mergeSerialized. The default deserializes and calls
merge(). PercentileAggregationFunction overrides it to read the doubles from
the CustomObject buffer straight into the merged list, with the same growth
as addAll(). MultistageAggregationExecutor and MultistageGroupByExecutor
merge
OBJECT columns of serialized blocks through it, instead of deserializing the
whole column into an Object[] first. Row heap blocks and the
leafReturnFinalResult path are unchanged.
Serialize: SerializedIntermediateResult gains getSize() and
writeTo(DataOutput), backed by the byte array by default. DataBlockBuilder
and
BaseDataTableBuilder write through them. PercentileAggregationFunction
returns
a subclass that writes the count and the doubles straight into the
destination stream, so the builders no longer need the full-size byte array.
getBytes() still materializes it for other callers. The wire format is
unchanged.
Measured on a list of 2^20 doubles (8 MB payload), as thread-allocated
bytes:
building a data block drops from 16.9 MB to 8.5 MB (the paged output only),
and merging into a list with room drops from 8.4 MB to about 0.
* Write and read primitive arrays in bulk for intermediate results
PinotOutputStream gains writeInts, writeLongs and writeDoubles for array
ranges. The default encodes chunks of at most 8 KB through big-endian typed
ByteBuffer views and writes each chunk with write(byte[], int, int).
PagedPinotOutputStream overrides them to put the values straight into the
current page through the same views. A value that does not fit in the rest
of
a page is written with writeInt/writeLong and split across the two pages:
page offsets depend on everything written before, so this happens with any
page size. The new PrimitiveArrayOutput helper picks the bulk path for a
PinotOutputStream and the chunked path for any other DataOutput, such as the
DataOutputStream behind data tables.
Doubles are written as raw bits, the same bytes as ByteBuffer.putDouble and
DOUBLE_ARRAY_LIST_SER_DE. DataOutput.writeDouble canonicalizes NaN, so the
previous per-value writeTo of PERCENTILE differed from the serializer for
NaN
payloads.
PERCENTILE writeTo uses the bulk writes. Its
mergeSerializedIntermediateResult
reads the doubles with one DoubleBuffer.get into the merged list.
INT/LONG/FLOAT/DOUBLE_ARRAY_LIST_SER_DE.deserialize(ByteBuffer) read with
one
typed-view get instead of a getX() loop, and advance the position exactly as
the loop did. They now reject a count larger than the remaining bytes before
allocating the array.
JMH on JDK 25 (Apple M4 Pro, 3 forks x 5 iterations, 2^20 values):
- Write doubles to PagedPinotOutputStream: 699 ops/s per value, 2,293 ops/s
bulk (3.3x). For ints: 385 vs 4,468 ops/s (11.6x).
- Write doubles to DataOutputStream: 214 ops/s per value, 2,140 ops/s
chunked
(10x).
- Deserialize doubles from a heap buffer: 1,935 ops/s per value, 6,416 ops/s
bulk (3.3x). For ints: 1,924 vs 12,821 ops/s (6.7x).
Allocation per operation is unchanged apart from the 8 KB chunk scratch.
* Read and write short primitive lists one value at a time
The bulk paths for primitive list intermediate results copy through
big-endian typed views (asDoubleBuffer().get/put and the like). On a
little-endian CPU, that copy swaps the bytes in Unsafe.copySwapMemory0,
a native call with a fixed cost of about 10-40 ns. A single-value
ByteBuffer access compiles to an inlined byte swap, so lists of a few
values got 2-9x slower than the per-value loop they replaced. GROUP BY
queries with many groups produce many such lists.
PrimitiveArrayOutput.MIN_BULK_VALUES (64) is the cutoff for the typed-view
paths:
- the IntArrayList, LongArrayList, FloatArrayList and DoubleArrayList
deserializers in ObjectSerDeUtils read shorter lists one value at a
time into the array they wrap;
- PercentileAggregationFunction.mergeSerializedIntermediateResult reads
shorter lists one value at a time into the merged list;
- PagedPinotOutputStream.writeInts/writeLongs/writeDoubles write one value
at a time when fewer values go to the current page.
The chunked writes to a plain DataOutput have a lower cutoff,
MIN_CHUNKED_VALUES (16), because the scratch buffer costs more than the
swap.
The tests cover lengths on both sides of each cutoff. The list
deserialize, array output and new PERCENTILE merge benchmarks run sizes
from 1 value up.
* Write short percentile lists in one call and snapshot kept serialized
results
- SerializedDoubleArrayList.writeTo writes lists shorter than
PrimitiveArrayOutput.MIN_BULK_VALUES as one
getBytes() array, like master. A write per value is a synchronized call
on a DataOutputStream, which made
data table writes of short lists 2-4.5x slower.
- PagedPinotOutputStream.writeInts/writeLongs/writeDoubles put the values
that fit in the current page with
absolute puts and update the page offset and the written size once,
instead of one writeLong call per value.
A value that straddles two pages still goes through writeInt/writeLong.
- GroupByUtils.serializeChunk takes a byte-array snapshot of each
serialized result. The derive tasks read it
later, so a result that serializes the live intermediate result on demand
would serialize it once per
grouping set. The SerializedIntermediateResult javadoc now says that a
caller that keeps the object must take
a snapshot.
* Write one-chunk percentile lists in one call only off PinotOutputStream
SerializedDoubleArrayList.writeTo wrote getBytes() for every list below
MIN_BULK_VALUES, so on the MSE path a short list never reached the
per-value puts of PagedPinotOutputStream. On the SSE data table path,
lists from 64 values up went through the chunked path, which saves no
memory while the list fits in one chunk and was slower than master.
writeTo now writes getBytes() only when the output is not a
PinotOutputStream and the list fits in one PrimitiveArrayOutput.CHUNK_SIZE
chunk. A PinotOutputStream gets the values straight into its pages at
any size.
mergeSerializedIntermediateResult never caps the target capacity below
the new size: the growth by half is capped at Integer.MAX_VALUE - 8, and
the new size, computed with Math.addExact, is the lower bound.
SerializedIntermediateResult.getSize throws an IllegalStateException
naming the class when an on-demand subclass does not override it,
instead of a bare NullPointerException.
PercentileIntermediateResultSerDeTest covers 1023, 1024 and 5000 values
and also writes each list to a PagedPinotOutputStream with small pages.
---------
Co-authored-by: Gonzalo Ortiz <[email protected]>
---
.../apache/pinot/core/common/ObjectSerDeUtils.java | 72 ++++--
.../core/common/datablock/DataBlockBuilder.java | 9 +-
.../common/datatable/BaseDataTableBuilder.java | 8 +-
.../aggregation/function/AggregationFunction.java | 51 ++++
.../function/AggregationFunctionUtils.java | 11 +
.../function/PercentileAggregationFunction.java | 88 ++++++-
.../org/apache/pinot/core/util/GroupByUtils.java | 6 +-
.../common/ObjectSerDeUtilsPrimitiveListTest.java | 263 +++++++++++++++++++
.../PercentileIntermediateResultSerDeTest.java | 281 ++++++++++++++++++++
.../perf/BenchmarkPercentileSerializedMerge.java | 98 +++++++
.../pinot/perf/BenchmarkPrimitiveArrayOutput.java | 159 ++++++++++++
.../perf/BenchmarkPrimitiveListDeserialize.java | 130 ++++++++++
.../query/runtime/operator/AggregateOperator.java | 24 +-
.../operator/MultistageAggregationExecutor.java | 11 +
.../operator/MultistageGroupByExecutor.java | 19 +-
.../SerializedIntermediateResultMergeTest.java | 141 +++++++++++
.../segment/spi/memory/PagedPinotOutputStream.java | 94 +++++++
.../segment/spi/memory/PinotOutputStream.java | 22 ++
.../segment/spi/memory/PrimitiveArrayOutput.java | 143 +++++++++++
.../spi/memory/PrimitiveArrayOutputTest.java | 282 +++++++++++++++++++++
20 files changed, 1878 insertions(+), 34 deletions(-)
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/common/ObjectSerDeUtils.java
b/pinot-core/src/main/java/org/apache/pinot/core/common/ObjectSerDeUtils.java
index b7b8063f982..c5427f76065 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/common/ObjectSerDeUtils.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/common/ObjectSerDeUtils.java
@@ -56,6 +56,7 @@ import it.unimi.dsi.fastutil.objects.ObjectSet;
import java.io.IOException;
import java.lang.foreign.MemorySegment;
import java.math.BigDecimal;
+import java.nio.BufferUnderflowException;
import java.nio.ByteBuffer;
import java.nio.IntBuffer;
import java.util.ArrayList;
@@ -98,6 +99,7 @@ import
org.apache.pinot.segment.local.customobject.tdigest.SerializedTDigestInpu
import org.apache.pinot.segment.local.customobject.tdigest.TDigest;
import org.apache.pinot.segment.local.customobject.tdigest.TDigestAccumulator;
import org.apache.pinot.segment.local.utils.GeometrySerializer;
+import org.apache.pinot.segment.spi.memory.PrimitiveArrayOutput;
import org.apache.pinot.spi.utils.BigDecimalUtils;
import org.apache.pinot.spi.utils.ByteArray;
import org.locationtech.jts.geom.Geometry;
@@ -395,6 +397,20 @@ public class ObjectSerDeUtils {
}
};
+ // The list deserializers below read short lists one value at a time and
longer ones in bulk; see
+ // PrimitiveArrayOutput.MIN_BULK_VALUES for the cutoff.
+
+ /// Checks that `byteBuffer` holds `numValues` values of `width` bytes,
before the array for them is allocated. Throws
+ /// the same exceptions as reading the values one by one into a list of
capacity `numValues` would.
+ private static void checkRemaining(ByteBuffer byteBuffer, int numValues, int
width) {
+ if (numValues < 0) {
+ throw new IllegalArgumentException("Initial capacity (" + numValues + ")
is negative");
+ }
+ if (numValues > byteBuffer.remaining() / width) {
+ throw new BufferUnderflowException();
+ }
+ }
+
public static final ObjectSerDe<IntArrayList> INT_ARRAY_LIST_SER_DE = new
ObjectSerDe<IntArrayList>() {
@Override
@@ -418,11 +434,17 @@ public class ObjectSerDeUtils {
@Override
public IntArrayList deserialize(ByteBuffer byteBuffer) {
int numValues = byteBuffer.getInt();
- IntArrayList intArrayList = new IntArrayList(numValues);
- for (int i = 0; i < numValues; i++) {
- intArrayList.add(byteBuffer.getInt());
+ checkRemaining(byteBuffer, numValues, Integer.BYTES);
+ int[] values = new int[numValues];
+ if (numValues < PrimitiveArrayOutput.MIN_BULK_VALUES) {
+ for (int i = 0; i < numValues; i++) {
+ values[i] = byteBuffer.getInt();
+ }
+ } else {
+ byteBuffer.asIntBuffer().get(values);
+ byteBuffer.position(byteBuffer.position() + numValues * Integer.BYTES);
}
- return intArrayList;
+ return IntArrayList.wrap(values);
}
};
@@ -449,11 +471,17 @@ public class ObjectSerDeUtils {
@Override
public LongArrayList deserialize(ByteBuffer byteBuffer) {
int numValues = byteBuffer.getInt();
- LongArrayList longArrayList = new LongArrayList(numValues);
- for (int i = 0; i < numValues; i++) {
- longArrayList.add(byteBuffer.getLong());
+ checkRemaining(byteBuffer, numValues, Long.BYTES);
+ long[] values = new long[numValues];
+ if (numValues < PrimitiveArrayOutput.MIN_BULK_VALUES) {
+ for (int i = 0; i < numValues; i++) {
+ values[i] = byteBuffer.getLong();
+ }
+ } else {
+ byteBuffer.asLongBuffer().get(values);
+ byteBuffer.position(byteBuffer.position() + numValues * Long.BYTES);
}
- return longArrayList;
+ return LongArrayList.wrap(values);
}
};
@@ -480,11 +508,17 @@ public class ObjectSerDeUtils {
@Override
public FloatArrayList deserialize(ByteBuffer byteBuffer) {
int numValues = byteBuffer.getInt();
- FloatArrayList floatArrayList = new FloatArrayList(numValues);
- for (int i = 0; i < numValues; i++) {
- floatArrayList.add(byteBuffer.getFloat());
+ checkRemaining(byteBuffer, numValues, Float.BYTES);
+ float[] values = new float[numValues];
+ if (numValues < PrimitiveArrayOutput.MIN_BULK_VALUES) {
+ for (int i = 0; i < numValues; i++) {
+ values[i] = byteBuffer.getFloat();
+ }
+ } else {
+ byteBuffer.asFloatBuffer().get(values);
+ byteBuffer.position(byteBuffer.position() + numValues * Float.BYTES);
}
- return floatArrayList;
+ return FloatArrayList.wrap(values);
}
};
@@ -511,11 +545,17 @@ public class ObjectSerDeUtils {
@Override
public DoubleArrayList deserialize(ByteBuffer byteBuffer) {
int numValues = byteBuffer.getInt();
- DoubleArrayList doubleArrayList = new DoubleArrayList(numValues);
- for (int i = 0; i < numValues; i++) {
- doubleArrayList.add(byteBuffer.getDouble());
+ checkRemaining(byteBuffer, numValues, Double.BYTES);
+ double[] values = new double[numValues];
+ if (numValues < PrimitiveArrayOutput.MIN_BULK_VALUES) {
+ for (int i = 0; i < numValues; i++) {
+ values[i] = byteBuffer.getDouble();
+ }
+ } else {
+ byteBuffer.asDoubleBuffer().get(values);
+ byteBuffer.position(byteBuffer.position() + numValues * Double.BYTES);
}
- return doubleArrayList;
+ return DoubleArrayList.wrap(values);
}
};
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/common/datablock/DataBlockBuilder.java
b/pinot-core/src/main/java/org/apache/pinot/core/common/datablock/DataBlockBuilder.java
index 6e418e572c4..6f0c0c40d4a 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/common/datablock/DataBlockBuilder.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/common/datablock/DataBlockBuilder.java
@@ -705,11 +705,10 @@ public class DataBlockBuilder {
AggregationFunction.SerializedIntermediateResult value)
throws IOException {
writeVarOffsetInFixed(fixedSize, varSize);
- int type = value.getType();
- byte[] bytes = value.getBytes();
- fixedSize.putInt(bytes.length);
- varSize.writeInt(type);
- varSize.write(bytes);
+ fixedSize.putInt(value.getSize());
+ varSize.writeInt(value.getType());
+ // Writes straight into the pages, so a large intermediate result never
needs a full-size byte array.
+ value.writeTo(varSize);
}
private static void setNull(ByteBuffer fixedSize, PagedPinotOutputStream
varSize)
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/BaseDataTableBuilder.java
b/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/BaseDataTableBuilder.java
index 441a6e4b70b..ca62f32c723 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/BaseDataTableBuilder.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/BaseDataTableBuilder.java
@@ -222,11 +222,9 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
- int type = value.getType();
- byte[] bytes = value.getBytes();
- _currentRowDataByteBuffer.putInt(bytes.length);
- _variableSizeDataOutputStream.writeInt(type);
- _variableSizeDataByteArrayOutputStream.write(bytes);
+ _currentRowDataByteBuffer.putInt(value.getSize());
+ _variableSizeDataOutputStream.writeInt(value.getType());
+ value.writeTo(_variableSizeDataOutputStream);
}
@Override
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java
index 94be372fd54..c11254ad96d 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java
@@ -18,6 +18,8 @@
*/
package org.apache.pinot.core.query.aggregation.function;
+import java.io.DataOutput;
+import java.io.IOException;
import java.util.List;
import java.util.Map;
import javax.annotation.Nullable;
@@ -199,6 +201,28 @@ public interface AggregationFunction<IntermediateResult,
FinalResult extends Com
/// [AggregationFunction].
IntermediateResult merge(IntermediateResult intermediateResult1,
IntermediateResult intermediateResult2);
+ /// Merges a serialized intermediate result into `intermediateResult`, which
may be `null`, and returns the merged
+ /// result. The result may be `intermediateResult` itself, modified in place.
+ ///
+ /// The default deserializes `serialized` and merges it with [#merge].
Override it when the function can read the
+ /// serialized form straight into `intermediateResult`, to skip a full-size
copy of the incoming value. This matters
+ /// for large intermediate results, such as the value list of an exact
percentile, which are humongous objects for G1.
+ ///
+ /// Call [AggregationFunctionUtils#mergeSerialized] instead of calling this
directly: it settles a `null` serialized
+ /// value.
+ @Nullable
+ default IntermediateResult mergeSerializedIntermediateResult(@Nullable
IntermediateResult intermediateResult,
+ CustomObject serialized) {
+ IntermediateResult incoming = deserializeIntermediateResult(serialized);
+ if (intermediateResult == null) {
+ return incoming;
+ }
+ if (incoming == null) {
+ return intermediateResult;
+ }
+ return merge(intermediateResult, incoming);
+ }
+
/// Returns the [ColumnDataType] of the intermediate result.
///
/// This column data type is used for transferring data in data table.
@@ -211,6 +235,12 @@ public interface AggregationFunction<IntermediateResult,
FinalResult extends Com
}
/// Serialized intermediate result. Type can be used to identify the
intermediate result type when deserializing it.
+ ///
+ /// Data table and data block builders write the value with [#getSize()] and
[#writeTo(DataOutput)], and do not call
+ /// [#getBytes()]. A subclass can therefore write a large intermediate
result straight into the destination, without
+ /// building the full byte array first. Such a subclass reads the
intermediate result when it is written, so it must
+ /// be written before the intermediate result changes. A caller that keeps
the object for later must take a snapshot
+ /// instead: `new SerializedIntermediateResult(serialized.getType(),
serialized.getBytes())`.
class SerializedIntermediateResult {
private final int _type;
private final byte[] _bytes;
@@ -220,6 +250,13 @@ public interface AggregationFunction<IntermediateResult,
FinalResult extends Com
_bytes = buffer;
}
+ /// Constructor for subclasses that serialize on demand. They must
override [#getBytes()], [#getSize()] and
+ /// [#writeTo(DataOutput)].
+ protected SerializedIntermediateResult(int type) {
+ _type = type;
+ _bytes = null;
+ }
+
public int getType() {
return _type;
}
@@ -227,6 +264,20 @@ public interface AggregationFunction<IntermediateResult,
FinalResult extends Com
public byte[] getBytes() {
return _bytes;
}
+
+ /// Returns the number of bytes [#writeTo(DataOutput)] writes, which is
the length of [#getBytes()].
+ public int getSize() {
+ if (_bytes == null) {
+ throw new IllegalStateException(getClass().getName() + " serializes on
demand and must override getSize()");
+ }
+ return _bytes.length;
+ }
+
+ /// Writes the bytes of [#getBytes()] to `output`.
+ public void writeTo(DataOutput output)
+ throws IOException {
+ output.write(_bytes);
+ }
}
/// Deserializes the intermediate result from the custom object. This method
should be implemented if the intermediate
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java
index 07c433112a6..c52f772cb14 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java
@@ -133,6 +133,17 @@ public class AggregationFunctionUtils {
return expressions;
}
+ /// Merges a serialized intermediate result into `intermediateResult` (see
+ /// [AggregationFunction#mergeSerializedIntermediateResult]). A `null`
serialized value does not change the result.
+ @Nullable
+ public static <I> I mergeSerialized(AggregationFunction<I, ?>
aggregationFunction, @Nullable I intermediateResult,
+ @Nullable CustomObject serialized) {
+ if (serialized == null) {
+ return intermediateResult;
+ }
+ return
aggregationFunction.mergeSerializedIntermediateResult(intermediateResult,
serialized);
+ }
+
/// Merges two intermediate results, either of which may be `null`.
///
/// A `null` intermediate result means nothing was aggregated, and is thus
the identity of merging, which means the
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunction.java
index 04dd79bcc10..f2952e00e9c 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunction.java
@@ -19,6 +19,10 @@
package org.apache.pinot.core.query.aggregation.function;
import it.unimi.dsi.fastutil.doubles.DoubleArrayList;
+import java.io.DataOutput;
+import java.io.IOException;
+import java.nio.BufferUnderflowException;
+import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.Map;
import javax.annotation.Nullable;
@@ -32,6 +36,8 @@ import
org.apache.pinot.core.query.aggregation.ObjectAggregationResultHolder;
import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder;
import
org.apache.pinot.core.query.aggregation.groupby.ObjectGroupByResultHolder;
import org.apache.pinot.segment.spi.AggregationFunctionType;
+import org.apache.pinot.segment.spi.memory.PinotOutputStream;
+import org.apache.pinot.segment.spi.memory.PrimitiveArrayOutput;
public class PercentileAggregationFunction extends
BaseSingleInputAggregationFunction<DoubleArrayList, Double> {
@@ -206,10 +212,88 @@ public class PercentileAggregationFunction extends
BaseSingleInputAggregationFun
return ColumnDataType.OBJECT;
}
+ @Override
+ public DoubleArrayList mergeSerializedIntermediateResult(@Nullable
DoubleArrayList intermediateResult,
+ CustomObject serialized) {
+ if (intermediateResult == null) {
+ return deserializeIntermediateResult(serialized);
+ }
+ // Reads the values in bulk straight into the merged list, in the format of
+ // ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE: an int count, then the
doubles, in the buffer's byte order. Same
+ // growth as addAll(), without the incoming list.
+ ByteBuffer buffer = serialized.getBuffer();
+ int numValues = buffer.getInt();
+ if (numValues < 0) {
+ throw new IllegalArgumentException("Negative number of values: " +
numValues);
+ }
+ if (numValues > buffer.remaining() / Double.BYTES) {
+ throw new BufferUnderflowException();
+ }
+ int size = intermediateResult.size();
+ int newSize = Math.addExact(size, numValues);
+ int capacity = intermediateResult.elements().length;
+ if (newSize > capacity) {
+ // Grow by half like addAll(), capped below the maximum array size, but
never below the size set next.
+ int grownCapacity = (int) Math.min((long) capacity + (capacity >> 1),
Integer.MAX_VALUE - 8);
+ intermediateResult.ensureCapacity(Math.max(newSize, grownCapacity));
+ }
+ // size() zero-fills the new range, so it must run before the values are
copied in.
+ intermediateResult.size(newSize);
+ double[] elements = intermediateResult.elements();
+ if (numValues < PrimitiveArrayOutput.MIN_BULK_VALUES) {
+ for (int i = size; i < newSize; i++) {
+ elements[i] = buffer.getDouble();
+ }
+ } else {
+ buffer.asDoubleBuffer().get(elements, size, numValues);
+ buffer.position(buffer.position() + numValues * Double.BYTES);
+ }
+ return intermediateResult;
+ }
+
@Override
public SerializedIntermediateResult
serializeIntermediateResult(DoubleArrayList doubleArrayList) {
- return new
SerializedIntermediateResult(ObjectSerDeUtils.ObjectType.DoubleArrayList.getValue(),
- ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE.serialize(doubleArrayList));
+ return new SerializedDoubleArrayList(doubleArrayList);
+ }
+
+ /// Serializes a [DoubleArrayList] on demand, with the bytes of
ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE: a
+ /// big-endian int count, then the big-endian doubles. [#writeTo] writes the
values in bulk straight into the
+ /// destination (see [PrimitiveArrayOutput]), so the builders never hold the
full-size byte array. On any other
+ /// [DataOutput] than a [PinotOutputStream], a list that fits in one
[PrimitiveArrayOutput#CHUNK_SIZE] chunk is
+ /// written as one small array. Not thread-safe; the list must not change
until it is written.
+ private static final class SerializedDoubleArrayList extends
SerializedIntermediateResult {
+ private final DoubleArrayList _values;
+
+ SerializedDoubleArrayList(DoubleArrayList values) {
+ super(ObjectSerDeUtils.ObjectType.DoubleArrayList.getValue());
+ _values = values;
+ }
+
+ @Override
+ public byte[] getBytes() {
+ return ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE.serialize(_values);
+ }
+
+ @Override
+ public int getSize() {
+ return Math.toIntExact(Integer.BYTES + (long) _values.size() *
Double.BYTES);
+ }
+
+ @Override
+ public void writeTo(DataOutput output)
+ throws IOException {
+ int size = _values.size();
+ if (!(output instanceof PinotOutputStream)
+ && Integer.BYTES + (long) size * Double.BYTES <=
PrimitiveArrayOutput.CHUNK_SIZE) {
+ // The chunked path saves no memory while the list fits in one chunk,
and one write of a small array is
+ // cheaper than chunk encoding or a write per value (a synchronized
call on a DataOutputStream). A
+ // PinotOutputStream writes the values straight into its pages, at any
size.
+ output.write(getBytes());
+ return;
+ }
+ output.writeInt(size);
+ PrimitiveArrayOutput.writeDoubles(output, _values.elements(), 0, size);
+ }
}
@Override
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java
b/pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java
index ca8ca72f464..492c690d77b 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java
@@ -358,7 +358,11 @@ public final class GroupByUtils {
Object intermediate = baseValues[numUnionColumns + i];
if (intermediate != null
&& aggregationFunctions[i].getIntermediateResultColumnType() ==
ColumnDataType.OBJECT) {
- row[i] =
aggregationFunctions[i].serializeIntermediateResult(intermediate);
+ // The derive tasks read these bytes later, so take a snapshot
instead of keeping a result that serializes
+ // the live intermediate result on demand.
+ SerializedIntermediateResult serializedResult =
+
aggregationFunctions[i].serializeIntermediateResult(intermediate);
+ row[i] = new
SerializedIntermediateResult(serializedResult.getType(),
serializedResult.getBytes());
}
}
serialized[e] = row;
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/common/ObjectSerDeUtilsPrimitiveListTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/common/ObjectSerDeUtilsPrimitiveListTest.java
new file mode 100644
index 00000000000..af53cb22479
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/core/common/ObjectSerDeUtilsPrimitiveListTest.java
@@ -0,0 +1,263 @@
+/**
+ * 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.common;
+
+import it.unimi.dsi.fastutil.doubles.DoubleArrayList;
+import it.unimi.dsi.fastutil.floats.FloatArrayList;
+import it.unimi.dsi.fastutil.ints.IntArrayList;
+import it.unimi.dsi.fastutil.longs.LongArrayList;
+import java.nio.BufferUnderflowException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
+import java.util.function.Function;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
+
+
+/// Checks that the bulk reads of the primitive array list SerDes return the
same values and leave the buffer at the
+/// same position as reading the values one by one, for heap, direct,
little-endian and sliced buffers.
+public class ObjectSerDeUtilsPrimitiveListTest {
+ private static final int PREFIX = 3;
+ private static final int SUFFIX = 5;
+
+ /// A list SerDe with a per-value oracle reader.
+ private enum ListType {
+ INT(Integer.BYTES) {
+ @Override
+ Object newList(Random random, int size) {
+ return new IntArrayList(random.ints(size).toArray());
+ }
+
+ @Override
+ ObjectSerDeUtils.ObjectSerDe<?> serDe() {
+ return ObjectSerDeUtils.INT_ARRAY_LIST_SER_DE;
+ }
+
+ @Override
+ void putValue(ByteBuffer buffer, Object list, int index) {
+ buffer.putInt(((IntArrayList) list).getInt(index));
+ }
+
+ @Override
+ Object readOneByOne(ByteBuffer buffer) {
+ int numValues = buffer.getInt();
+ IntArrayList list = new IntArrayList(numValues);
+ for (int i = 0; i < numValues; i++) {
+ list.add(buffer.getInt());
+ }
+ return list;
+ }
+ },
+ LONG(Long.BYTES) {
+ @Override
+ Object newList(Random random, int size) {
+ return new LongArrayList(random.longs(size).toArray());
+ }
+
+ @Override
+ ObjectSerDeUtils.ObjectSerDe<?> serDe() {
+ return ObjectSerDeUtils.LONG_ARRAY_LIST_SER_DE;
+ }
+
+ @Override
+ void putValue(ByteBuffer buffer, Object list, int index) {
+ buffer.putLong(((LongArrayList) list).getLong(index));
+ }
+
+ @Override
+ Object readOneByOne(ByteBuffer buffer) {
+ int numValues = buffer.getInt();
+ LongArrayList list = new LongArrayList(numValues);
+ for (int i = 0; i < numValues; i++) {
+ list.add(buffer.getLong());
+ }
+ return list;
+ }
+ },
+ FLOAT(Float.BYTES) {
+ @Override
+ Object newList(Random random, int size) {
+ FloatArrayList list = new FloatArrayList(size);
+ for (int i = 0; i < size; i++) {
+ list.add(i == 0 ? Float.intBitsToFloat(0x7fc0_0123) :
random.nextFloat() - 0.5f);
+ }
+ return list;
+ }
+
+ @Override
+ ObjectSerDeUtils.ObjectSerDe<?> serDe() {
+ return ObjectSerDeUtils.FLOAT_ARRAY_LIST_SER_DE;
+ }
+
+ @Override
+ void putValue(ByteBuffer buffer, Object list, int index) {
+ buffer.putFloat(((FloatArrayList) list).getFloat(index));
+ }
+
+ @Override
+ Object readOneByOne(ByteBuffer buffer) {
+ int numValues = buffer.getInt();
+ FloatArrayList list = new FloatArrayList(numValues);
+ for (int i = 0; i < numValues; i++) {
+ list.add(buffer.getFloat());
+ }
+ return list;
+ }
+ },
+ DOUBLE(Double.BYTES) {
+ @Override
+ Object newList(Random random, int size) {
+ DoubleArrayList list = new DoubleArrayList(size);
+ for (int i = 0; i < size; i++) {
+ // A NaN with a payload keeps its raw bits on both paths.
+ list.add(i == 0 ? Double.longBitsToDouble(0x7ff8_0000_0000_0123L) :
random.nextDouble() - 0.5);
+ }
+ return list;
+ }
+
+ @Override
+ ObjectSerDeUtils.ObjectSerDe<?> serDe() {
+ return ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE;
+ }
+
+ @Override
+ void putValue(ByteBuffer buffer, Object list, int index) {
+ buffer.putDouble(((DoubleArrayList) list).getDouble(index));
+ }
+
+ @Override
+ Object readOneByOne(ByteBuffer buffer) {
+ int numValues = buffer.getInt();
+ DoubleArrayList list = new DoubleArrayList(numValues);
+ for (int i = 0; i < numValues; i++) {
+ list.add(buffer.getDouble());
+ }
+ return list;
+ }
+ };
+
+ final int _width;
+
+ ListType(int width) {
+ _width = width;
+ }
+
+ abstract Object newList(Random random, int size);
+
+ abstract ObjectSerDeUtils.ObjectSerDe<?> serDe();
+
+ abstract void putValue(ByteBuffer buffer, Object list, int index);
+
+ abstract Object readOneByOne(ByteBuffer buffer);
+ }
+
+ /// How the serialized bytes are laid out in the buffer that is read.
+ private enum Layout {
+ HEAP(size -> ByteBuffer.allocate(size)),
+ DIRECT(ByteBuffer::allocateDirect),
+ // The SerDes read in the buffer's byte order, whatever it is.
+ LITTLE_ENDIAN(size ->
ByteBuffer.allocate(size).order(ByteOrder.LITTLE_ENDIAN));
+
+ final Function<Integer, ByteBuffer> _allocator;
+
+ Layout(Function<Integer, ByteBuffer> allocator) {
+ _allocator = allocator;
+ }
+ }
+
+ @DataProvider
+ public Object[][] cases() {
+ List<Object[]> cases = new ArrayList<>();
+ for (ListType type : ListType.values()) {
+ for (Layout layout : Layout.values()) {
+ for (int size : new int[]{0, 1, 63, 64, 1000}) {
+ for (boolean slice : new boolean[]{false, true}) {
+ cases.add(new Object[]{type, layout, size, slice});
+ }
+ }
+ }
+ }
+ return cases.toArray(new Object[0][]);
+ }
+
+ /// Writes `size` values with surrounding bytes, the way a Map or List SerDe
nests values: [PREFIX bytes][count]
+ /// [values][SUFFIX bytes]. The returned buffer is positioned at the count.
+ private static ByteBuffer write(ListType type, Layout layout, Object list,
int size, boolean slice) {
+ ByteBuffer buffer = layout._allocator.apply(PREFIX + Integer.BYTES + size
* type._width + SUFFIX);
+ for (int i = 0; i < PREFIX; i++) {
+ buffer.put((byte) 0x55);
+ }
+ buffer.putInt(size);
+ for (int i = 0; i < size; i++) {
+ type.putValue(buffer, list, i);
+ }
+ while (buffer.hasRemaining()) {
+ buffer.put((byte) 0x66);
+ }
+ buffer.position(PREFIX);
+ return slice ? buffer.slice().order(buffer.order()) : buffer;
+ }
+
+ @Test(dataProvider = "cases")
+ public void testBulkReadMatchesOneByOne(ListType type, Layout layout, int
size, boolean slice) {
+ Object list = type.newList(new Random(size), size);
+ ByteBuffer bulkBuffer = write(type, layout, list, size, slice);
+ ByteBuffer oracleBuffer = bulkBuffer.duplicate().order(bulkBuffer.order());
+
+ @SuppressWarnings("unchecked")
+ ObjectSerDeUtils.ObjectSerDe<Object> serDe =
(ObjectSerDeUtils.ObjectSerDe<Object>) type.serDe();
+ Object actual = serDe.deserialize(bulkBuffer);
+ Object expected = type.readOneByOne(oracleBuffer);
+
+ // Compares the serialized bytes, so floating point values must match bit
for bit, NaN payloads included.
+ assertEquals(serDe.serialize(actual), serDe.serialize(expected));
+ assertEquals(bulkBuffer.position(), oracleBuffer.position());
+ assertEquals(bulkBuffer.remaining(), SUFFIX);
+ if (layout != Layout.LITTLE_ENDIAN) {
+ assertEquals(serDe.serialize(actual), serDe.serialize(list));
+ }
+ }
+
+ @Test(dataProvider = "cases")
+ public void testSerializeThenDeserializeRoundTrips(ListType type, Layout
layout, int size, boolean slice) {
+ Object list = type.newList(new Random(size), size);
+ @SuppressWarnings("unchecked")
+ ObjectSerDeUtils.ObjectSerDe<Object> serDe =
(ObjectSerDeUtils.ObjectSerDe<Object>) type.serDe();
+ byte[] bytes = serDe.serialize(list);
+ assertEquals(serDe.serialize(serDe.deserialize(bytes)), bytes);
+ }
+
+ @Test
+ public void testTruncatedAndNegativeCounts() {
+ for (ListType type : ListType.values()) {
+ ByteBuffer truncated = ByteBuffer.allocate(Integer.BYTES + 2 *
type._width).putInt(3);
+ truncated.flip().limit(truncated.capacity());
+ assertThrows(BufferUnderflowException.class, () ->
type.serDe().deserialize(truncated));
+
+ ByteBuffer negative = ByteBuffer.allocate(Integer.BYTES).putInt(0, -1);
+ assertThrows(IllegalArgumentException.class, () ->
type.serDe().deserialize(negative));
+ }
+ }
+}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileIntermediateResultSerDeTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileIntermediateResultSerDeTest.java
new file mode 100644
index 00000000000..122e01607f1
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileIntermediateResultSerDeTest.java
@@ -0,0 +1,281 @@
+/**
+ * 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 com.sun.management.ThreadMXBean;
+import it.unimi.dsi.fastutil.doubles.DoubleArrayList;
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.lang.management.ManagementFactory;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.pinot.common.CustomObject;
+import org.apache.pinot.common.datablock.DataBlock;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.common.ObjectSerDeUtils;
+import org.apache.pinot.core.common.datablock.DataBlockBuilder;
+import org.apache.pinot.core.common.datatable.DataTableBuilder;
+import org.apache.pinot.core.common.datatable.DataTableBuilderFactory;
+import org.apache.pinot.segment.spi.memory.PagedPinotOutputStream;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertTrue;
+
+
+/// Checks that PERCENTILE writes its intermediate result straight into data
blocks and data tables, and merges it
+/// straight from its serialized form, with the same bytes and results as the
byte array based defaults.
+public class PercentileIntermediateResultSerDeTest {
+ private static final DataSchema SCHEMA = new DataSchema(new String[]{"key",
"percentile"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.OBJECT});
+ /// 8 MB of doubles: a humongous object for G1 with any region size up to 16
MB.
+ private static final int LARGE_SIZE = 1 << 20;
+
+ private final PercentileAggregationFunction _function = newFunction();
+ /// The behavior before the change: serializes to a full byte array, and
merges the deserialized value.
+ private final PercentileAggregationFunction _byteArrayFunction =
+ new PercentileAggregationFunction(ExpressionContext.forIdentifier("$1"),
50, true) {
+ @Override
+ public SerializedIntermediateResult
serializeIntermediateResult(DoubleArrayList values) {
+ return new
SerializedIntermediateResult(ObjectSerDeUtils.ObjectType.DoubleArrayList.getValue(),
+ ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE.serialize(values));
+ }
+
+ @Override
+ public DoubleArrayList
mergeSerializedIntermediateResult(DoubleArrayList intermediateResult,
+ CustomObject serialized) {
+ DoubleArrayList incoming = deserializeIntermediateResult(serialized);
+ return intermediateResult == null ? incoming :
merge(intermediateResult, incoming);
+ }
+ };
+
+ private static PercentileAggregationFunction newFunction() {
+ return new
PercentileAggregationFunction(ExpressionContext.forIdentifier("$1"), 50, true);
+ }
+
+ private static DoubleArrayList sequence(int size) {
+ DoubleArrayList values = new DoubleArrayList(size);
+ for (int i = 0; i < size; i++) {
+ values.add(i * 0.5 - 7);
+ }
+ return values;
+ }
+
+ private static CustomObject serialize(DoubleArrayList values) {
+ return new
CustomObject(ObjectSerDeUtils.ObjectType.DoubleArrayList.getValue(),
+
ByteBuffer.wrap(ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE.serialize(values)));
+ }
+
+ @DataProvider
+ public Object[][] mergeCases() {
+ return new Object[][]{
+ {null, sequence(3)},
+ {sequence(0), sequence(0)},
+ {sequence(0), sequence(5)},
+ {sequence(4), sequence(0)},
+ {sequence(10), sequence(10)},
+ // Fills the exact capacity, then grows past it.
+ {new DoubleArrayList(new double[]{1, 2}), sequence(1)},
+ {sequence(1000), sequence(3000)},
+ // Both sides of PrimitiveArrayOutput.MIN_BULK_VALUES: read one value
at a time, then in bulk.
+ {sequence(7), sequence(63)},
+ {sequence(7), sequence(64)},
+ {sequence(2), new DoubleArrayList(new
double[]{Double.longBitsToDouble(0x7ff8_0000_0000_0123L), -0.0})}
+ };
+ }
+
+ @Test(dataProvider = "mergeCases")
+ public void testMergeSerializedMatchesMergeOfDeserialized(DoubleArrayList
state, DoubleArrayList incoming) {
+ DoubleArrayList expectedState = state == null ? null : new
DoubleArrayList(state);
+ DoubleArrayList expected =
_byteArrayFunction.mergeSerializedIntermediateResult(expectedState,
serialize(incoming));
+
+ CustomObject serialized = serialize(incoming);
+ DoubleArrayList actual =
_function.mergeSerializedIntermediateResult(state, serialized);
+
+ // Compares the serialized bytes, so the values must match bit for bit,
NaN payloads included.
+ assertEquals(ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE.serialize(actual),
+ ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE.serialize(expected));
+ assertEquals(serialized.getBuffer().remaining(), 0, "Consumes the whole
value");
+ if (state != null) {
+ assertSame(actual, state, "Merges in place");
+ }
+ }
+
+ @Test
+ public void testMergeSerializedSkipsNull() {
+ DoubleArrayList state = sequence(2);
+ assertSame(AggregationFunctionUtils.mergeSerialized(_function, state,
null), state);
+ }
+
+ @Test
+ public void testSerializedIntermediateResultMatchesByteArray()
+ throws IOException {
+ // A NaN with a payload must keep its raw bits, like ByteBuffer.putDouble
in the byte array serializer.
+ DoubleArrayList specialValues = new DoubleArrayList(
+ new double[]{Double.NaN,
Double.longBitsToDouble(0x7ff8_0000_0000_0123L), -0.0,
Double.NEGATIVE_INFINITY});
+ // On a DataOutputStream, 1023 values are the most that fit in one
PrimitiveArrayOutput.CHUNK_SIZE chunk (written
+ // as one small array), and 1024 go through the chunked path. On a
PagedPinotOutputStream, both sides of
+ // PrimitiveArrayOutput.MIN_BULK_VALUES (63 and 64) switch from per-value
puts to the bulk view.
+ for (DoubleArrayList values : List.of(sequence(0), sequence(1),
sequence(63), sequence(64), sequence(1000),
+ sequence(1023), sequence(1024), sequence(5000), specialValues)) {
+ AggregationFunction.SerializedIntermediateResult expected =
+ _byteArrayFunction.serializeIntermediateResult(values);
+ AggregationFunction.SerializedIntermediateResult actual =
_function.serializeIntermediateResult(values);
+
+ assertEquals(actual.getType(), expected.getType());
+ assertEquals(actual.getSize(), expected.getBytes().length);
+ assertEquals(actual.getBytes(), expected.getBytes());
+ ByteArrayOutputStream written = new ByteArrayOutputStream();
+ actual.writeTo(new DataOutputStream(written));
+ assertEquals(written.toByteArray(), expected.getBytes());
+ // Small pages, so the bulk path also splits values across page
boundaries.
+ try (PagedPinotOutputStream paged =
+ new PagedPinotOutputStream(new
PagedPinotOutputStream.HeapPageAllocator(100))) {
+ actual.writeTo(paged);
+ assertEquals(concat(List.of(paged.getPages())), expected.getBytes());
+ }
+ }
+ }
+
+ private static List<Object[]> rows(int largeSize) {
+ List<Object[]> rows = new ArrayList<>();
+ rows.add(new Object[]{1, sequence(3)});
+ rows.add(new Object[]{2, null});
+ rows.add(new Object[]{3, sequence(0)});
+ // Spans several pages of the variable size output stream.
+ rows.add(new Object[]{4, sequence(largeSize)});
+ rows.add(new Object[]{5, sequence(17)});
+ return rows;
+ }
+
+ private static byte[] concat(List<ByteBuffer> buffers) {
+ ByteArrayOutputStream out = new ByteArrayOutputStream();
+ for (ByteBuffer buffer : buffers) {
+ byte[] bytes = new byte[buffer.remaining()];
+ buffer.duplicate().get(bytes);
+ out.writeBytes(bytes);
+ }
+ return out.toByteArray();
+ }
+
+ private static byte[] dataBlockBytes(List<Object[]> rows,
AggregationFunction function)
+ throws IOException {
+ DataBlock dataBlock = DataBlockBuilder.buildFromRows(rows, SCHEMA, new
AggregationFunction[]{function});
+ return concat(dataBlock.serialize());
+ }
+
+ /// Builds a columnar data block of the state column only: buildFromColumns
does not support key columns together
+ /// with aggregation functions.
+ private static byte[] columnarDataBlockBytes(List<Object[]> rows,
AggregationFunction function)
+ throws IOException {
+ Object[] values = new Object[rows.size()];
+ for (int i = 0; i < rows.size(); i++) {
+ values[i] = rows.get(i)[1];
+ }
+ DataSchema schema = new DataSchema(new String[]{"percentile"}, new
ColumnDataType[]{ColumnDataType.OBJECT});
+ DataBlock dataBlock =
DataBlockBuilder.buildFromColumns(List.<Object[]>of(values), schema,
+ new AggregationFunction[]{function});
+ return concat(dataBlock.serialize());
+ }
+
+ private static byte[] dataTableBytes(List<Object[]> rows,
AggregationFunction function)
+ throws IOException {
+ DataTableBuilder builder =
DataTableBuilderFactory.getDataTableBuilder(SCHEMA);
+ for (Object[] row : rows) {
+ builder.startRow();
+ builder.setColumn(0, (int) row[0]);
+ if (row[1] == null) {
+ builder.setNull(1);
+ } else {
+ builder.setColumn(1, function.serializeIntermediateResult(row[1]));
+ }
+ builder.finishRow();
+ }
+ return builder.build().toBytes();
+ }
+
+ @Test
+ public void testDataBlockAndDataTableBytesMatchByteArraySerialization()
+ throws IOException {
+ List<Object[]> rows = rows(100_000);
+ assertEquals(dataBlockBytes(rows, _function), dataBlockBytes(rows,
_byteArrayFunction));
+ assertEquals(columnarDataBlockBytes(rows, _function),
columnarDataBlockBytes(rows, _byteArrayFunction));
+ assertEquals(dataTableBytes(rows, _function), dataTableBytes(rows,
_byteArrayFunction));
+ }
+
+ /// Building a data block writes the large list into the paged output stream
(8 MB of small pages) without the 8 MB
+ /// byte array the byte array path allocates first.
+ @Test
+ public void testDataBlockBuildDoesNotAllocateFullSizeByteArray() {
+ List<Object[]> rows = new ArrayList<>();
+ rows.add(new Object[]{1, sequence(LARGE_SIZE)});
+ long payloadBytes = Integer.BYTES + (long) LARGE_SIZE * Double.BYTES;
+
+ long byteArrayAllocated = allocatedBytes(
+ () -> DataBlockBuilder.buildFromRows(rows, SCHEMA, new
AggregationFunction[]{_byteArrayFunction}));
+ long allocated =
+ allocatedBytes(() -> DataBlockBuilder.buildFromRows(rows, SCHEMA, new
AggregationFunction[]{_function}));
+
+ assertTrue(byteArrayAllocated > payloadBytes * 3 / 2, "byte array path: "
+ byteArrayAllocated);
+ assertTrue(allocated < payloadBytes * 3 / 2, "stream path: " + allocated);
+ }
+
+ /// Merging reads the doubles straight into the merged list, which already
has room for them, instead of
+ /// deserializing an 8 MB list first.
+ @Test
+ public void testMergeSerializedDoesNotAllocateIncomingList() {
+ CustomObject serialized = serialize(sequence(LARGE_SIZE));
+ long payloadBytes = serialized.getBuffer().remaining();
+
+ long byteArrayAllocated = allocatedBytes(() ->
_byteArrayFunction.mergeSerializedIntermediateResult(
+ new DoubleArrayList(LARGE_SIZE), new
CustomObject(serialized.getType(), serialized.getBuffer().duplicate())));
+ long allocated = allocatedBytes(() ->
_function.mergeSerializedIntermediateResult(
+ new DoubleArrayList(LARGE_SIZE), new
CustomObject(serialized.getType(), serialized.getBuffer().duplicate())));
+
+ // Both allocate the 8 MB state of the merge target.
+ assertTrue(byteArrayAllocated > payloadBytes * 3 / 2, "deserialize path: "
+ byteArrayAllocated);
+ assertTrue(allocated < payloadBytes * 3 / 2, "in place path: " +
allocated);
+ }
+
+ /// Heap bytes the current thread allocates while running `action`. A
warm-up run keeps class loading and lazy
+ /// initialization out of the measurement.
+ private static long allocatedBytes(ThrowingRunnable action) {
+ ThreadMXBean threadBean = (ThreadMXBean)
ManagementFactory.getThreadMXBean();
+ long threadId = Thread.currentThread().threadId();
+ try {
+ action.run();
+ long before = threadBean.getThreadAllocatedBytes(threadId);
+ action.run();
+ return threadBean.getThreadAllocatedBytes(threadId) - before;
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ private interface ThrowingRunnable {
+ void run()
+ throws Exception;
+ }
+}
diff --git
a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkPercentileSerializedMerge.java
b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkPercentileSerializedMerge.java
new file mode 100644
index 00000000000..0a08c2d79a6
--- /dev/null
+++
b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkPercentileSerializedMerge.java
@@ -0,0 +1,98 @@
+/**
+ * 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;
+
+import it.unimi.dsi.fastutil.doubles.DoubleArrayList;
+import java.nio.ByteBuffer;
+import java.util.Random;
+import java.util.concurrent.TimeUnit;
+import org.apache.pinot.common.CustomObject;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.core.common.ObjectSerDeUtils;
+import
org.apache.pinot.core.query.aggregation.function.PercentileAggregationFunction;
+import org.openjdk.jmh.annotations.Benchmark;
+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.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+
+/// Compares two ways to merge a serialized PERCENTILE intermediate result (a
DoubleArrayList) into a merged list that
+/// already holds `size` values:
+/// - `deserializeAndAddAll`: reads the incoming list one value at a time,
then calls `merge` (`addAll`).
+/// - `mergeSerialized`:
[PercentileAggregationFunction#mergeSerializedIntermediateResult], which reads
the values
+/// straight into the merged list (in bulk from
PrimitiveArrayOutput.MIN_BULK_VALUES values up).
+///
+/// Each operation first truncates the merged list back to `size` values, so
it never reallocates.
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@Fork(2)
+@Warmup(iterations = 3, time = 1)
+@Measurement(iterations = 5, time = 1)
+@State(Scope.Benchmark)
+public class BenchmarkPercentileSerializedMerge {
+ @Param({"1", "4", "16", "64", "1024", "65536"})
+ private int _size;
+
+ private PercentileAggregationFunction _function;
+ private ByteBuffer _serialized;
+ private DoubleArrayList _merged;
+
+ @Setup
+ public void setUp() {
+ _function = new
PercentileAggregationFunction(ExpressionContext.forIdentifier("col"), 50,
false);
+ double[] values = new Random(42).doubles(_size).toArray();
+ _serialized =
ByteBuffer.wrap(ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE.serialize(new
DoubleArrayList(values)));
+ _merged = new DoubleArrayList(2 * _size);
+ _merged.addElements(0, values);
+ }
+
+ @Benchmark
+ public DoubleArrayList deserializeAndAddAll() {
+ _merged.size(_size);
+ ByteBuffer buffer = _serialized.duplicate();
+ int numValues = buffer.getInt();
+ DoubleArrayList incoming = new DoubleArrayList(numValues);
+ for (int i = 0; i < numValues; i++) {
+ incoming.add(buffer.getDouble());
+ }
+ return _function.merge(_merged, incoming);
+ }
+
+ @Benchmark
+ public DoubleArrayList mergeSerialized() {
+ _merged.size(_size);
+ CustomObject serialized =
+ new
CustomObject(ObjectSerDeUtils.ObjectType.DoubleArrayList.getValue(),
_serialized.duplicate());
+ return _function.mergeSerializedIntermediateResult(_merged, serialized);
+ }
+
+ public static void main(String[] args)
+ throws Exception {
+ new Runner(new
OptionsBuilder().include(BenchmarkPercentileSerializedMerge.class.getSimpleName()).build()).run();
+ }
+}
diff --git
a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkPrimitiveArrayOutput.java
b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkPrimitiveArrayOutput.java
new file mode 100644
index 00000000000..b8daa3f7ff2
--- /dev/null
+++
b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkPrimitiveArrayOutput.java
@@ -0,0 +1,159 @@
+/**
+ * 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;
+
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutput;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.Random;
+import java.util.concurrent.TimeUnit;
+import org.apache.pinot.segment.spi.memory.PagedPinotOutputStream;
+import org.apache.pinot.segment.spi.memory.PrimitiveArrayOutput;
+import org.openjdk.jmh.annotations.Benchmark;
+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.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+
+/// Compares ways to write a primitive array to the outputs data blocks and
data tables use:
+/// - `*PerValue`: one writeInt/writeLong/writeDouble call per value.
+/// - `*Chunked`: [PrimitiveArrayOutput] through a plain [DataOutput], which
encodes small chunks into a scratch array.
+/// - `pagedBulk`: the [PagedPinotOutputStream] override, which puts the
values straight into the pages.
+///
+/// Each operation writes `size` values into a new output:
[PagedPinotOutputStream] with small heap pages (as data
+/// blocks use), or a [DataOutputStream] over a presized
[ByteArrayOutputStream] (as data tables use).
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Fork(3)
+@Warmup(iterations = 3, time = 1)
+@Measurement(iterations = 5, time = 1)
+@State(Scope.Benchmark)
+public class BenchmarkPrimitiveArrayOutput {
+ @Param({"INT", "LONG", "DOUBLE"})
+ private String _type;
+
+ @Param({"1", "4", "16", "64", "1024", "65536", "1048576"})
+ private int _size;
+
+ private int[] _ints;
+ private long[] _longs;
+ private double[] _doubles;
+ private int _numBytes;
+
+ @Setup
+ public void setUp() {
+ Random random = new Random(42);
+ _ints = random.ints(_size).toArray();
+ _longs = random.longs(_size).toArray();
+ _doubles = random.doubles(_size).toArray();
+ _numBytes = _size * ("INT".equals(_type) ? Integer.BYTES : Long.BYTES);
+ }
+
+ @Benchmark
+ public PagedPinotOutputStream pagedPerValue()
+ throws IOException {
+ PagedPinotOutputStream output = PagedPinotOutputStream.createHeap();
+ writePerValue(output);
+ return output;
+ }
+
+ @Benchmark
+ public PagedPinotOutputStream pagedChunked()
+ throws IOException {
+ PagedPinotOutputStream output = PagedPinotOutputStream.createHeap();
+ // A DataOutputStream wrapper is not a PinotOutputStream, so
PrimitiveArrayOutput uses the chunked encoding.
+ writeDispatched(new DataOutputStream(output));
+ return output;
+ }
+
+ @Benchmark
+ public PagedPinotOutputStream pagedBulk()
+ throws IOException {
+ PagedPinotOutputStream output = PagedPinotOutputStream.createHeap();
+ writeDispatched(output);
+ return output;
+ }
+
+ @Benchmark
+ public ByteArrayOutputStream dataOutputStreamPerValue()
+ throws IOException {
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream(_numBytes);
+ writePerValue(new DataOutputStream(bytes));
+ return bytes;
+ }
+
+ @Benchmark
+ public ByteArrayOutputStream dataOutputStreamChunked()
+ throws IOException {
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream(_numBytes);
+ writeDispatched(new DataOutputStream(bytes));
+ return bytes;
+ }
+
+ private void writePerValue(DataOutput output)
+ throws IOException {
+ switch (_type) {
+ case "INT":
+ for (int value : _ints) {
+ output.writeInt(value);
+ }
+ break;
+ case "LONG":
+ for (long value : _longs) {
+ output.writeLong(value);
+ }
+ break;
+ default:
+ for (double value : _doubles) {
+ output.writeDouble(value);
+ }
+ break;
+ }
+ }
+
+ private void writeDispatched(DataOutput output)
+ throws IOException {
+ switch (_type) {
+ case "INT":
+ PrimitiveArrayOutput.writeInts(output, _ints, 0, _size);
+ break;
+ case "LONG":
+ PrimitiveArrayOutput.writeLongs(output, _longs, 0, _size);
+ break;
+ default:
+ PrimitiveArrayOutput.writeDoubles(output, _doubles, 0, _size);
+ break;
+ }
+ }
+
+ public static void main(String[] args)
+ throws Exception {
+ new Runner(new
OptionsBuilder().include(BenchmarkPrimitiveArrayOutput.class.getSimpleName()).build()).run();
+ }
+}
diff --git
a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkPrimitiveListDeserialize.java
b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkPrimitiveListDeserialize.java
new file mode 100644
index 00000000000..569bc9855d1
--- /dev/null
+++
b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkPrimitiveListDeserialize.java
@@ -0,0 +1,130 @@
+/**
+ * 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;
+
+import it.unimi.dsi.fastutil.doubles.DoubleArrayList;
+import it.unimi.dsi.fastutil.ints.IntArrayList;
+import it.unimi.dsi.fastutil.longs.LongArrayList;
+import java.nio.ByteBuffer;
+import java.util.Random;
+import java.util.concurrent.TimeUnit;
+import org.apache.pinot.core.common.ObjectSerDeUtils;
+import org.openjdk.jmh.annotations.Benchmark;
+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.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+
+/// Compares reading a serialized primitive array list one value at a time
(`perValue`, the loop the list SerDes used
+/// before) with the ObjectSerDeUtils list SerDes (`bulk`), from a heap or a
direct buffer. The SerDes read in bulk from
+/// PrimitiveArrayOutput.MIN_BULK_VALUES values up, and one value at a time
below it.
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Fork(3)
+@Warmup(iterations = 3, time = 1)
+@Measurement(iterations = 5, time = 1)
+@State(Scope.Benchmark)
+public class BenchmarkPrimitiveListDeserialize {
+ @Param({"INT", "LONG", "DOUBLE"})
+ private String _type;
+
+ @Param({"1", "4", "16", "64", "1024", "65536", "1048576"})
+ private int _size;
+
+ @Param({"HEAP", "DIRECT"})
+ private String _buffer;
+
+ private ByteBuffer _serialized;
+
+ @Setup
+ public void setUp() {
+ Random random = new Random(42);
+ byte[] bytes;
+ switch (_type) {
+ case "INT":
+ bytes = ObjectSerDeUtils.INT_ARRAY_LIST_SER_DE.serialize(new
IntArrayList(random.ints(_size).toArray()));
+ break;
+ case "LONG":
+ bytes = ObjectSerDeUtils.LONG_ARRAY_LIST_SER_DE.serialize(new
LongArrayList(random.longs(_size).toArray()));
+ break;
+ default:
+ bytes = ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE.serialize(
+ new DoubleArrayList(random.doubles(_size).toArray()));
+ break;
+ }
+ _serialized = "HEAP".equals(_buffer) ? ByteBuffer.allocate(bytes.length) :
ByteBuffer.allocateDirect(bytes.length);
+ _serialized.put(bytes).flip();
+ }
+
+ @Benchmark
+ public Object perValue() {
+ ByteBuffer buffer = _serialized.duplicate();
+ int numValues = buffer.getInt();
+ switch (_type) {
+ case "INT": {
+ IntArrayList list = new IntArrayList(numValues);
+ for (int i = 0; i < numValues; i++) {
+ list.add(buffer.getInt());
+ }
+ return list;
+ }
+ case "LONG": {
+ LongArrayList list = new LongArrayList(numValues);
+ for (int i = 0; i < numValues; i++) {
+ list.add(buffer.getLong());
+ }
+ return list;
+ }
+ default: {
+ DoubleArrayList list = new DoubleArrayList(numValues);
+ for (int i = 0; i < numValues; i++) {
+ list.add(buffer.getDouble());
+ }
+ return list;
+ }
+ }
+ }
+
+ @Benchmark
+ public Object bulk() {
+ ByteBuffer buffer = _serialized.duplicate();
+ switch (_type) {
+ case "INT":
+ return ObjectSerDeUtils.INT_ARRAY_LIST_SER_DE.deserialize(buffer);
+ case "LONG":
+ return ObjectSerDeUtils.LONG_ARRAY_LIST_SER_DE.deserialize(buffer);
+ default:
+ return ObjectSerDeUtils.DOUBLE_ARRAY_LIST_SER_DE.deserialize(buffer);
+ }
+ }
+
+ public static void main(String[] args)
+ throws Exception {
+ new Runner(new
OptionsBuilder().include(BenchmarkPrimitiveListDeserialize.class.getSimpleName()).build()).run();
+ }
+}
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
index c13f7b09890..2b4cfa5e50f 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
@@ -32,6 +32,7 @@ import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.FunctionContext;
import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.common.utils.config.QueryOptionsUtils;
import org.apache.pinot.core.common.BlockValSet;
import org.apache.pinot.core.operator.docvalsets.DataBlockValSet;
@@ -455,11 +456,30 @@ public class AggregateOperator extends MultiStageOperator
{
return blockValSetMap;
}
- static Object[] getIntermediateResults(AggregationFunction<?, ?>
aggFunction, MseBlock.Data block) {
+ /// Returns the id of the column that holds the intermediate results of
`aggFunction` in a merge block.
+ static int getIntermediateResultColId(AggregationFunction<?, ?> aggFunction)
{
ExpressionContext firstArgument = aggFunction.getInputExpressions().get(0);
Preconditions.checkState(firstArgument.getType() ==
ExpressionContext.Type.IDENTIFIER,
"Expected the first argument to be IDENTIFIER, got: %s",
firstArgument.getType());
- int colId = fromIdentifierToColId(firstArgument.getIdentifier());
+ return fromIdentifierToColId(firstArgument.getIdentifier());
+ }
+
+ /// Returns the serialized data block of `block` when the intermediate
results of `aggFunction` are serialized
+ /// custom objects in it, or `null` otherwise. Merging such a block with
+ /// `AggregationFunctionUtils.mergeSerialized` skips deserializing each
intermediate result into a separate object.
+ @Nullable
+ static DataBlock getSerializedIntermediateResults(AggregationFunction<?, ?>
aggFunction, MseBlock.Data block) {
+ if (block.isRowHeap()) {
+ return null;
+ }
+ DataBlock dataBlock = block.asSerialized().getDataBlock();
+ ColumnDataType storedType =
+
dataBlock.getDataSchema().getColumnDataType(getIntermediateResultColId(aggFunction)).getStoredType();
+ return storedType == ColumnDataType.OBJECT ? dataBlock : null;
+ }
+
+ static Object[] getIntermediateResults(AggregationFunction<?, ?>
aggFunction, MseBlock.Data block) {
+ int colId = getIntermediateResultColId(aggFunction);
int numRows = block.getNumRows();
if (block.isRowHeap()) {
Object[] values = new Object[numRows];
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java
index f985a59a392..4f059bf573d 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java
@@ -20,6 +20,7 @@ package org.apache.pinot.query.runtime.operator;
import java.util.List;
import java.util.Map;
+import org.apache.pinot.common.datablock.DataBlock;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.common.BlockValSet;
@@ -145,6 +146,16 @@ public class MultistageAggregationExecutor {
private void processMerge(MseBlock.Data block) {
for (int i = 0; i < _aggFunctions.length; i++) {
AggregationFunction aggFunction = _aggFunctions[i];
+ DataBlock serializedResults =
AggregateOperator.getSerializedIntermediateResults(aggFunction, block);
+ if (serializedResults != null) {
+ int colId = AggregateOperator.getIntermediateResultColId(aggFunction);
+ int numRows = serializedResults.getNumberOfRows();
+ for (int rowId = 0; rowId < numRows; rowId++) {
+ _mergeResultHolder[i] =
AggregationFunctionUtils.mergeSerialized(aggFunction, _mergeResultHolder[i],
+ serializedResults.getCustomObject(rowId, colId));
+ }
+ continue;
+ }
Object[] intermediateResults =
AggregateOperator.getIntermediateResults(aggFunction, block);
for (Object intermediateResult : intermediateResults) {
_mergeResultHolder[i] =
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
index af296abf93e..d7cb6814f32 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
@@ -465,8 +465,18 @@ public class MultistageGroupByExecutor {
int numRows = groupByKeys.length;
int numFunctions = _aggFunctions.length;
Object[][] intermediateResults = new Object[numFunctions][];
+ // Serialized custom object intermediate results are merged straight from
the data block (see mergeSerialized).
+ DataBlock[] serializedResults = new DataBlock[numFunctions];
+ int[] serializedResultColIds = new int[numFunctions];
for (int i = 0; i < numFunctions; i++) {
- intermediateResults[i] =
AggregateOperator.getIntermediateResults(_aggFunctions[i], block);
+ if (!_leafReturnFinalResult) {
+ serializedResults[i] =
AggregateOperator.getSerializedIntermediateResults(_aggFunctions[i], block);
+ }
+ if (serializedResults[i] != null) {
+ serializedResultColIds[i] =
AggregateOperator.getIntermediateResultColId(_aggFunctions[i]);
+ } else {
+ intermediateResults[i] =
AggregateOperator.getIntermediateResults(_aggFunctions[i], block);
+ }
}
if (_leafReturnFinalResult) {
for (int i = 0; i < numRows; i++) {
@@ -500,8 +510,11 @@ public class MultistageGroupByExecutor {
mergedResults = _mergeResultHolder.get(groupByKey);
}
for (int j = 0; j < numFunctions; j++) {
- mergedResults[j] =
- AggregationFunctionUtils.merge(_aggFunctions[j],
mergedResults[j], intermediateResults[j][i]);
+ DataBlock serialized = serializedResults[j];
+ mergedResults[j] = serialized != null
+ ? AggregationFunctionUtils.mergeSerialized(_aggFunctions[j],
mergedResults[j],
+ serialized.getCustomObject(i, serializedResultColIds[j]))
+ : AggregationFunctionUtils.merge(_aggFunctions[j],
mergedResults[j], intermediateResults[j][i]);
}
}
}
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SerializedIntermediateResultMergeTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SerializedIntermediateResultMergeTest.java
new file mode 100644
index 00000000000..00932c93219
--- /dev/null
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SerializedIntermediateResultMergeTest.java
@@ -0,0 +1,141 @@
+/**
+ * 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.query.runtime.operator;
+
+import it.unimi.dsi.fastutil.doubles.DoubleArrayList;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
+import org.apache.pinot.core.query.aggregation.function.AvgAggregationFunction;
+import
org.apache.pinot.core.query.aggregation.function.PercentileAggregationFunction;
+import org.apache.pinot.query.planner.plannode.AggregateNode.AggType;
+import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock;
+import org.apache.pinot.segment.local.customobject.AvgPair;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+/// Checks that merging serialized custom object intermediate results straight
from a data block gives the same result
+/// as merging the deserialized values of a row heap block.
+public class SerializedIntermediateResultMergeTest {
+ private static final DataSchema INPUT_SCHEMA = new DataSchema(new
String[]{"key", "percentile", "avg"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.OBJECT,
ColumnDataType.OBJECT});
+ private static final DataSchema GROUP_BY_RESULT_SCHEMA = new DataSchema(new
String[]{"key", "percentile", "avg"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.DOUBLE,
ColumnDataType.DOUBLE});
+ private static final DataSchema AGGREGATION_RESULT_SCHEMA = new
DataSchema(new String[]{"percentile", "avg"},
+ new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.DOUBLE});
+
+ @Test
+ public void testGroupByMergeOfSerializedBlocksMatchesRowHeapBlocks() {
+ MultistageGroupByExecutor serializedExecutor = newGroupByExecutor();
+ MultistageGroupByExecutor rowHeapExecutor = newGroupByExecutor();
+ for (List<Object[]> rows : inputBlocks()) {
+ serializedExecutor.processBlock(newBlock(rows).asSerialized());
+ rowHeapExecutor.processBlock(newBlock(copyStates(rows)));
+ }
+
+ List<Object[]> expected = rowHeapExecutor.getResult(100);
+ List<Object[]> actual = serializedExecutor.getResult(100);
+ assertEquals(actual.size(), 3);
+ assertEquals(toMap(actual), toMap(expected));
+ }
+
+ @Test
+ public void testAggregationMergeOfSerializedBlocksMatchesRowHeapBlocks() {
+ MultistageAggregationExecutor serializedExecutor =
newAggregationExecutor();
+ MultistageAggregationExecutor rowHeapExecutor = newAggregationExecutor();
+ for (List<Object[]> rows : inputBlocks()) {
+ serializedExecutor.processBlock(newBlock(rows).asSerialized());
+ rowHeapExecutor.processBlock(newBlock(copyStates(rows)));
+ }
+
+ Object[] expected = rowHeapExecutor.getResult().get(0);
+ Object[] actual = serializedExecutor.getResult().get(0);
+ assertTrue(actual[0] instanceof Double && actual[1] instanceof Double);
+ assertEquals(actual, expected);
+ }
+
+ /// Three blocks. Group 1 and 2 span blocks, a null state is skipped, and
group 3 only has an empty list.
+ private static List<List<Object[]>> inputBlocks() {
+ return List.of(
+ List.of(
+ new Object[]{1, values(5, 1, 9), new AvgPair(15, 3)},
+ new Object[]{2, values(100), new AvgPair(100, 1)},
+ new Object[]{1, null, null}),
+ List.of(),
+ List.of(
+ new Object[]{2, values(-3, 7, 7, 42), new AvgPair(53, 4)},
+ new Object[]{1, values(2), new AvgPair(2, 1)},
+ new Object[]{3, values(), new AvgPair(0, 0)}));
+ }
+
+ private static DoubleArrayList values(double... values) {
+ return new DoubleArrayList(values);
+ }
+
+ /// The executors merge in place, so each executor gets its own copy of the
mutable states.
+ private static List<Object[]> copyStates(List<Object[]> rows) {
+ List<Object[]> copies = new ArrayList<>(rows.size());
+ for (Object[] row : rows) {
+ Object[] copy = row.clone();
+ if (copy[1] != null) {
+ copy[1] = new DoubleArrayList((DoubleArrayList) copy[1]);
+ }
+ copies.add(copy);
+ }
+ return copies;
+ }
+
+ private static RowHeapDataBlock newBlock(List<Object[]> rows) {
+ return new RowHeapDataBlock(rows, INPUT_SCHEMA, newFunctions());
+ }
+
+ private static AggregationFunction[] newFunctions() {
+ return new AggregationFunction[]{
+ new
PercentileAggregationFunction(ExpressionContext.forIdentifier("$1"), 50, true),
+ new
AvgAggregationFunction(List.of(ExpressionContext.forIdentifier("$2")), true)
+ };
+ }
+
+ private static MultistageGroupByExecutor newGroupByExecutor() {
+ return new MultistageGroupByExecutor(new int[]{0}, newFunctions(), new
int[]{-1, -1}, -1, AggType.FINAL, false,
+ GROUP_BY_RESULT_SCHEMA, Map.of(), null);
+ }
+
+ private static MultistageAggregationExecutor newAggregationExecutor() {
+ return new MultistageAggregationExecutor(newFunctions(), new int[]{-1,
-1}, -1, AggType.FINAL,
+ AGGREGATION_RESULT_SCHEMA);
+ }
+
+ private static Map<Object, List<Object>> toMap(List<Object[]> rows) {
+ Map<Object, List<Object>> result = new HashMap<>();
+ for (Object[] row : rows) {
+ result.put(row[0], Arrays.asList(row).subList(1, row.length));
+ }
+ return result;
+ }
+}
diff --git
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PagedPinotOutputStream.java
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PagedPinotOutputStream.java
index 8a6b3719487..cb610cc063e 100644
---
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PagedPinotOutputStream.java
+++
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PagedPinotOutputStream.java
@@ -23,6 +23,7 @@ import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.ArrayList;
+import java.util.Objects;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -212,6 +213,99 @@ public class PagedPinotOutputStream extends
PinotOutputStream {
_written = Math.max(_written, _offsetInPage + _currentPageStartOffset);
}
+ /// Puts as many values as fit in the current page straight into it, then
updates the page offset once. When at least
+ /// [PrimitiveArrayOutput#MIN_BULK_VALUES] values go to the page, they are
copied through a big-endian typed view;
+ /// fewer values are put one by one, because the typed-view copy has a fixed
cost that a few values do not pay back.
+ /// A value that does not fit in the rest of the page (the page offset
depends on everything written before, so it
+ /// is not aligned to the value width) is written with [#writeInt(int)],
which splits it across the two pages.
+ @Override
+ public void writeInts(int[] values, int offset, int length)
+ throws IOException {
+ Objects.checkFromIndexSize(offset, length, values.length);
+ int end = offset + length;
+ int i = offset;
+ while (i < end) {
+ int numValues = Math.min(end - i, remainingInPage() / Integer.BYTES);
+ if (numValues == 0) {
+ writeInt(values[i++]);
+ continue;
+ }
+ if (numValues < PrimitiveArrayOutput.MIN_BULK_VALUES) {
+ int position = _offsetInPage;
+ for (int j = i; j < i + numValues; j++, position += Integer.BYTES) {
+ _currentPage.putInt(position, values[j]);
+ }
+ } else {
+ currentPageView().asIntBuffer().put(values, i, numValues);
+ }
+ i += numValues;
+ _offsetInPage += numValues * Integer.BYTES;
+ _written = Math.max(_written, _offsetInPage + _currentPageStartOffset);
+ }
+ }
+
+ /// Same as [#writeInts(int[], int, int)], for longs.
+ @Override
+ public void writeLongs(long[] values, int offset, int length)
+ throws IOException {
+ Objects.checkFromIndexSize(offset, length, values.length);
+ int end = offset + length;
+ int i = offset;
+ while (i < end) {
+ int numValues = Math.min(end - i, remainingInPage() / Long.BYTES);
+ if (numValues == 0) {
+ writeLong(values[i++]);
+ continue;
+ }
+ if (numValues < PrimitiveArrayOutput.MIN_BULK_VALUES) {
+ int position = _offsetInPage;
+ for (int j = i; j < i + numValues; j++, position += Long.BYTES) {
+ _currentPage.putLong(position, values[j]);
+ }
+ } else {
+ currentPageView().asLongBuffer().put(values, i, numValues);
+ }
+ i += numValues;
+ _offsetInPage += numValues * Long.BYTES;
+ _written = Math.max(_written, _offsetInPage + _currentPageStartOffset);
+ }
+ }
+
+ /// Same as [#writeInts(int[], int, int)], for doubles.
+ @Override
+ public void writeDoubles(double[] values, int offset, int length)
+ throws IOException {
+ Objects.checkFromIndexSize(offset, length, values.length);
+ int end = offset + length;
+ int i = offset;
+ while (i < end) {
+ int numValues = Math.min(end - i, remainingInPage() / Double.BYTES);
+ if (numValues == 0) {
+ // Raw bits, like the typed view (and ByteBuffer#putDouble):
writeDouble() would canonicalize NaN.
+ writeLong(Double.doubleToRawLongBits(values[i++]));
+ continue;
+ }
+ if (numValues < PrimitiveArrayOutput.MIN_BULK_VALUES) {
+ int position = _offsetInPage;
+ for (int j = i; j < i + numValues; j++, position += Double.BYTES) {
+ _currentPage.putLong(position,
Double.doubleToRawLongBits(values[j]));
+ }
+ } else {
+ currentPageView().asDoubleBuffer().put(values, i, numValues);
+ }
+ i += numValues;
+ _offsetInPage += numValues * Double.BYTES;
+ _written = Math.max(_written, _offsetInPage + _currentPageStartOffset);
+ }
+ }
+
+ /// Returns a big-endian view of the current page that starts at the current
offset. The page itself is not changed.
+ private ByteBuffer currentPageView() {
+ ByteBuffer view = _currentPage.duplicate();
+ view.position(_offsetInPage);
+ return view.order(ByteOrder.BIG_ENDIAN);
+ }
+
@Override
public void write(DataBuffer input, long offset, long length)
throws IOException {
diff --git
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PinotOutputStream.java
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PinotOutputStream.java
index 8fd496fd9fe..b0f5b8795c4 100644
---
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PinotOutputStream.java
+++
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PinotOutputStream.java
@@ -89,6 +89,28 @@ public abstract class PinotOutputStream extends OutputStream
implements DataOutp
write(bytes, 0, bytes.length);
}
+ /// Writes `length` ints of `values` from `offset`, with the same big-endian
bytes as calling [#writeInt(int)] on
+ /// each one. The default encodes the values in small chunks; subclasses can
write them in bulk.
+ public void writeInts(int[] values, int offset, int length)
+ throws IOException {
+ PrimitiveArrayOutput.writeIntsChunked(this, values, offset, length);
+ }
+
+ /// Writes `length` longs of `values` from `offset`, with the same
big-endian bytes as calling [#writeLong(long)] on
+ /// each one. The default encodes the values in small chunks; subclasses can
write them in bulk.
+ public void writeLongs(long[] values, int offset, int length)
+ throws IOException {
+ PrimitiveArrayOutput.writeLongsChunked(this, values, offset, length);
+ }
+
+ /// Writes `length` doubles of `values` from `offset` as big-endian raw bits
([Double#doubleToRawLongBits]), the same
+ /// bytes as `ByteBuffer.putDouble`. Unlike [#writeDouble(double)], a NaN
keeps its payload. The default
+ /// encodes the values in small chunks; subclasses can write them in bulk.
+ public void writeDoubles(double[] values, int offset, int length)
+ throws IOException {
+ PrimitiveArrayOutput.writeDoublesChunked(this, values, offset, length);
+ }
+
public void writeShort(int v) throws IOException {
write(0xFF & (v >> 8));
write(0xFF & v);
diff --git
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PrimitiveArrayOutput.java
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PrimitiveArrayOutput.java
new file mode 100644
index 00000000000..da7d372a96f
--- /dev/null
+++
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PrimitiveArrayOutput.java
@@ -0,0 +1,143 @@
+/**
+ * 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.spi.memory;
+
+import java.io.DataOutput;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.Objects;
+
+
+/// Writes ranges of primitive arrays to a [DataOutput], with the same
big-endian bytes as one [ByteBuffer#putInt(int)],
+/// [ByteBuffer#putLong(long)] or [ByteBuffer#putDouble(double)] call per
value. For ints and longs that is also what
+/// [DataOutput#writeInt(int)] and [DataOutput#writeLong(long)] write. Doubles
keep their raw bits, while
+/// [DataOutput#writeDouble(double)] canonicalizes NaN.
+///
+/// A [PinotOutputStream] writes the values in bulk (see
[PinotOutputStream#writeDoubles]). Any other [DataOutput] gets
+/// the values encoded in chunks of at most [#CHUNK_SIZE] bytes, so a large
array never needs a byte array of its full
+/// size. The class is stateless and thread-safe.
+public final class PrimitiveArrayOutput {
+ /// Size in bytes of the scratch buffer used to encode a chunk. Small enough
to be allocated in a TLAB.
+ public static final int CHUNK_SIZE = 8 * 1024;
+
+ /// Below this number of values, reads and writes through a big-endian typed
view (`asIntBuffer().get(int[])` and
+ /// the like) are slower than one access per value. On a little-endian CPU
the bulk copy swaps the bytes in a native
+ /// call with a fixed cost of about 10-40 ns, while a single-value access
compiles to an inlined byte swap. Measured
+ /// with BenchmarkPrimitiveListDeserialize and BenchmarkPrimitiveArrayOutput
on aarch64 and x86_64.
+ public static final int MIN_BULK_VALUES = 64;
+
+ /// Below this number of values, the chunked writes call the [DataOutput]
once per value: encoding a chunk costs a
+ /// scratch buffer and a typed-view copy, which is more than a few small
writes.
+ static final int MIN_CHUNKED_VALUES = 16;
+
+ private PrimitiveArrayOutput() {
+ }
+
+ public static void writeInts(DataOutput output, int[] values, int offset,
int length)
+ throws IOException {
+ if (output instanceof PinotOutputStream) {
+ ((PinotOutputStream) output).writeInts(values, offset, length);
+ } else {
+ writeIntsChunked(output, values, offset, length);
+ }
+ }
+
+ public static void writeLongs(DataOutput output, long[] values, int offset,
int length)
+ throws IOException {
+ if (output instanceof PinotOutputStream) {
+ ((PinotOutputStream) output).writeLongs(values, offset, length);
+ } else {
+ writeLongsChunked(output, values, offset, length);
+ }
+ }
+
+ public static void writeDoubles(DataOutput output, double[] values, int
offset, int length)
+ throws IOException {
+ if (output instanceof PinotOutputStream) {
+ ((PinotOutputStream) output).writeDoubles(values, offset, length);
+ } else {
+ writeDoublesChunked(output, values, offset, length);
+ }
+ }
+
+ static void writeIntsChunked(DataOutput output, int[] values, int offset,
int length)
+ throws IOException {
+ Objects.checkFromIndexSize(offset, length, values.length);
+ if (length < MIN_CHUNKED_VALUES) {
+ for (int i = offset; i < offset + length; i++) {
+ output.writeInt(values[i]);
+ }
+ return;
+ }
+ ByteBuffer scratch = newScratch(length, Integer.BYTES);
+ int chunkValues = scratch.capacity() / Integer.BYTES;
+ int end = offset + length;
+ for (int i = offset; i < end; i += chunkValues) {
+ int numValues = Math.min(chunkValues, end - i);
+ scratch.asIntBuffer().put(values, i, numValues);
+ output.write(scratch.array(), 0, numValues * Integer.BYTES);
+ }
+ }
+
+ static void writeLongsChunked(DataOutput output, long[] values, int offset,
int length)
+ throws IOException {
+ Objects.checkFromIndexSize(offset, length, values.length);
+ if (length < MIN_CHUNKED_VALUES) {
+ for (int i = offset; i < offset + length; i++) {
+ output.writeLong(values[i]);
+ }
+ return;
+ }
+ ByteBuffer scratch = newScratch(length, Long.BYTES);
+ int chunkValues = scratch.capacity() / Long.BYTES;
+ int end = offset + length;
+ for (int i = offset; i < end; i += chunkValues) {
+ int numValues = Math.min(chunkValues, end - i);
+ scratch.asLongBuffer().put(values, i, numValues);
+ output.write(scratch.array(), 0, numValues * Long.BYTES);
+ }
+ }
+
+ static void writeDoublesChunked(DataOutput output, double[] values, int
offset, int length)
+ throws IOException {
+ Objects.checkFromIndexSize(offset, length, values.length);
+ if (length < MIN_CHUNKED_VALUES) {
+ for (int i = offset; i < offset + length; i++) {
+ // Raw bits, like the chunked path: writeDouble() would canonicalize
NaN.
+ output.writeLong(Double.doubleToRawLongBits(values[i]));
+ }
+ return;
+ }
+ ByteBuffer scratch = newScratch(length, Double.BYTES);
+ int chunkValues = scratch.capacity() / Double.BYTES;
+ int end = offset + length;
+ for (int i = offset; i < end; i += chunkValues) {
+ int numValues = Math.min(chunkValues, end - i);
+ scratch.asDoubleBuffer().put(values, i, numValues);
+ output.write(scratch.array(), 0, numValues * Double.BYTES);
+ }
+ }
+
+ /// Returns a big-endian heap buffer for one chunk: [#CHUNK_SIZE] bytes, or
less when all the values fit.
+ private static ByteBuffer newScratch(int length, int width) {
+ int capacity = (int) Math.min(CHUNK_SIZE, (long) length * width);
+ return ByteBuffer.allocate(capacity).order(ByteOrder.BIG_ENDIAN);
+ }
+}
diff --git
a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/memory/PrimitiveArrayOutputTest.java
b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/memory/PrimitiveArrayOutputTest.java
new file mode 100644
index 00000000000..dfe6702181e
--- /dev/null
+++
b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/memory/PrimitiveArrayOutputTest.java
@@ -0,0 +1,282 @@
+/**
+ * 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.spi.memory;
+
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutput;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
+
+
+/// Checks that the bulk primitive array writes of [PinotOutputStream] and
[PrimitiveArrayOutput] write the same bytes
+/// as one [ByteBuffer] put per value, in pages of any size, from any offset
in a page.
+public class PrimitiveArrayOutputTest {
+ private static final int NUM_VALUES = 5000;
+
+ /// One primitive type: how to write a range of an array in bulk, and the
per-value oracle.
+ private enum Type {
+ INT(Integer.BYTES) {
+ @Override
+ Object newArray(Random random) {
+ int[] values = random.ints(NUM_VALUES).toArray();
+ values[0] = Integer.MIN_VALUE;
+ values[1] = -1;
+ return values;
+ }
+
+ @Override
+ void put(ByteBuffer buffer, Object values, int index) {
+ buffer.putInt(((int[]) values)[index]);
+ }
+
+ @Override
+ void writeBulk(PinotOutputStream output, Object values, int offset, int
length)
+ throws IOException {
+ output.writeInts((int[]) values, offset, length);
+ }
+
+ @Override
+ void writeChunked(DataOutput output, Object values, int offset, int
length)
+ throws IOException {
+ PrimitiveArrayOutput.writeIntsChunked(output, (int[]) values, offset,
length);
+ }
+
+ @Override
+ void writeDispatched(DataOutput output, Object values, int offset, int
length)
+ throws IOException {
+ PrimitiveArrayOutput.writeInts(output, (int[]) values, offset, length);
+ }
+ },
+ LONG(Long.BYTES) {
+ @Override
+ Object newArray(Random random) {
+ long[] values = random.longs(NUM_VALUES).toArray();
+ values[0] = Long.MIN_VALUE;
+ values[1] = -1;
+ return values;
+ }
+
+ @Override
+ void put(ByteBuffer buffer, Object values, int index) {
+ buffer.putLong(((long[]) values)[index]);
+ }
+
+ @Override
+ void writeBulk(PinotOutputStream output, Object values, int offset, int
length)
+ throws IOException {
+ output.writeLongs((long[]) values, offset, length);
+ }
+
+ @Override
+ void writeChunked(DataOutput output, Object values, int offset, int
length)
+ throws IOException {
+ PrimitiveArrayOutput.writeLongsChunked(output, (long[]) values,
offset, length);
+ }
+
+ @Override
+ void writeDispatched(DataOutput output, Object values, int offset, int
length)
+ throws IOException {
+ PrimitiveArrayOutput.writeLongs(output, (long[]) values, offset,
length);
+ }
+ },
+ DOUBLE(Double.BYTES) {
+ @Override
+ Object newArray(Random random) {
+ double[] values = random.doubles(NUM_VALUES).map(d -> (d - 0.5) *
1e6).toArray();
+ values[0] = Double.NaN;
+ // A NaN with a payload keeps its raw bits, like ByteBuffer.putDouble.
+ values[1] = Double.longBitsToDouble(0x7ff8_0000_0000_0123L);
+ values[2] = -0.0;
+ values[3] = Double.NEGATIVE_INFINITY;
+ return values;
+ }
+
+ @Override
+ void put(ByteBuffer buffer, Object values, int index) {
+ buffer.putDouble(((double[]) values)[index]);
+ }
+
+ @Override
+ void writeBulk(PinotOutputStream output, Object values, int offset, int
length)
+ throws IOException {
+ output.writeDoubles((double[]) values, offset, length);
+ }
+
+ @Override
+ void writeChunked(DataOutput output, Object values, int offset, int
length)
+ throws IOException {
+ PrimitiveArrayOutput.writeDoublesChunked(output, (double[]) values,
offset, length);
+ }
+
+ @Override
+ void writeDispatched(DataOutput output, Object values, int offset, int
length)
+ throws IOException {
+ PrimitiveArrayOutput.writeDoubles(output, (double[]) values, offset,
length);
+ }
+ };
+
+ final int _width;
+
+ Type(int width) {
+ _width = width;
+ }
+
+ abstract Object newArray(Random random);
+
+ abstract void put(ByteBuffer buffer, Object values, int index);
+
+ abstract void writeBulk(PinotOutputStream output, Object values, int
offset, int length)
+ throws IOException;
+
+ abstract void writeChunked(DataOutput output, Object values, int offset,
int length)
+ throws IOException;
+
+ abstract void writeDispatched(DataOutput output, Object values, int
offset, int length)
+ throws IOException;
+
+ byte[] expected(int prefix, Object values, int offset, int length) {
+ ByteBuffer buffer = ByteBuffer.allocate(prefix + length * _width +
Integer.BYTES).order(ByteOrder.BIG_ENDIAN);
+ for (int i = 0; i < prefix; i++) {
+ buffer.put((byte) i);
+ }
+ for (int i = offset; i < offset + length; i++) {
+ put(buffer, values, i);
+ }
+ buffer.putInt(0xCAFEBABE);
+ return buffer.array();
+ }
+ }
+
+ @DataProvider
+ public Object[][] cases() {
+ // Page sizes: one value per page or less (3, 7), not a multiple of 8 (13,
100), the default.
+ int[] pageSizes = {3, 7, 13, 100,
PagedPinotOutputStream.PageAllocator.MIN_RECOMMENDED_PAGE_SIZE};
+ // Bytes written before the values, so the values start at different
offsets in a page.
+ int[] prefixes = {0, 1, 5};
+ // Lengths on both sides of MIN_CHUNKED_VALUES (16) and MIN_BULK_VALUES
(64), where the per-value path takes over.
+ int[][] ranges = {
+ {0, 0}, {7, 0}, {0, NUM_VALUES}, {3, NUM_VALUES - 10}, {NUM_VALUES -
1, 1}, {2, 15}, {2, 16}, {9, 63}, {9, 64}
+ };
+ List<Object[]> cases = new ArrayList<>();
+ for (Type type : Type.values()) {
+ for (int pageSize : pageSizes) {
+ for (int prefix : prefixes) {
+ for (int[] range : ranges) {
+ cases.add(new Object[]{type, pageSize, prefix, range[0],
range[1]});
+ }
+ }
+ }
+ }
+ return cases.toArray(new Object[0][]);
+ }
+
+ @Test(dataProvider = "cases")
+ public void testPagedBulkWriteMatchesPerValuePuts(Type type, int pageSize,
int prefix, int offset, int length)
+ throws IOException {
+ Object values = type.newArray(new Random(pageSize * 31L + prefix));
+ byte[] expected = type.expected(prefix, values, offset, length);
+ try (PagedPinotOutputStream output = newStream(pageSize)) {
+ writePrefix(output, prefix);
+ type.writeBulk(output, values, offset, length);
+ output.writeInt(0xCAFEBABE);
+ assertEquals(output.getCurrentOffset(), expected.length);
+ assertEquals(toBytes(output), expected);
+ }
+ }
+
+ @Test(dataProvider = "cases")
+ public void testChunkedWriteMatchesPerValuePuts(Type type, int pageSize, int
prefix, int offset, int length)
+ throws IOException {
+ Object values = type.newArray(new Random(pageSize * 31L + prefix));
+ byte[] expected = type.expected(prefix, values, offset, length);
+ // The default implementation, into pages.
+ try (PagedPinotOutputStream output = newStream(pageSize)) {
+ writePrefix(output, prefix);
+ type.writeChunked(output, values, offset, length);
+ output.writeInt(0xCAFEBABE);
+ assertEquals(toBytes(output), expected);
+ }
+ // A plain DataOutput, as used by data tables.
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ DataOutputStream output = new DataOutputStream(bytes);
+ writePrefix(output, prefix);
+ type.writeDispatched(output, values, offset, length);
+ output.writeInt(0xCAFEBABE);
+ assertEquals(bytes.toByteArray(), expected);
+ }
+
+ /// Writing over earlier bytes after a seek back does not move the end of
the written data.
+ @Test
+ public void testBulkWriteAfterSeekBackKeepsWrittenSize()
+ throws IOException {
+ try (PagedPinotOutputStream output = newStream(13)) {
+ output.writeLongs(new long[]{1, 2, 3, 4}, 0, 4);
+ output.seek(4);
+ output.writeInts(new int[]{-1}, 0, 1);
+ assertEquals(output.getCurrentOffset(), 8);
+ ByteBuffer expected =
ByteBuffer.allocate(32).putLong(1).putLong(2).putLong(3).putLong(4);
+ expected.putInt(4, -1);
+ assertEquals(toBytes(output), expected.array());
+ }
+ }
+
+ @Test
+ public void testInvalidRangeThrows()
+ throws IOException {
+ try (PagedPinotOutputStream output = newStream(16)) {
+ assertThrows(IndexOutOfBoundsException.class, () -> output.writeInts(new
int[2], 1, 2));
+ assertThrows(IndexOutOfBoundsException.class, () ->
output.writeLongs(new long[2], -1, 1));
+ assertThrows(IndexOutOfBoundsException.class,
+ () -> PrimitiveArrayOutput.writeDoubles(new DataOutputStream(new
ByteArrayOutputStream()), new double[2], 0,
+ 3));
+ assertEquals(output.getCurrentOffset(), 0);
+ }
+ }
+
+ private static PagedPinotOutputStream newStream(int pageSize) {
+ return new PagedPinotOutputStream(new
PagedPinotOutputStream.HeapPageAllocator(pageSize));
+ }
+
+ private static void writePrefix(DataOutput output, int prefix)
+ throws IOException {
+ for (int i = 0; i < prefix; i++) {
+ output.write(i);
+ }
+ }
+
+ private static byte[] toBytes(PagedPinotOutputStream output) {
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ for (ByteBuffer page : output.getPages()) {
+ byte[] pageBytes = new byte[page.remaining()];
+ page.get(pageBytes);
+ bytes.writeBytes(pageBytes);
+ }
+ return bytes.toByteArray();
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
