xiangfu0 commented on code in PR #19017:
URL: https://github.com/apache/pinot/pull/19017#discussion_r3619719638
##########
pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAccumulator.java:
##########
@@ -476,19 +477,36 @@ public double compression() {
return _compression;
}
+ /// Returns the length of the [#serialize()] bytes rather than the
materialized [MergingDigest] byte size, so
+ /// generic `TDigest` serializers ([#byteSize()] followed by
[#asBytes(ByteBuffer)]) emit the capacity-preserving
+ /// encoding. Re-encoding through a materialized [MergingDigest] can produce
a verbose digest with more centroids
+ /// than a freshly allocated [MergingDigest] of the same compression can
hold, which readers reject with
+ /// [ArrayIndexOutOfBoundsException]. The length is computed without
materializing the bytes by mirroring the
+ /// [#serialize()] branches; the mutations here (flush, capacity
normalization) are idempotent, so the following
+ /// [#asBytes(ByteBuffer)] call writes exactly this many bytes.
@Override
public int byteSize() {
- return toTDigest().byteSize();
+ if (_pendingSerializedTDigest != null) {
+ return _pendingSerializedTDigest.length;
+ }
+ flush();
+ if (requiresCapacityPreservingEncoding()) {
+ return SMALL_HEADER_SIZE + SMALL_CENTROID_SIZE * _numCentroids;
+ }
+ normalizeCentroidCapacity();
+ return VERBOSE_HEADER_SIZE + VERBOSE_CENTROID_SIZE * _numCentroids;
}
@Override
public int smallByteSize() {
return toTDigest().smallByteSize();
}
+ /// Writes the [#serialize()] bytes; see [#byteSize()]. Bytes are always
written in big-endian order (the t-digest
+ /// wire order) regardless of the destination buffer's byte order, unlike
[MergingDigest#asBytes(ByteBuffer)].
@Override
public void asBytes(ByteBuffer buffer) {
- toTDigest().asBytes(buffer);
+ buffer.put(serialize());
Review Comment:
Done — `asBytes()` now writes the pending serialized bytes straight into the
destination buffer (keeping the same validation `serialize()` performed),
skipping the defensive clone.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileSmartTDigestAggregationFunction.java:
##########
@@ -267,28 +267,35 @@ public Object extractGroupByResult(GroupByResultHolder
groupByResultHolder, int
@Override
public Object merge(Object intermediateResult1, Object intermediateResult2) {
if (intermediateResult1 instanceof TDigest) {
- return mergeIntoTDigest((TDigest) intermediateResult1,
intermediateResult2);
+ return mergeIntoAccumulator((TDigest) intermediateResult1,
intermediateResult2);
}
if (intermediateResult2 instanceof TDigest) {
- return mergeIntoTDigest((TDigest) intermediateResult2,
intermediateResult1);
+ return mergeIntoAccumulator((TDigest) intermediateResult2,
intermediateResult1);
}
DoubleArrayList valueList1 = (DoubleArrayList) intermediateResult1;
DoubleArrayList valueList2 = (DoubleArrayList) intermediateResult2;
valueList1.addAll(valueList2);
return valueList1.size() > _threshold ?
convertValueListToTDigest(valueList1) : valueList1;
}
- private static TDigest mergeIntoTDigest(TDigest tDigest, Object
intermediateResult) {
+ private static TDigest mergeIntoAccumulator(TDigest tDigest, Object
intermediateResult) {
Review Comment:
In practice yes: both producers of TDigest intermediates in this function
(`convertValueListToTDigest` and `deserializeIntermediateResult`) create
accumulators, so the wrap branch is a safety net for plain digests only
(mirroring the same pattern in `PercentileTDigestAggregationFunction.merge`).
Added a comment stating the invariant.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileSmartTDigestAggregationFunction.java:
##########
@@ -267,28 +267,35 @@ public Object extractGroupByResult(GroupByResultHolder
groupByResultHolder, int
@Override
public Object merge(Object intermediateResult1, Object intermediateResult2) {
Review Comment:
Good catch — it could NPE on a null intermediate. Added null short-circuits
at the top of `merge` (return the other argument) plus a test.
##########
pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileSmartTDigestAggregationFunctionTest.java:
##########
@@ -18,8 +18,113 @@
*/
package org.apache.pinot.core.query.aggregation.function;
+import com.tdunning.math.stats.MergingDigest;
+import com.tdunning.math.stats.TDigest;
+import it.unimi.dsi.fastutil.doubles.DoubleArrayList;
+import java.nio.ByteBuffer;
+import java.util.List;
+import org.apache.pinot.common.CustomObject;
+import org.apache.pinot.common.request.Literal;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.core.common.ObjectSerDeUtils;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
public class PercentileSmartTDigestAggregationFunctionTest {
+ private static final ExpressionContext EXPRESSION =
ExpressionContext.forIdentifier("col");
+
+ /// The intermediate result must round-trip through the
[PercentileTDigestAccumulator] so that
+ /// capacity-preserving digests (more centroids than a fresh [MergingDigest]
of the same
+ /// compression can hold) are not re-encoded into bytes a receiving server
cannot read.
+ @Test
+ public void
testIntermediateResultRoundTripPreservesCapacityPreservingState() {
+ int numCentroids = 51;
+ double compression = 20.0;
+ byte[] small =
+
PercentileRawTDigestAggregationFunctionTest.createSmallUnitCentroidDigest(numCentroids,
compression, 60, 100);
+ PercentileSmartTDigestAggregationFunction function = new
PercentileSmartTDigestAggregationFunction(
+ List.of(EXPRESSION,
ExpressionContext.forLiteral(Literal.doubleValue(50.0)),
+
ExpressionContext.forLiteral(Literal.stringValue("THRESHOLD=1;COMPRESSION=20"))),
false);
+
+ Object intermediateResult = function.deserializeIntermediateResult(
+ new CustomObject(ObjectSerDeUtils.ObjectType.TDigest.getValue(),
ByteBuffer.wrap(small)));
+ Assert.assertTrue(intermediateResult instanceof
PercentileTDigestAccumulator);
Review Comment:
Done — switched both new test classes to static `Assert` imports.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
