xiangfu0 commented on code in PR #19795:
URL: https://github.com/apache/pinot/pull/19795#discussion_r4253011592


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/TDigestAccumulator.java:
##########
@@ -104,64 +110,201 @@ public TDigestAccumulator(int compression) {
     this((double) compression, true, true);
   }
 
+  public TDigestAccumulator(double compression) {
+    this(compression, true, true);
+  }
+
   private TDigestAccumulator(double compression, boolean allocateRawBuffer, 
boolean allocateMergeBuffers) {
-    if (!(compression > 0.0) || !Double.isFinite(compression)) {
-      throw new IllegalArgumentException("TDigest compression must be 
positive: " + compression);
-    }
-    super.setScaleFunction(ScaleFunction.K_1);
-    _compression = Math.max(MIN_COMPRESSION, compression);
+    this(compression, allocateRawBuffer, allocateMergeBuffers, false);
+  }
+
+  private TDigestAccumulator(double compression, boolean allocateRawBuffer, 
boolean allocateMergeBuffers,
+      boolean useTwoLevelRawCompression) {
+    _compression = TDigestCodec.normalizeCompression(compression);
+    _useTwoLevelRawCompression = useTwoLevelRawCompression;
     _workingCompression = TWO_LEVEL_COMPRESSION_MULTIPLIER * _compression;
     _publicMaxWeightScale = calculateK1MaxWeightScale(_compression);
     _workingMaxWeightScale = calculateK1MaxWeightScale(_workingCompression);
     long roundedCompression = (long) Math.ceil(_compression);
-    _centroidCapacity = getCentroidCapacity(_compression);
-    _pendingCentroidCapacity = getPendingCentroidCapacity(_centroidCapacity);
+    _centroidCapacity = TDigestCodec.getDefaultCentroidCapacity(_compression);
+    _pendingCentroidCapacity = getPendingCentroidCapacity(_centroidCapacity, 
useTwoLevelRawCompression);
     if (allocateRawBuffer) {
-      _rawValues = new double[getRawBufferSize(roundedCompression)];
+      _rawValues = new double[getInitialRawBufferSize(roundedCompression)];
     }
     if (allocateMergeBuffers) {
-      _centroidMeans = new double[_centroidCapacity];
-      _centroidWeights = new double[_centroidCapacity];
-      _outputMeans = new double[_centroidCapacity];
-      _outputWeights = new double[_centroidCapacity];
+      int initialCapacity = Math.min(_centroidCapacity, MIN_RAW_BUFFER_SIZE);
+      _centroidMeans = new double[initialCapacity];
+      _centroidWeights = new double[initialCapacity];
+      _outputMeans = new double[initialCapacity];
+      _outputWeights = new double[initialCapacity];
     }
   }
 
-  public static TDigestAccumulator forSerializedTDigest(byte[] bytes) {
-    return new TDigestAccumulator(readCompression(bytes), false, false);
+  /// Creates a lazy digest with the legacy two-level compression used by 
Star-tree and general digest aggregation.
+  public static TDigestAccumulator forLegacyAggregation(double compression) {
+    return new TDigestAccumulator(compression, false, false, true);
   }
 
-  public static TDigestAccumulator forSerializedTDigest(SerializedTDigestInput 
input) {
-    return new TDigestAccumulator(input._compression, false, false);
+  /// Reads and retains one legacy digest, preserving the general aggregation 
compression policy.
+  public static TDigestAccumulator fromBytes(byte[] bytes) {
+    return fromBytes(ByteBuffer.wrap(bytes));
   }
 
-  public static TDigestAccumulator forSerializedTDigestWithMergeBuffers(byte[] 
bytes) {
-    return new TDigestAccumulator(readCompression(bytes), false, true);
+  /// Reads one payload and advances the buffer past it.
+  public static TDigestAccumulator fromBytes(ByteBuffer bytes) {
+    SerializedTDigestInput input = new SerializedTDigestInput();
+    input.reset(bytes);
+    return fromBytes(input);
   }
 
-  public static TDigestAccumulator forSerializedTDigest(ByteBuffer input) {
-    byte[] bytes = new byte[input.remaining()];
-    input.get(bytes);
-    TDigestAccumulator accumulator = forSerializedTDigest(bytes);
-    accumulator.addSerializedTDigest(bytes);
-    return accumulator;
+  /// Retains an already validated input without rescanning its centroid bytes.
+  public static TDigestAccumulator fromBytes(SerializedTDigestInput input) {
+    TDigestAccumulator digest = forLegacyAggregation(input.getCompression());
+    digest.addSerializedTDigest(input);
+    return digest;
   }
 
   public static TDigestAccumulator forReduction(double compression) {
     return new TDigestAccumulator(compression, false, false);
   }
 
+  /// Snapshots buffered inputs and centroid state without materializing or 
compressing the source. Mutable buffers
+  /// are copied; immutable retained bytes and metadata are shared. Merge and 
sort scratch buffers remain lazy.
+  TDigestAccumulator copy() {
+    TDigestAccumulator copy =
+        new TDigestAccumulator(_compression, false, false, 
_useTwoLevelRawCompression);
+    copy._rawValues = _rawValues != null ? _rawValues.clone() : null;

Review Comment:
   Addressed in 564fec0064. Serialization snapshots now copy only active raw 
values; ordinary copies retain capacity and future flush timing. All 10 local 
snapshot byte hashes match, and the existing regression now covers continued 
updates to a copy. Allocation for the one-value buffered RAW case fell from 
5,976 to 3,936 bytes in the local probe.



##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/TDigestAccumulator.java:
##########
@@ -104,64 +110,201 @@ public TDigestAccumulator(int compression) {
     this((double) compression, true, true);
   }
 
+  public TDigestAccumulator(double compression) {
+    this(compression, true, true);
+  }
+
   private TDigestAccumulator(double compression, boolean allocateRawBuffer, 
boolean allocateMergeBuffers) {
-    if (!(compression > 0.0) || !Double.isFinite(compression)) {
-      throw new IllegalArgumentException("TDigest compression must be 
positive: " + compression);
-    }
-    super.setScaleFunction(ScaleFunction.K_1);
-    _compression = Math.max(MIN_COMPRESSION, compression);
+    this(compression, allocateRawBuffer, allocateMergeBuffers, false);
+  }
+
+  private TDigestAccumulator(double compression, boolean allocateRawBuffer, 
boolean allocateMergeBuffers,
+      boolean useTwoLevelRawCompression) {
+    _compression = TDigestCodec.normalizeCompression(compression);
+    _useTwoLevelRawCompression = useTwoLevelRawCompression;
     _workingCompression = TWO_LEVEL_COMPRESSION_MULTIPLIER * _compression;
     _publicMaxWeightScale = calculateK1MaxWeightScale(_compression);
     _workingMaxWeightScale = calculateK1MaxWeightScale(_workingCompression);
     long roundedCompression = (long) Math.ceil(_compression);
-    _centroidCapacity = getCentroidCapacity(_compression);
-    _pendingCentroidCapacity = getPendingCentroidCapacity(_centroidCapacity);
+    _centroidCapacity = TDigestCodec.getDefaultCentroidCapacity(_compression);
+    _pendingCentroidCapacity = getPendingCentroidCapacity(_centroidCapacity, 
useTwoLevelRawCompression);
     if (allocateRawBuffer) {
-      _rawValues = new double[getRawBufferSize(roundedCompression)];
+      _rawValues = new double[getInitialRawBufferSize(roundedCompression)];
     }
     if (allocateMergeBuffers) {
-      _centroidMeans = new double[_centroidCapacity];
-      _centroidWeights = new double[_centroidCapacity];
-      _outputMeans = new double[_centroidCapacity];
-      _outputWeights = new double[_centroidCapacity];
+      int initialCapacity = Math.min(_centroidCapacity, MIN_RAW_BUFFER_SIZE);
+      _centroidMeans = new double[initialCapacity];
+      _centroidWeights = new double[initialCapacity];
+      _outputMeans = new double[initialCapacity];
+      _outputWeights = new double[initialCapacity];
     }
   }
 
-  public static TDigestAccumulator forSerializedTDigest(byte[] bytes) {
-    return new TDigestAccumulator(readCompression(bytes), false, false);
+  /// Creates a lazy digest with the legacy two-level compression used by 
Star-tree and general digest aggregation.
+  public static TDigestAccumulator forLegacyAggregation(double compression) {
+    return new TDigestAccumulator(compression, false, false, true);
   }
 
-  public static TDigestAccumulator forSerializedTDigest(SerializedTDigestInput 
input) {
-    return new TDigestAccumulator(input._compression, false, false);
+  /// Reads and retains one legacy digest, preserving the general aggregation 
compression policy.
+  public static TDigestAccumulator fromBytes(byte[] bytes) {
+    return fromBytes(ByteBuffer.wrap(bytes));
   }
 
-  public static TDigestAccumulator forSerializedTDigestWithMergeBuffers(byte[] 
bytes) {
-    return new TDigestAccumulator(readCompression(bytes), false, true);
+  /// Reads one payload and advances the buffer past it.
+  public static TDigestAccumulator fromBytes(ByteBuffer bytes) {
+    SerializedTDigestInput input = new SerializedTDigestInput();
+    input.reset(bytes);
+    return fromBytes(input);
   }
 
-  public static TDigestAccumulator forSerializedTDigest(ByteBuffer input) {
-    byte[] bytes = new byte[input.remaining()];
-    input.get(bytes);
-    TDigestAccumulator accumulator = forSerializedTDigest(bytes);
-    accumulator.addSerializedTDigest(bytes);
-    return accumulator;
+  /// Retains an already validated input without rescanning its centroid bytes.
+  public static TDigestAccumulator fromBytes(SerializedTDigestInput input) {
+    TDigestAccumulator digest = forLegacyAggregation(input.getCompression());
+    digest.addSerializedTDigest(input);
+    return digest;
   }
 
   public static TDigestAccumulator forReduction(double compression) {
     return new TDigestAccumulator(compression, false, false);
   }
 
+  /// Snapshots buffered inputs and centroid state without materializing or 
compressing the source. Mutable buffers
+  /// are copied; immutable retained bytes and metadata are shared. Merge and 
sort scratch buffers remain lazy.
+  TDigestAccumulator copy() {
+    TDigestAccumulator copy =
+        new TDigestAccumulator(_compression, false, false, 
_useTwoLevelRawCompression);
+    copy._rawValues = _rawValues != null ? _rawValues.clone() : null;
+    copy._centroidMeans = _numCentroids > 0 ? Arrays.copyOf(_centroidMeans, 
_numCentroids) : null;
+    copy._centroidWeights = _numCentroids > 0 ? 
Arrays.copyOf(_centroidWeights, _numCentroids) : null;
+    copy._incomingMeans = _numIncomingCentroids > 0 ? 
Arrays.copyOf(_incomingMeans, _numIncomingCentroids) : null;
+    copy._incomingWeights = _numIncomingCentroids > 0 ? 
Arrays.copyOf(_incomingWeights, _numIncomingCentroids) : null;
+    copy._pendingSerializedTDigest = _pendingSerializedTDigest;
+    copy._pendingSerializedMetadata = _pendingSerializedMetadata;
+    copy._originalFractionalBytes = _originalFractionalBytes;
+    copy._originalBoundaryBytes = _originalBoundaryBytes;
+    copy._historicalFractionalMinMean = _historicalFractionalMinMean;
+    copy._historicalFractionalMaxMean = _historicalFractionalMaxMean;
+    copy._freshFractionalMinMean = _freshFractionalMinMean;
+    copy._freshFractionalMaxMean = _freshFractionalMaxMean;
+    copy._legacyDegraded = _legacyDegraded;
+    copy._hasFractionalWeights = _hasFractionalWeights;
+    copy._deferBoundaryNormalization = _deferBoundaryNormalization;
+    copy._numRawValues = _numRawValues;
+    copy._numCentroids = _numCentroids;
+    copy._numIncomingCentroids = _numIncomingCentroids;
+    copy._mergeCount = _mergeCount;
+    copy._publiclyCompressed = _publiclyCompressed;
+    copy._incomingCentroidsSorted = _incomingCentroidsSorted;
+    copy._totalWeight = _totalWeight;
+    copy._incomingWeight = _incomingWeight;
+    copy._min = _min;
+    copy._max = _max;
+    return copy;
+  }
+
+  /// Copies primitive centroids without encoding and decoding them. Zero 
weights are ignored, and unordered
+  /// means are sorted when initializing a stored, already compressed 
distribution; caller arrays are not mutated.
+  /// Subsequent inputs use the normal buffered merge. Positive weights and 
consistent extrema are required.
+  public void addCentroids(double[] means, double[] weights, int count, double 
min, double max,
+      boolean alreadyCompressed) {
+    addCentroids(means, weights, count, min, max, alreadyCompressed, false);
+  }
+
+  /// Adds validated legacy centroids without marking their existing 
fractional mass as fresh input.
+  void addHistoricalCentroids(double[] means, double[] weights, int count, 
double min, double max,
+      boolean alreadyCompressed) {
+    addCentroids(means, weights, count, min, max, alreadyCompressed, true);
+  }
+
+  private void addCentroids(double[] means, double[] weights, int count, 
double min, double max,
+      boolean alreadyCompressed, boolean historical) {
+    requireMutable();
+    if (count < 0 || count > means.length || count > weights.length) {
+      throw new IllegalArgumentException("Invalid TDigest centroid count: " + 
count);
+    }
+    double totalWeight = 0.0;
+    int positiveCount = 0;
+    boolean sorted = true;
+    boolean fractionalWeights = false;
+    double previousMean = Double.NEGATIVE_INFINITY;
+    for (int i = 0; i < count; i++) {
+      double weight = weights[i];
+      if (weight == 0.0) {
+        continue;
+      }
+      double mean = means[i];
+      if (!(weight > 0.0) || !Double.isFinite(weight) || !(mean >= min && mean 
<= max)) {
+        throw new IllegalArgumentException("Invalid TDigest centroid");
+      }
+      totalWeight += weight;
+      sorted &= mean >= previousMean;
+      previousMean = mean;
+      fractionalWeights |= weight != Math.rint(weight);
+      positiveCount++;
+    }
+    checkTotalWeight(getTotalWeight() + totalWeight);
+    if (positiveCount == 0) {
+      return;
+    }
+    prepareForMutation();
+    _originalFractionalBytes = null;
+    _hasFractionalWeights |= fractionalWeights;
+    if (!historical && fractionalWeights) {
+      for (int i = 0; i < count; i++) {
+        recordFreshFractionalMean(means[i], weights[i]);
+      }
+    }
+    if (hasNoInputs() && alreadyCompressed) {
+      ensureCentroidCapacity(Math.addExact(positiveCount, 2));
+      for (int i = 0; i < count; i++) {
+        if (weights[i] > 0.0) {
+          _centroidMeans[_numCentroids] = means[i];
+          _centroidWeights[_numCentroids++] = weights[i];
+        }
+      }
+      if (!sorted) {
+        sortCentroids(_centroidMeans, _centroidWeights, _numCentroids);
+      }
+      _totalWeight = totalWeight;
+      _min = min;
+      _max = max;
+      if (!historical) {
+        normalizeBoundaryCentroids();
+      }
+      _publiclyCompressed = true;
+    } else {
+      if (sorted && positiveCount == count) {
+        bufferIncomingCentroids(means, weights, count, totalWeight, min, max);
+      } else {
+        // Caller-owned arrays are copied only when filtering or sorting is 
necessary.
+        double[] sortedMeans = new double[positiveCount];
+        double[] sortedWeights = new double[positiveCount];
+        int index = 0;
+        for (int i = 0; i < count; i++) {
+          if (weights[i] > 0.0) {
+            sortedMeans[index] = means[i];
+            sortedWeights[index++] = weights[i];
+          }
+        }
+        if (!sorted) {
+          sortCentroids(sortedMeans, sortedWeights, positiveCount);
+        }
+        bufferIncomingCentroids(sortedMeans, sortedWeights, positiveCount, 
totalWeight, min, max);
+      }
+      _min = Math.min(_min, min);
+      _max = Math.max(_max, max);

Review Comment:
   Addressed in 564fec0064. Removed both duplicate outer extrema updates. 
bufferIncomingCentroids owns those updates on both paths; the separate 
serialized-merge extrema handling remains. The affected extrema, infinity and 
merge regressions pass.



-- 
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]

Reply via email to