Marks end criteria reached for the segment if the Index cannot consume more rows

pinot-common/src/main/java/org/apache/pinot/common/protocols/SegmentCompletionProtocol.java

@@ -149,6 +149,7 @@ public enum ControllerResponseStatus {
   public static final String REASON_END_OF_PARTITION_GROUP = "endOfPartitionGroup";
   // Stop reason sent by server as force commit message received
   public static final String REASON_FORCE_COMMIT_MESSAGE_RECEIVED = "forceCommitMessageReceived";
+  public static final String REASON_NUM_OF_COL_VALUES_ABOVE_THRESHOLD = "numColValuesAboveThreshold";
 
   // Canned responses
   public static final Response RESP_NOT_LEADER =

pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java

@@ -309,6 +309,7 @@ public void deleteSegmentFile() {
   private String _stopReason = null;
   private final Semaphore _segBuildSemaphore;
   private final boolean _isOffHeap;
+  private final boolean _thresholdForNumOfColValuesEnabled;
   /**
    * Whether null handling is enabled by default. This value is only used if
    * {@link Schema#isEnableColumnBasedNullHandling()} is false.
@@ -362,6 +363,13 @@ private boolean endCriteriaReached() {
               _numRowsConsumed, _numRowsIndexed);
           _stopReason = SegmentCompletionProtocol.REASON_FORCE_COMMIT_MESSAGE_RECEIVED;
           return true;
+        } else if (_thresholdForNumOfColValuesEnabled && _realtimeSegment.isNumOfColValuesAboveThreshold()) {
+          _segmentLogger.info(
+              "Stopping consumption as num of values for a column is above threshold - numRowsConsumed={} "
+                  + "numRowsIndexed={}",
+              _numRowsConsumed, _numRowsIndexed);
+          _stopReason = SegmentCompletionProtocol.REASON_NUM_OF_COL_VALUES_ABOVE_THRESHOLD;
+          return true;
         }
         return false;
 
@@ -1233,6 +1241,7 @@ private static class ConsumptionStopIndicator {
     final Logger _logger;
     final ServerSegmentCompletionProtocolHandler _protocolHandler;
     final String _reason;
+
     private ConsumptionStopIndicator(StreamPartitionMsgOffset offset, String segmentName, String instanceId,
         ServerSegmentCompletionProtocolHandler protocolHandler, String reason, Logger logger) {
       _offset = offset;
@@ -1529,6 +1538,7 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf
 
     _isOffHeap = indexLoadingConfig.isRealtimeOffHeapAllocation();
     _defaultNullHandlingEnabled = indexingConfig.isNullHandlingEnabled();
+    _thresholdForNumOfColValuesEnabled = tableConfig.getValidationConfig().isThresholdForNumOfColValuesEnabled();
 
     // Start new realtime segment
     String consumerDir = realtimeTableDataManager.getConsumerDir();
@@ -1552,7 +1562,8 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf
             .setUpsertDropOutOfOrderRecord(tableConfig.isDropOutOfOrderRecord())
             .setPartitionDedupMetadataManager(partitionDedupMetadataManager)
             .setDedupTimeColumn(tableConfig.getDedupTimeColumn())
-            .setFieldConfigList(tableConfig.getFieldConfigList());
+            .setFieldConfigList(tableConfig.getFieldConfigList())
+            .setThresholdForNumOfColValuesEnabled(_thresholdForNumOfColValuesEnabled);
 
     // Create message decoder
     Set<String> fieldsToRead = IngestionUtils.getFieldsForRecordExtractor(_tableConfig.getIngestionConfig(), _schema);
@@ -1625,7 +1636,7 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf
           "Failed to initialize segment data manager", e));
       _segmentLogger.warn(
           "Scheduling task to call controller to mark the segment as OFFLINE in Ideal State due"
-           + " to initialization error: '{}'",
+              + " to initialization error: '{}'",
           e.getMessage());
       // Since we are going to throw exception from this thread (helix execution thread), the externalview
       // entry for this segment will be ERROR. We allow time for Helix to make this transition, and then

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java

@@ -128,6 +128,7 @@ public class MutableSegmentImpl implements MutableSegment {
   private static final int EXPECTED_COMPRESSION = 1000;
   private static final int MIN_ROWS_TO_INDEX = 1000_000; // Min size of recordIdMap for updatable metrics.
   private static final int MIN_RECORD_ID_MAP_CACHE_SIZE = 10000; // Min overflow map size for updatable metrics.
+  private final static int DEFAULT_THRESHOLD_FOR_NUM_OF_VALUES_PER_COLUMN = 2_000_000_000;
 
   private final Logger _logger;
   private final long _startTimeMillis = System.currentTimeMillis();
@@ -147,8 +148,10 @@ public class MutableSegmentImpl implements MutableSegment {
   private final int _mainPartitionId; // partition id designated for this consuming segment
   private final boolean _defaultNullHandlingEnabled;
   private final File _consumerDir;
+  private final boolean _thresholdForNumOfColValuesEnabled;
 
   private final Map<String, IndexContainer> _indexContainerMap = new HashMap<>();
+  private boolean _numOfColValuesLimitBreached = false;
 
   private final IdMap<FixedIntArray> _recordIdMap;
 
@@ -225,6 +228,7 @@ public boolean isMutableSegment() {
     _mainPartitionId = config.getPartitionId();
     _defaultNullHandlingEnabled = config.isNullHandlingEnabled();
     _consumerDir = new File(config.getConsumerDir());
+    _thresholdForNumOfColValuesEnabled = config.isThresholdForNumOfColValuesEnabled();
 
     Collection<FieldSpec> allFieldSpecs = _schema.getAllFieldSpecs();
     List<FieldSpec> physicalFieldSpecs = new ArrayList<>(allFieldSpecs.size());
@@ -797,6 +801,18 @@ private void addNewRow(int docId, GenericRow row) {
             recordIndexingError(indexEntry.getKey(), e);
           }
         }
+
+        if (_thresholdForNumOfColValuesEnabled) {
+          int prevCount = indexContainer._valuesInfo.getNumValues();
+          long newCount = prevCount + 1L + values.length;
+
+          if (newCount > DEFAULT_THRESHOLD_FOR_NUM_OF_VALUES_PER_COLUMN) {
+            _logger.warn("Number of total values for column {} is {} and has breached the threshold limit {}",
+                column, newCount, DEFAULT_THRESHOLD_FOR_NUM_OF_VALUES_PER_COLUMN);
+            _numOfColValuesLimitBreached = true;
+          }
+        }
+
         indexContainer._valuesInfo.updateMVNumValues(values.length);
       }
     }
@@ -1229,6 +1245,10 @@ private boolean isAggregateMetricsEnabled() {
     return _recordIdMap != null;
   }
 
+  public boolean isNumOfColValuesAboveThreshold() {
+    return _numOfColValuesLimitBreached;
+  }
+
   // NOTE: Okay for single-writer
   @SuppressWarnings("NonAtomicOperationOnVolatileField")
   private static class ValuesInfo {
@@ -1285,6 +1305,10 @@ void updateVarByteMVMaxRowLengthInBytes(Object entry, DataType dataType) {
           throw new IllegalStateException("Invalid type=" + dataType);
       }
     }
+
+    int getNumValues() {
+      return _numValues;
+    }
   }
 
   private static Map<String, Pair<String, ValueAggregator>> getMetricsAggregators(RealtimeSegmentConfig segmentConfig) {

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/RealtimeSegmentConfig.java

@@ -73,6 +73,7 @@ public class RealtimeSegmentConfig {
   private final String _consumerDir;
   private final List<FieldConfig> _fieldConfigList;
   private final List<AggregationConfig> _ingestionAggregationConfigs;
+  private final boolean _thresholdForNumOfColValuesEnabled;
 
   // TODO: Clean up this constructor. Most of these things can be extracted from tableConfig.
 
@@ -89,7 +90,8 @@ private RealtimeSegmentConfig(String tableNameWithType, String segmentName, Stri
       List<String> upsertComparisonColumns, String upsertDeleteRecordColumn, String upsertOutOfOrderRecordColumn,
       boolean upsertDropOutOfOrderRecord, PartitionUpsertMetadataManager partitionUpsertMetadataManager,
       String dedupTimeColumn, PartitionDedupMetadataManager partitionDedupMetadataManager,
-      List<FieldConfig> fieldConfigList, List<AggregationConfig> ingestionAggregationConfigs) {
+      List<FieldConfig> fieldConfigList, List<AggregationConfig> ingestionAggregationConfigs,
+      boolean enableThresholdForNumOfValues) {
     _tableNameWithType = tableNameWithType;
     _segmentName = segmentName;
     _streamName = streamName;
@@ -119,6 +121,7 @@ private RealtimeSegmentConfig(String tableNameWithType, String segmentName, Stri
     _partitionDedupMetadataManager = partitionDedupMetadataManager;
     _fieldConfigList = fieldConfigList;
     _ingestionAggregationConfigs = ingestionAggregationConfigs;
+    _thresholdForNumOfColValuesEnabled = enableThresholdForNumOfValues;
   }
 
   public String getTableNameWithType() {
@@ -241,6 +244,10 @@ public List<AggregationConfig> getIngestionAggregationConfigs() {
     return _ingestionAggregationConfigs;
   }
 
+  public boolean isThresholdForNumOfColValuesEnabled() {
+    return _thresholdForNumOfColValuesEnabled;
+  }
+
   public static class Builder {
     private String _tableNameWithType;
     private String _segmentName;
@@ -275,6 +282,7 @@ public static class Builder {
     private PartitionDedupMetadataManager _partitionDedupMetadataManager;
     private List<FieldConfig> _fieldConfigList;
     private List<AggregationConfig> _ingestionAggregationConfigs;
+    private boolean _thresholdForNumOfColValuesEnabled = false;
 
     public Builder() {
       _indexConfigByCol = new HashMap<>();
@@ -475,6 +483,11 @@ public Builder setIngestionAggregationConfigs(List<AggregationConfig> ingestionA
       return this;
     }
 
+    public Builder setThresholdForNumOfColValuesEnabled(boolean thresholdForNumOfColValuesEnabled) {
+      _thresholdForNumOfColValuesEnabled = thresholdForNumOfColValuesEnabled;
+      return this;
+    }
+
     public RealtimeSegmentConfig build() {
       Map<String, FieldIndexConfigs> indexConfigByCol = Maps.newHashMapWithExpectedSize(_indexConfigByCol.size());
       for (Map.Entry<String, FieldIndexConfigs.Builder> entry : _indexConfigByCol.entrySet()) {
@@ -487,7 +500,7 @@ public RealtimeSegmentConfig build() {
           _defaultNullHandlingEnabled, _consumerDir, _upsertMode, _upsertConsistencyMode, _upsertComparisonColumns,
           _upsertDeleteRecordColumn, _upsertOutOfOrderRecordColumn, _upsertDropOutOfOrderRecord,
           _partitionUpsertMetadataManager, _dedupTimeColumn, _partitionDedupMetadataManager, _fieldConfigList,
-          _ingestionAggregationConfigs);
+          _ingestionAggregationConfigs, _thresholdForNumOfColValuesEnabled);
     }
   }
 }

pinot-spi/src/main/java/org/apache/pinot/spi/config/table/SegmentsValidationAndRetentionConfig.java

@@ -51,6 +51,7 @@ public class SegmentsValidationAndRetentionConfig extends BaseJsonConfig {
   // using the specified download scheme. Both realtime tables and offline tables can set this field.
   // For more usage of this field, please refer to this design doc: https://tinyurl.com/f63ru4sb
   private String _peerSegmentDownloadScheme;
+  private boolean _thresholdForNumOfColValuesEnabled;
 
   @Deprecated
   public String getSegmentAssignmentStrategy() {
@@ -218,6 +219,14 @@ public void setPeerSegmentDownloadScheme(String peerSegmentDownloadScheme) {
     _peerSegmentDownloadScheme = peerSegmentDownloadScheme;
   }
 
+  public boolean isThresholdForNumOfColValuesEnabled() {
+    return _thresholdForNumOfColValuesEnabled;
+  }
+
+  public void setThresholdForNumOfColValuesEnabled(boolean thresholdForNumOfColValuesEnabled) {
+    _thresholdForNumOfColValuesEnabled = thresholdForNumOfColValuesEnabled;
+  }
+
   public String getCrypterClassName() {
     return _crypterClassName;
   }

pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/TableConfigBuilder.java

@@ -81,6 +81,7 @@ public class TableConfigBuilder {
   private ReplicaGroupStrategyConfig _replicaGroupStrategyConfig;
   private CompletionConfig _completionConfig;
   private String _crypterClassName;
+  private boolean _thresholdForNumOfColValuesEnabled = false;
 
   // Tenant config related
   private String _brokerTenant;
@@ -439,6 +440,7 @@ public TableConfig build() {
     validationConfig.setReplication(_numReplicas);
     validationConfig.setPeerSegmentDownloadScheme(_peerSegmentDownloadScheme);
     validationConfig.setCrypterClassName(_crypterClassName);
+    validationConfig.setThresholdForNumOfColValuesEnabled(_thresholdForNumOfColValuesEnabled);
 
     // Tenant config
     TenantConfig tenantConfig = new TenantConfig(_brokerTenant, _serverTenant, _tagOverrideConfig);