From e2dba55c3b1cac0402f7852375d2facadcc63a22 Mon Sep 17 00:00:00 2001
From: ankitsultana <ankitsultana@uber.com>
Date: Wed, 20 Nov 2024 19:58:34 +0000
Subject: [PATCH] add limit for max data points

---
 .../combine/merger/TimeSeriesAggResultsBlockMerger.java  | 9 ++++++++-
 .../timeseries/TimeSeriesAggregationOperator.java        | 5 +++++
 .../pinot/tsdb/spi/series/TimeSeriesBuilderFactory.java  | 8 ++++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/merger/TimeSeriesAggResultsBlockMerger.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/merger/TimeSeriesAggResultsBlockMerger.java
index 124cad46aff6..428cfde55511 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/merger/TimeSeriesAggResultsBlockMerger.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/merger/TimeSeriesAggResultsBlockMerger.java
@@ -30,11 +30,13 @@ public class TimeSeriesAggResultsBlockMerger implements ResultsBlockMerger<TimeS
   private final TimeSeriesBuilderFactory _seriesBuilderFactory;
   private final AggInfo _aggInfo;
   private final int _maxSeriesLimit;
+  private final long _maxDataPointsLimit;
 
   public TimeSeriesAggResultsBlockMerger(TimeSeriesBuilderFactory seriesBuilderFactory, AggInfo aggInfo) {
     _seriesBuilderFactory = seriesBuilderFactory;
     _aggInfo = aggInfo;
     _maxSeriesLimit = _seriesBuilderFactory.getMaxUniqueSeriesPerServerLimit();
+    _maxDataPointsLimit = _seriesBuilderFactory.getMaxDataPointsPerServerLimit();
   }
 
   @Override
@@ -47,7 +49,12 @@ public void mergeResultsBlocks(TimeSeriesResultsBlock mergedBlock, TimeSeriesRes
       BaseTimeSeriesBuilder newTimeSeriesToMerge = entry.getValue();
       if (currentTimeSeriesBuilder == null) {
         currentTimeSeriesBlock.getSeriesBuilderMap().put(seriesHash, newTimeSeriesToMerge);
-        Preconditions.checkState(currentTimeSeriesBlock.getSeriesBuilderMap().size() <= _maxSeriesLimit,
+        final long currentUniqueSeries = currentTimeSeriesBlock.getSeriesBuilderMap().size();
+        final long numBuckets = currentTimeSeriesBlock.getTimeBuckets().getNumBuckets();
+        Preconditions.checkState(currentUniqueSeries * numBuckets <= _maxDataPointsLimit,
+            "Max data points limit reached in combine operator. Limit: %s. Current count: %s",
+            _maxDataPointsLimit, currentUniqueSeries * numBuckets);
+        Preconditions.checkState(currentUniqueSeries <= _maxSeriesLimit,
             "Max series limit reached in combine operator. Limit: %s. Current Size: %s",
             _maxSeriesLimit, currentTimeSeriesBlock.getSeriesBuilderMap().size());
       } else {
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/timeseries/TimeSeriesAggregationOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/timeseries/TimeSeriesAggregationOperator.java
index 749ab35edee9..ee7c39c8c855 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/timeseries/TimeSeriesAggregationOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/timeseries/TimeSeriesAggregationOperator.java
@@ -62,6 +62,7 @@ public class TimeSeriesAggregationOperator extends BaseOperator<TimeSeriesResult
   private final TimeBuckets _timeBuckets;
   private final TimeSeriesBuilderFactory _seriesBuilderFactory;
   private final int _maxSeriesLimit;
+  private final long _maxDataPointsLimit;
   private final long _numTotalDocs;
   private long _numDocsScanned = 0;
 
@@ -86,6 +87,7 @@ public TimeSeriesAggregationOperator(
     _timeBuckets = timeBuckets;
     _seriesBuilderFactory = seriesBuilderFactory;
     _maxSeriesLimit = _seriesBuilderFactory.getMaxUniqueSeriesPerServerLimit();
+    _maxDataPointsLimit = _seriesBuilderFactory.getMaxDataPointsPerServerLimit();
     _numTotalDocs = segmentMetadata.getTotalDocs();
   }
 
@@ -138,6 +140,9 @@ protected TimeSeriesResultsBlock getNextBlock() {
           throw new IllegalStateException(
               "Don't yet support value expression of type: " + valueExpressionBlockValSet.getValueType());
       }
+      Preconditions.checkState(seriesBuilderMap.size() * (long) _timeBuckets.getNumBuckets() <= _maxDataPointsLimit,
+          "Query exceed max data point limit per server. Limit: %s. Data points in current segment so far: %s",
+          _maxDataPointsLimit, seriesBuilderMap.size() * _timeBuckets.getNumBuckets());
       Preconditions.checkState(seriesBuilderMap.size() <= _maxSeriesLimit,
           "Query exceeded max unique series limit per server. Limit: %s. Series in current segment so far: %s",
           _maxSeriesLimit, seriesBuilderMap.size());
diff --git a/pinot-timeseries/pinot-timeseries-spi/src/main/java/org/apache/pinot/tsdb/spi/series/TimeSeriesBuilderFactory.java b/pinot-timeseries/pinot-timeseries-spi/src/main/java/org/apache/pinot/tsdb/spi/series/TimeSeriesBuilderFactory.java
index 66d71faa26f9..2db2d4c49d37 100644
--- a/pinot-timeseries/pinot-timeseries-spi/src/main/java/org/apache/pinot/tsdb/spi/series/TimeSeriesBuilderFactory.java
+++ b/pinot-timeseries/pinot-timeseries-spi/src/main/java/org/apache/pinot/tsdb/spi/series/TimeSeriesBuilderFactory.java
@@ -26,6 +26,10 @@
 
 public abstract class TimeSeriesBuilderFactory {
   private static final int DEFAULT_MAX_UNIQUE_SERIES_PER_SERVER_LIMIT = 100_000;
+  /**
+   * Default limit for the total number of values across all series.
+   */
+  private static final long DEFAULT_MAX_DATA_POINTS_PER_SERVER_LIMIT = 100_000_000;
 
   public abstract BaseTimeSeriesBuilder newTimeSeriesBuilder(
       AggInfo aggInfo,
@@ -38,5 +42,9 @@ public int getMaxUniqueSeriesPerServerLimit() {
     return DEFAULT_MAX_UNIQUE_SERIES_PER_SERVER_LIMIT;
   }
 
+  public long getMaxDataPointsPerServerLimit() {
+    return DEFAULT_MAX_DATA_POINTS_PER_SERVER_LIMIT;
+  }
+
   public abstract void init(PinotConfiguration pinotConfiguration);
 }