Fix and clean-up upsert and dedup config (apache#15528)

Author: Jackie-Jiang

pinot-common/src/test/java/org/apache/pinot/common/utils/config/TableConfigSerDeTest.java

@@ -262,7 +262,8 @@ public void testSerDe()
     }
     {
       // with dedup config - without metadata ttl and metadata time column
-      DedupConfig dedupConfig = new DedupConfig(true, HashFunction.MD5);
+      DedupConfig dedupConfig = new DedupConfig();
+      dedupConfig.setHashFunction(HashFunction.MD5);
       TableConfig tableConfig = tableConfigBuilder.setDedupConfig(dedupConfig).build();
       // Serialize then de-serialize
       checkTableConfigWithDedupConfigWithoutTTL(
@@ -272,7 +273,10 @@ public void testSerDe()
     }
     {
       // with dedup config - with metadata ttl and metadata time column
-      DedupConfig dedupConfig = new DedupConfig(true, HashFunction.MD5, null, null, 10, "dedupTimeColumn", false);
+      DedupConfig dedupConfig = new DedupConfig();
+      dedupConfig.setHashFunction(HashFunction.MD5);
+      dedupConfig.setMetadataTTL(10);
+      dedupConfig.setDedupTimeColumn("dedupTimeColumn");
       TableConfig tableConfig = tableConfigBuilder.setDedupConfig(dedupConfig).build();
       // Serialize then de-serialize
       checkTableConfigWithDedupConfigWithTTL(JsonUtils.stringToObject(tableConfig.toJsonString(), TableConfig.class));

pinot-common/src/test/java/org/apache/pinot/common/utils/config/TableConfigTest.java

@@ -91,14 +91,17 @@ public void testCopyConstructor() {
     ingestionConfig.setRowTimeValueCheck(true);
     ingestionConfig.setSegmentTimeValueCheck(false);
 
+    DedupConfig dedupConfig = new DedupConfig();
+    dedupConfig.setHashFunction(HashFunction.MD5);
+
     TableConfig config = new TableConfigBuilder(TableType.OFFLINE)
         .setTableName(RAW_TABLE_NAME)
         .setAggregateMetrics(true)
         .setRetentionTimeValue("5")
         .setRetentionTimeUnit("DAYS")
         .setNumReplicas(2)
-        .setDedupConfig(new DedupConfig(true, HashFunction.MD5))
         .setIngestionConfig(ingestionConfig)
+        .setDedupConfig(dedupConfig)
         .setQueryConfig(new QueryConfig(2000L, true, false, Collections.emptyMap(), 100_000L, 100_000L))
         .setTierConfigList(List.of(new TierConfig("name", "type", null, null, "storageType", null, null, null)))
         .build();

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java

@@ -2546,7 +2546,7 @@ private Map<InstancePartitionsType, InstancePartitions> fetchOrComputeInstancePa
           InstancePartitionsUtils.fetchOrComputeInstancePartitions(_helixZkManager, tableConfig,
               InstancePartitionsType.OFFLINE));
     }
-    if (tableConfig.getUpsertMode() != UpsertConfig.Mode.NONE) {
+    if (tableConfig.isUpsertEnabled()) {
       // In an upsert enabled LLC realtime table, all segments of the same partition are collocated on the same server
       // -- consuming or completed. So it is fine to use CONSUMING as the InstancePartitionsType.
       return Collections.singletonMap(InstancePartitionsType.CONSUMING,

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceConfig.java

@@ -86,6 +86,7 @@ public class RebalanceConfig {
   // Whether to run Minimal Data Movement Algorithm, overriding the minimizeDataMovement flag in table config. If set
   // to default, the minimizeDataMovement flag in table config will be used to determine whether to run the Minimal
   // Data Movement Algorithm.
+  // TODO: Replace this with Enablement
   @ApiModel
   public enum MinimizeDataMovementOptions {
     ENABLE, DISABLE, DEFAULT

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

@@ -54,6 +54,7 @@
 import org.apache.pinot.common.utils.TarCompressionUtils;
 import org.apache.pinot.core.data.manager.realtime.RealtimeConsumptionRateManager.ConsumptionRateLimiter;
 import org.apache.pinot.segment.local.data.manager.SegmentDataManager;
+import org.apache.pinot.segment.local.dedup.DedupContext;
 import org.apache.pinot.segment.local.dedup.PartitionDedupMetadataManager;
 import org.apache.pinot.segment.local.indexsegment.mutable.MutableSegmentImpl;
 import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager;
@@ -63,6 +64,7 @@
 import org.apache.pinot.segment.local.segment.creator.TransformPipeline;
 import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig;
 import org.apache.pinot.segment.local.upsert.PartitionUpsertMetadataManager;
+import org.apache.pinot.segment.local.upsert.UpsertContext;
 import org.apache.pinot.segment.local.utils.IngestionUtils;
 import org.apache.pinot.segment.spi.MutableSegment;
 import org.apache.pinot.segment.spi.V1Constants;
@@ -74,7 +76,6 @@
 import org.apache.pinot.spi.config.instance.InstanceDataManagerConfig;
 import org.apache.pinot.spi.config.table.ColumnPartitionConfig;
 import org.apache.pinot.spi.config.table.CompletionConfig;
-import org.apache.pinot.spi.config.table.DedupConfig;
 import org.apache.pinot.spi.config.table.IndexingConfig;
 import org.apache.pinot.spi.config.table.SegmentPartitionConfig;
 import org.apache.pinot.spi.config.table.SegmentZKPropsConfig;
@@ -761,7 +762,7 @@ public void run() {
         //   persisted.
         // Take upsert snapshot before starting consuming events
         if (_partitionUpsertMetadataManager != null) {
-          if (_tableConfig.getUpsertMetadataTTL() > 0) {
+          if (_partitionUpsertMetadataManager.getContext().getMetadataTTL() > 0) {
             // If upsertMetadataTTL is enabled, we will remove expired primary keys from upsertMetadata
             // AFTER taking a snapshot. Taking the snapshot first is crucial to capture the final
             // state of each key before it exits the TTL window. Out-of-TTL segments are skipped in
@@ -779,7 +780,8 @@ public void run() {
           }
         }
 
-        if (_partitionDedupMetadataManager != null && _tableConfig.getDedupMetadataTTL() > 0) {
+        if (_partitionDedupMetadataManager != null
+            && _partitionDedupMetadataManager.getContext().getMetadataTTL() > 0) {
           _partitionDedupMetadataManager.removeExpiredPrimaryKeys();
         }
 
@@ -1662,27 +1664,24 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf
 
     // Start new realtime segment
     String consumerDir = realtimeTableDataManager.getConsumerDir();
-    RealtimeSegmentConfig.Builder realtimeSegmentConfigBuilder =
-        new RealtimeSegmentConfig.Builder(indexLoadingConfig).setTableNameWithType(_tableNameWithType)
-            .setSegmentName(_segmentNameStr)
-            .setStreamName(streamTopic).setSchema(_schema).setTimeColumnName(timeColumnName)
-            .setCapacity(_segmentMaxRowCount).setAvgNumMultiValues(indexLoadingConfig.getRealtimeAvgMultiValueCount())
-            .setSegmentZKMetadata(segmentZKMetadata)
-            .setOffHeap(_isOffHeap).setMemoryManager(_memoryManager)
-            .setStatsHistory(realtimeTableDataManager.getStatsHistory())
-            .setAggregateMetrics(indexingConfig.isAggregateMetrics())
-            .setIngestionAggregationConfigs(IngestionConfigUtils.getAggregationConfigs(tableConfig))
-            .setDefaultNullHandlingEnabled(_defaultNullHandlingEnabled)
-            .setConsumerDir(consumerDir).setUpsertMode(tableConfig.getUpsertMode())
-            .setUpsertConsistencyMode(tableConfig.getUpsertConsistencyMode())
-            .setPartitionUpsertMetadataManager(partitionUpsertMetadataManager)
-            .setUpsertComparisonColumns(tableConfig.getUpsertComparisonColumns())
-            .setUpsertDeleteRecordColumn(tableConfig.getUpsertDeleteRecordColumn())
-            .setUpsertOutOfOrderRecordColumn(tableConfig.getOutOfOrderRecordColumn())
-            .setUpsertDropOutOfOrderRecord(tableConfig.isDropOutOfOrderRecord())
-            .setPartitionDedupMetadataManager(partitionDedupMetadataManager)
-            .setDedupTimeColumn(tableConfig.getDedupTimeColumn())
-            .setFieldConfigList(tableConfig.getFieldConfigList());
+    RealtimeSegmentConfig.Builder realtimeSegmentConfigBuilder = new RealtimeSegmentConfig.Builder(indexLoadingConfig)
+        .setTableNameWithType(_tableNameWithType)
+        .setSegmentName(_segmentNameStr)
+        .setStreamName(streamTopic)
+        .setSchema(_schema)
+        .setTimeColumnName(timeColumnName)
+        .setCapacity(_segmentMaxRowCount)
+        .setAvgNumMultiValues(indexLoadingConfig.getRealtimeAvgMultiValueCount())
+        .setSegmentZKMetadata(segmentZKMetadata)
+        .setOffHeap(_isOffHeap)
+        .setMemoryManager(_memoryManager)
+        .setStatsHistory(realtimeTableDataManager.getStatsHistory())
+        .setAggregateMetrics(indexingConfig.isAggregateMetrics())
+        .setIngestionAggregationConfigs(IngestionConfigUtils.getAggregationConfigs(tableConfig))
+        .setDefaultNullHandlingEnabled(_defaultNullHandlingEnabled)
+        .setPartitionUpsertMetadataManager(partitionUpsertMetadataManager)
+        .setPartitionDedupMetadataManager(partitionDedupMetadataManager)
+        .setConsumerDir(consumerDir);
 
     // Create message decoder
     Set<String> fieldsToRead = IngestionUtils.getFieldsForRecordExtractor(_tableConfig, _schema);
@@ -1761,8 +1760,7 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf
     }
   }
 
-  @VisibleForTesting
-  ParallelSegmentConsumptionPolicy getParallelConsumptionPolicy() {
+  private ParallelSegmentConsumptionPolicy getParallelConsumptionPolicy() {
     IngestionConfig ingestionConfig = _tableConfig.getIngestionConfig();
     ParallelSegmentConsumptionPolicy parallelSegmentConsumptionPolicy = null;
     boolean pauseless = false;
@@ -1782,19 +1780,18 @@ ParallelSegmentConsumptionPolicy getParallelConsumptionPolicy() {
     //   - For pauseless tables, allow consumption during build, but disallow consumption during download
     //   - For non-pauseless tables, disallow consumption during build and download
     //     TODO: Revisit the non-pauseless handling
-    if (_realtimeTableDataManager.isDedupEnabled()) {
-      DedupConfig dedupConfig = _tableConfig.getDedupConfig();
-      assert dedupConfig != null;
-      if (dedupConfig.isAllowDedupConsumptionDuringCommit()) {
+    if (_partitionUpsertMetadataManager != null) {
+      UpsertContext upsertContext = _partitionUpsertMetadataManager.getContext();
+      if (upsertContext.isAllowPartialUpsertConsumptionDuringCommit()
+          || upsertContext.getUpsertMode() != UpsertConfig.Mode.PARTIAL) {
         return ParallelSegmentConsumptionPolicy.ALLOW_ALWAYS;
       }
       return pauseless ? ParallelSegmentConsumptionPolicy.ALLOW_DURING_BUILD_ONLY
           : ParallelSegmentConsumptionPolicy.DISALLOW_ALWAYS;
     }
-    if (_realtimeTableDataManager.isPartialUpsertEnabled()) {
-      UpsertConfig upsertConfig = _tableConfig.getUpsertConfig();
-      assert upsertConfig != null;
-      if (upsertConfig.isAllowPartialUpsertConsumptionDuringCommit()) {
+    if (_partitionDedupMetadataManager != null) {
+      DedupContext dedupContext = _partitionDedupMetadataManager.getContext();
+      if (dedupContext.isAllowDedupConsumptionDuringCommit()) {
         return ParallelSegmentConsumptionPolicy.ALLOW_ALWAYS;
       }
       return pauseless ? ParallelSegmentConsumptionPolicy.ALLOW_DURING_BUILD_ONLY

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

@@ -70,7 +70,6 @@
 import org.apache.pinot.segment.spi.ImmutableSegment;
 import org.apache.pinot.segment.spi.IndexSegment;
 import org.apache.pinot.segment.spi.SegmentContext;
-import org.apache.pinot.spi.config.table.DedupConfig;
 import org.apache.pinot.spi.config.table.IndexingConfig;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.config.table.UpsertConfig;
@@ -203,27 +202,22 @@ protected void doInit() {
     IndexLoadingConfig indexLoadingConfig = _indexLoadingConfig;
     TableConfig tableConfig = indexLoadingConfig.getTableConfig();
     assert tableConfig != null;
-    DedupConfig dedupConfig = tableConfig.getDedupConfig();
-    boolean dedupEnabled = dedupConfig != null && dedupConfig.isDedupEnabled();
-    if (dedupEnabled) {
+    if (tableConfig.isDedupEnabled()) {
       Schema schema = indexLoadingConfig.getSchema();
       assert schema != null;
-      List<String> primaryKeyColumns = schema.getPrimaryKeyColumns();
-      Preconditions.checkState(!CollectionUtils.isEmpty(primaryKeyColumns),
-          "Primary key columns must be configured for dedup");
-      _tableDedupMetadataManager = TableDedupMetadataManagerFactory.create(tableConfig, schema, this, _serverMetrics,
-          _instanceDataManagerConfig.getDedupConfig());
+      _tableDedupMetadataManager =
+          TableDedupMetadataManagerFactory.create(_instanceDataManagerConfig.getDedupConfig(), tableConfig, schema,
+              this);
     }
 
-    UpsertConfig upsertConfig = tableConfig.getUpsertConfig();
-    if (upsertConfig != null && upsertConfig.getMode() != UpsertConfig.Mode.NONE) {
-      Preconditions.checkState(!dedupEnabled, "Dedup and upsert cannot be both enabled for table: %s",
-          _tableUpsertMetadataManager);
+    if (tableConfig.isUpsertEnabled()) {
+      Preconditions.checkState(_tableDedupMetadataManager == null,
+          "Dedup and upsert cannot be both enabled for table: %s", _tableNameWithType);
       Schema schema = indexLoadingConfig.getSchema();
       assert schema != null;
       _tableUpsertMetadataManager =
-          TableUpsertMetadataManagerFactory.create(tableConfig, _instanceDataManagerConfig.getUpsertConfig());
-      _tableUpsertMetadataManager.init(tableConfig, schema, this);
+          TableUpsertMetadataManagerFactory.create(_instanceDataManagerConfig.getUpsertConfig(), tableConfig, schema,
+              this);
     }
 
     _enforceConsumptionInOrder = isEnforceConsumptionInOrder();
@@ -423,7 +417,7 @@ public boolean isUpsertEnabled() {
 
   public boolean isPartialUpsertEnabled() {
     return _tableUpsertMetadataManager != null
-        && _tableUpsertMetadataManager.getUpsertMode() == UpsertConfig.Mode.PARTIAL;
+        && _tableUpsertMetadataManager.getContext().getUpsertMode() == UpsertConfig.Mode.PARTIAL;
   }
 
   private void handleSegmentPreload(SegmentZKMetadata zkMetadata, IndexLoadingConfig indexLoadingConfig) {
@@ -439,7 +433,7 @@ private void handleSegmentPreload(SegmentZKMetadata zkMetadata, IndexLoadingConf
    * Handles upsert preload if the upsert preload is enabled.
    */
   private void handleUpsertPreload(SegmentZKMetadata zkMetadata, IndexLoadingConfig indexLoadingConfig) {
-    if (_tableUpsertMetadataManager == null || !_tableUpsertMetadataManager.isEnablePreload()) {
+    if (_tableUpsertMetadataManager == null || !_tableUpsertMetadataManager.getContext().isPreloadEnabled()) {
       return;
     }
     String segmentName = zkMetadata.getSegmentName();
@@ -453,7 +447,7 @@ private void handleUpsertPreload(SegmentZKMetadata zkMetadata, IndexLoadingConfi
    * Handles dedup preload if the dedup preload is enabled.
    */
   private void handleDedupPreload(SegmentZKMetadata zkMetadata, IndexLoadingConfig indexLoadingConfig) {
-    if (_tableDedupMetadataManager == null || !_tableDedupMetadataManager.isEnablePreload()) {
+    if (_tableDedupMetadataManager == null || !_tableDedupMetadataManager.getContext().isPreloadEnabled()) {
       return;
     }
     String segmentName = zkMetadata.getSegmentName();
@@ -654,12 +648,12 @@ private long getDownloadTimeoutMs(TableConfig tableConfig) {
   @VisibleForTesting
   protected RealtimeSegmentDataManager createRealtimeSegmentDataManager(SegmentZKMetadata zkMetadata,
       TableConfig tableConfig, IndexLoadingConfig indexLoadingConfig, Schema schema, LLCSegmentName llcSegmentName,
-      ConsumerCoordinator consumerCoordinator, PartitionUpsertMetadataManager partitionUpsertMetadataManager,
-      PartitionDedupMetadataManager partitionDedupMetadataManager, BooleanSupplier isTableReadyToConsumeData)
+      ConsumerCoordinator consumerCoordinator, @Nullable PartitionUpsertMetadataManager partitionUpsertMetadataManager,
+      @Nullable PartitionDedupMetadataManager partitionDedupMetadataManager, BooleanSupplier isTableReadyToConsumeData)
       throws AttemptsExceededException, RetriableOperationException {
     return new RealtimeSegmentDataManager(zkMetadata, tableConfig, this, _indexDir.getAbsolutePath(),
-        indexLoadingConfig, schema, llcSegmentName, consumerCoordinator, _serverMetrics,
-        partitionUpsertMetadataManager, partitionDedupMetadataManager, isTableReadyToConsumeData);
+        indexLoadingConfig, schema, llcSegmentName, consumerCoordinator, _serverMetrics, partitionUpsertMetadataManager,
+        partitionDedupMetadataManager, isTableReadyToConsumeData);
   }
 
   /**
@@ -784,7 +778,7 @@ private void replaceUpsertSegment(String segmentName, SegmentDataManager oldSegm
     // being filled by partitionUpsertMetadataManager, making the queries see less valid docs than expected.
     IndexSegment oldSegment = oldSegmentManager.getSegment();
     ImmutableSegment immutableSegment = newSegmentManager.getSegment();
-    UpsertConfig.ConsistencyMode consistencyMode = _tableUpsertMetadataManager.getUpsertConsistencyMode();
+    UpsertConfig.ConsistencyMode consistencyMode = _tableUpsertMetadataManager.getContext().getConsistencyMode();
     if (consistencyMode == UpsertConfig.ConsistencyMode.NONE) {
       partitionUpsertMetadataManager.replaceSegment(immutableSegment, oldSegment);
       registerSegment(segmentName, newSegmentManager, partitionUpsertMetadataManager);

pinot-core/src/main/java/org/apache/pinot/core/query/executor/ServerQueryExecutorV1Impl.java

@@ -227,7 +227,7 @@ private InstanceResponseBlock executeInternal(ServerQueryRequest queryRequest, E
       RealtimeTableDataManager rtdm = (RealtimeTableDataManager) tableDataManager;
       TableUpsertMetadataManager tumm = rtdm.getTableUpsertMetadataManager();
       boolean isUsingConsistencyMode =
-          rtdm.getTableUpsertMetadataManager().getUpsertConsistencyMode() != UpsertConfig.ConsistencyMode.NONE;
+          rtdm.getTableUpsertMetadataManager().getContext().getConsistencyMode() != UpsertConfig.ConsistencyMode.NONE;
       if (isUsingConsistencyMode) {
         tumm.lockForSegmentContexts();
       }