xiangfu0 commented on code in PR #19795: URL: https://github.com/apache/pinot/pull/19795#discussion_r4242893258
########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/SerializedTDigestInput.java: ########## @@ -0,0 +1,169 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.segment.local.customobject.tdigest; + +import java.nio.ByteBuffer; +import org.apache.pinot.segment.local.customobject.tdigest.TDigestCodec.SerializedTDigestMetadata; +import org.apache.pinot.segment.local.customobject.tdigest.TDigestCodec.SerializedTDigestState; + +/// Reusable, invocation-local view of one serialized TDigest for group-by-MV fanout. +/// +/// [#reset(byte\[\])] validates the header once per input row. The first decode combines numerical validation +/// with filling reusable centroid arrays, so all groups for a row merge the same primitive input without sharing +/// mutable accumulator state. A pending accumulator requests an immutable byte snapshot, created at most once per +/// reset +/// and shared by the row's group fanout. The input must remain thread-confined and must not outlive the aggregation +/// call that owns it. Pending-only inputs inspect their bytes without allocating centroid arrays; their first +/// later materialization reads the already validated values without another numerical validation pass. +public final class SerializedTDigestInput { + byte[] _bytes; + byte[] _retainedBytes; + int _numCentroids; + double[] _means; + double[] _weights; + private boolean _decoded; + private boolean _metadataInspected; + + final SerializedTDigestState _metadata = new SerializedTDigestState(); + private SerializedTDigestMetadata _metadataSnapshot; + + public void reset(byte[] bytes) { + _metadata.reset(bytes); + TDigestCodec.readSerializedHeader(_metadata, true); + resetState(bytes); + _metadataSnapshot = null; + } + + /// Owns one snapshot of the encoded digest and advances the source buffer past it. + public void reset(ByteBuffer bytes) { + reset(bytes, true); + } + + /// Reads a source digest, optionally enforcing the historical decoder's declared centroid capacity. + public void reset(ByteBuffer bytes, boolean checkCapacity) { + SerializedTDigestMetadata metadata = TDigestCodec.readSerializedHeader(bytes, checkCapacity); + byte[] snapshot = new byte[metadata.encodedLength()]; + bytes.get(snapshot); + reset(snapshot, metadata); + _retainedBytes = snapshot; + } + + void reset(byte[] bytes, SerializedTDigestMetadata metadata) { + _metadata.reset(bytes); + _metadata.copyFrom(metadata); + resetState(bytes); + _metadataSnapshot = _metadataInspected ? metadata : null; + } + + private void resetState(byte[] bytes) { + _bytes = bytes; + _retainedBytes = null; + _numCentroids = _metadata._centroidCount; + _decoded = false; + _metadataInspected = !Double.isNaN(_metadata._totalWeight); + } + + public double getCompression() { + return _metadata._compression; + } + + /// Returns validated metadata without allocating centroid arrays for a digest that will remain pending. + public SerializedTDigestMetadata getMetadata() { + inspectMetadata(); + if (_metadataSnapshot == null) { + _metadataSnapshot = _metadata.snapshot(); + } + return _metadataSnapshot; + } + + void inspectMetadata() { + if (!_metadataInspected) { + TDigestCodec.inspectSerialized(_metadata, null, null); + _metadataInspected = true; + } + } + + /// Propagates endpoint provenance cached during the one numerical inspection of this serialized distribution. + /// The encoded endpoint tuple remains valid when a previous fanout decode has split boundary centroids. + boolean recordHistoricalFractionalBoundaries(PercentileTDigestAccumulator target) { + inspectMetadata(); + if (!_metadata._fractionalWeights || _metadata._needsLegacyFallback) { + return false; + } + target.requireMutable(); + double minMean = _metadata._historicalFractionalMinMean; + double maxMean = _metadata._historicalFractionalMaxMean; + target.inheritHistoricalFractionalBoundaries(minMean, maxMean); + return !Double.isNaN(minMean) || !Double.isNaN(maxMean); + } + + byte[] retainBytes() { + if (_retainedBytes == null) { + _retainedBytes = _bytes.clone(); + } + return _retainedBytes; + } + + void decode() { + if (_decoded) { + return; + } + int encodedCentroidCount = _numCentroids; + ensureCapacity(Math.addExact(encodedCentroidCount, 2)); + if (!_metadataInspected) { + TDigestCodec.inspectSerialized(_metadata, _means, _weights); + _metadataInspected = true; + } else { + // Pending state was already inspected before retaining bytes. Its later first read only materializes + // those validated centroids; regular merged inputs compute flags and weights in the decode pass above. + TDigestCodec.decodeSerializedCentroids(_metadata, _means, _weights); + } + if (!_metadata._needsLegacyFallback) { + if (_metadata._hasZeroWeightCentroids) { + int nonZeroCount = 0; + for (int i = 0; i < encodedCentroidCount; i++) { + if (_weights[i] != 0.0) { + _means[nonZeroCount] = _means[i]; + _weights[nonZeroCount++] = _weights[i]; + } + } + encodedCentroidCount = nonZeroCount; + } + if (_metadata._unorderedMeans) { + PercentileTDigestAccumulator.sortCentroids(_means, _weights, encodedCentroidCount); + } + _numCentroids = encodedCentroidCount; + } + _decoded = true; + } + + /// Repairs endpoints only when this input is being combined with another distribution. + void normalizeBoundaries() { + _numCentroids = PercentileTDigestAccumulator.normalizeSerializedBoundaries(_means, _weights, _numCentroids, Review Comment: **Reproduced crash (group-by MV fanout):** repeated `normalizeBoundaries()` on one decoded input fails after the first call splits a weighted boundary. `decode()` sizes `_means/_weights` to exactly `encodedCount + 2` (`ensureCapacity(Math.addExact(encodedCentroidCount, 2))`). The first target's normalization grows `_numCentroids` by 1-2 (boundary split). When `aggregateGroupByMV` fans the same row's input to a second group key, `prepareSerializedTDigest` calls `input.normalizeBoundaries()` again, and `normalizeSerializedBoundaries` unconditionally requires `count + 2 <= means.length` against the *grown* count - even though the now-unit endpoints need no further split - throwing `IllegalArgumentException: TDigest boundary repair requires two spare centroid entries`. Repro: `SELECT PERCENTILETDIGEST(digestCol, 50) ... GROUP BY mvCol` over a BYTES t-digest column where a row's digest has serialized first/last centroid weight > 1 (typical of t-digest 3.2-era segments, e.g. centroids `{w=3,m=0},{w=2,m=5},{w=3,m=10}`) and its mvCol maps to 2+ groups that already hold digests. Stack: `normalizeSerializedBoundaries:1496 -> SerializedTDigestInput.normalizeBoundaries:158 -> prepareSerializedTDigest:564 -> addSerializedTDigestDirect:523`. It is also intermittent: if the reusable input's arrays are oversized from an earlier larger row, the same call succeeds - a flaky query failure. Fix options: make the capacity precondition account for already-normalized state (skip when endpoints are already unit weight / track a `_boundariesNormalized` flag on the input), or size `decode()`'s arrays for the post-normalization count. -- 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]
