Make upsert metadata manager pluggable

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

@@ -57,9 +57,9 @@
 import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig;
 import org.apache.pinot.segment.local.segment.index.loader.LoaderUtils;
 import org.apache.pinot.segment.local.segment.virtualcolumn.VirtualColumnProviderFactory;
-import org.apache.pinot.segment.local.upsert.PartialUpsertHandler;
 import org.apache.pinot.segment.local.upsert.PartitionUpsertMetadataManager;
 import org.apache.pinot.segment.local.upsert.TableUpsertMetadataManager;
+import org.apache.pinot.segment.local.upsert.TableUpsertMetadataManagerFactory;
 import org.apache.pinot.segment.local.utils.SchemaUtils;
 import org.apache.pinot.segment.local.utils.tablestate.TableStateUtils;
 import org.apache.pinot.segment.spi.ImmutableSegment;
@@ -188,26 +188,7 @@ protected void doInit() {
           _tableUpsertMetadataManager);
       Schema schema = ZKMetadataProvider.getTableSchema(_propertyStore, _tableNameWithType);
       Preconditions.checkState(schema != null, "Failed to find schema for table: %s", _tableNameWithType);
-
-      List<String> primaryKeyColumns = schema.getPrimaryKeyColumns();
-      Preconditions.checkState(!CollectionUtils.isEmpty(primaryKeyColumns),
-          "Primary key columns must be configured for upsert");
-
-      String comparisonColumn = upsertConfig.getComparisonColumn();
-      if (comparisonColumn == null) {
-        comparisonColumn = tableConfig.getValidationConfig().getTimeColumnName();
-      }
-
-      PartialUpsertHandler partialUpsertHandler = null;
-      if (upsertConfig.getMode() == UpsertConfig.Mode.PARTIAL) {
-        assert upsertConfig.getPartialUpsertStrategies() != null;
-        partialUpsertHandler = new PartialUpsertHandler(schema, upsertConfig.getPartialUpsertStrategies(),
-            upsertConfig.getDefaultPartialUpsertStrategy(), comparisonColumn);
-      }
-
-      _tableUpsertMetadataManager =
-          new TableUpsertMetadataManager(_tableNameWithType, primaryKeyColumns, comparisonColumn,
-              upsertConfig.getHashFunction(), partialUpsertHandler, _serverMetrics);
+      _tableUpsertMetadataManager = TableUpsertMetadataManagerFactory.create(tableConfig, schema, this, _serverMetrics);
     }
   }
 
@@ -264,7 +245,8 @@ public boolean isUpsertEnabled() {
   }
 
   public boolean isPartialUpsertEnabled() {
-    return _tableUpsertMetadataManager != null && _tableUpsertMetadataManager.isPartialUpsertEnabled();
+    return _tableUpsertMetadataManager != null
+        && _tableUpsertMetadataManager.getUpsertMode() == UpsertConfig.Mode.PARTIAL;
   }
 
   /*

pinot-core/src/test/java/org/apache/pinot/core/plan/maker/MetadataAndDictionaryAggregationPlanMakerTest.java

@@ -38,7 +38,7 @@
 import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentImpl;
 import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
 import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
-import org.apache.pinot.segment.local.upsert.PartitionUpsertMetadataManager;
+import org.apache.pinot.segment.local.upsert.ConcurrentMapPartitionUpsertMetadataManager;
 import org.apache.pinot.segment.spi.IndexSegment;
 import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
 import org.apache.pinot.segment.spi.creator.SegmentIndexCreationDriver;
@@ -125,7 +125,7 @@ public void loadSegment()
     ServerMetrics serverMetrics = Mockito.mock(ServerMetrics.class);
     _upsertIndexSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME), ReadMode.heap);
     ((ImmutableSegmentImpl) _upsertIndexSegment).enableUpsert(
-        new PartitionUpsertMetadataManager("testTable_REALTIME", 0, Collections.singletonList("column6"),
+        new ConcurrentMapPartitionUpsertMetadataManager("testTable_REALTIME", 0, Collections.singletonList("column6"),
             "daysSinceEpoch", HashFunction.NONE, null, serverMetrics), new ThreadSafeMutableRoaringBitmap());
   }
 

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

@@ -62,11 +62,11 @@
 import org.apache.pinot.segment.local.segment.virtualcolumn.VirtualColumnProvider;
 import org.apache.pinot.segment.local.segment.virtualcolumn.VirtualColumnProviderFactory;
 import org.apache.pinot.segment.local.upsert.PartitionUpsertMetadataManager;
+import org.apache.pinot.segment.local.upsert.RecordInfo;
 import org.apache.pinot.segment.local.utils.FixedIntArrayOffHeapIdMap;
 import org.apache.pinot.segment.local.utils.GeometrySerializer;
 import org.apache.pinot.segment.local.utils.IdMap;
 import org.apache.pinot.segment.local.utils.IngestionUtils;
-import org.apache.pinot.segment.local.utils.RecordInfo;
 import org.apache.pinot.segment.local.utils.TableConfigUtils;
 import org.apache.pinot.segment.spi.AggregationFunctionType;
 import org.apache.pinot.segment.spi.MutableSegment;

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BaseTableUpsertMetadataManager.java

@@ -0,0 +1,79 @@
+/**
+ * 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.upsert;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.concurrent.ThreadSafe;
+import org.apache.commons.collections.CollectionUtils;
+import org.apache.pinot.common.metrics.ServerMetrics;
+import org.apache.pinot.segment.local.data.manager.TableDataManager;
+import org.apache.pinot.spi.config.table.HashFunction;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.UpsertConfig;
+import org.apache.pinot.spi.data.Schema;
+
+
+@ThreadSafe
+public abstract class BaseTableUpsertMetadataManager implements TableUpsertMetadataManager {
+  protected String _tableNameWithType;
+  protected List<String> _primaryKeyColumns;
+  protected String _comparisonColumn;
+  protected HashFunction _hashFunction;
+  protected PartialUpsertHandler _partialUpsertHandler;
+  protected ServerMetrics _serverMetrics;
+
+  @Override
+  public void init(TableConfig tableConfig, Schema schema, TableDataManager tableDataManager,
+      ServerMetrics serverMetrics) {
+    _tableNameWithType = tableConfig.getTableName();
+
+    UpsertConfig upsertConfig = tableConfig.getUpsertConfig();
+    Preconditions.checkArgument(upsertConfig != null && upsertConfig.getMode() != UpsertConfig.Mode.NONE,
+        "Upsert must be enabled for table: %s", _tableNameWithType);
+
+    _primaryKeyColumns = schema.getPrimaryKeyColumns();
+    Preconditions.checkArgument(!CollectionUtils.isEmpty(_primaryKeyColumns),
+        "Primary key columns must be configured for upsert enabled table: %s", _tableNameWithType);
+
+    _comparisonColumn = upsertConfig.getComparisonColumn();
+    if (_comparisonColumn == null) {
+      _comparisonColumn = tableConfig.getValidationConfig().getTimeColumnName();
+    }
+
+    _hashFunction = upsertConfig.getHashFunction();
+
+    if (upsertConfig.getMode() == UpsertConfig.Mode.PARTIAL) {
+      Map<String, UpsertConfig.Strategy> partialUpsertStrategies = upsertConfig.getPartialUpsertStrategies();
+      Preconditions.checkArgument(partialUpsertStrategies != null,
+          "Partial-upsert strategies must be configured for partial-upsert enabled table: %s", _tableNameWithType);
+      _partialUpsertHandler =
+          new PartialUpsertHandler(schema, partialUpsertStrategies, upsertConfig.getDefaultPartialUpsertStrategy(),
+              _comparisonColumn);
+    }
+
+    _serverMetrics = serverMetrics;
+  }
+
+  @Override
+  public UpsertConfig.Mode getUpsertMode() {
+    return _partialUpsertHandler == null ? UpsertConfig.Mode.FULL : UpsertConfig.Mode.PARTIAL;
+  }
+}