#13049 Refactored RecordTransformer & merged RecordEnricher

pinot-common/src/main/java/org/apache/pinot/common/utils/config/TableConfigUtils.java

@@ -62,12 +62,11 @@
 
 public class TableConfigUtils {
   private static final Logger LOGGER = LoggerFactory.getLogger(TableConfigUtils.class);
+  private static final String FIELD_MISSING_MESSAGE_TEMPLATE = "Mandatory field '%s' is missing";
 
   private TableConfigUtils() {
   }
 
-  private static final String FIELD_MISSING_MESSAGE_TEMPLATE = "Mandatory field '%s' is missing";
-
   public static TableConfig fromZNRecord(ZNRecord znRecord)
       throws IOException {
     Map<String, String> simpleFields = znRecord.getSimpleFields();
@@ -80,8 +79,8 @@ public static TableConfig fromZNRecord(ZNRecord znRecord)
     Preconditions.checkState(tableType != null, FIELD_MISSING_MESSAGE_TEMPLATE, TableConfig.TABLE_TYPE_KEY);
 
     String validationConfigString = simpleFields.get(TableConfig.VALIDATION_CONFIG_KEY);
-    Preconditions
-        .checkState(validationConfigString != null, FIELD_MISSING_MESSAGE_TEMPLATE, TableConfig.VALIDATION_CONFIG_KEY);
+    Preconditions.checkState(validationConfigString != null, FIELD_MISSING_MESSAGE_TEMPLATE,
+        TableConfig.VALIDATION_CONFIG_KEY);
     SegmentsValidationAndRetentionConfig validationConfig =
         JsonUtils.stringToObject(validationConfigString, SegmentsValidationAndRetentionConfig.class);
 
@@ -90,8 +89,8 @@ public static TableConfig fromZNRecord(ZNRecord znRecord)
     TenantConfig tenantConfig = JsonUtils.stringToObject(tenantConfigString, TenantConfig.class);
 
     String indexingConfigString = simpleFields.get(TableConfig.INDEXING_CONFIG_KEY);
-    Preconditions
-        .checkState(indexingConfigString != null, FIELD_MISSING_MESSAGE_TEMPLATE, TableConfig.INDEXING_CONFIG_KEY);
+    Preconditions.checkState(indexingConfigString != null, FIELD_MISSING_MESSAGE_TEMPLATE,
+        TableConfig.INDEXING_CONFIG_KEY);
     IndexingConfig indexingConfig = JsonUtils.stringToObject(indexingConfigString, IndexingConfig.class);
 
     String customConfigString = simpleFields.get(TableConfig.CUSTOM_CONFIG_KEY);
@@ -180,14 +179,16 @@ public static TableConfig fromZNRecord(ZNRecord znRecord)
     String instancePartitionsMapString = simpleFields.get(TableConfig.INSTANCE_PARTITIONS_MAP_CONFIG_KEY);
     if (instancePartitionsMapString != null) {
       instancePartitionsMap = JsonUtils.stringToObject(instancePartitionsMapString,
-          new TypeReference<Map<InstancePartitionsType, String>>() { });
+          new TypeReference<Map<InstancePartitionsType, String>>() {
+          });
     }
 
     Map<String, SegmentAssignmentConfig> segmentAssignmentConfigMap = null;
     String segmentAssignmentConfigMapString = simpleFields.get(TableConfig.SEGMENT_ASSIGNMENT_CONFIG_MAP_KEY);
     if (segmentAssignmentConfigMapString != null) {
       segmentAssignmentConfigMap = JsonUtils.stringToObject(segmentAssignmentConfigMapString,
-          new TypeReference<Map<String, SegmentAssignmentConfig>>() { });
+          new TypeReference<Map<String, SegmentAssignmentConfig>>() {
+          });
     }
 
     return new TableConfig(tableName, tableType, validationConfig, tenantConfig, indexingConfig, customConfig,
@@ -228,8 +229,8 @@ public static ZNRecord toZNRecord(TableConfig tableConfig)
     }
     Map<String, InstanceAssignmentConfig> instanceAssignmentConfigMap = tableConfig.getInstanceAssignmentConfigMap();
     if (instanceAssignmentConfigMap != null) {
-      simpleFields
-          .put(TableConfig.INSTANCE_ASSIGNMENT_CONFIG_MAP_KEY, JsonUtils.objectToString(instanceAssignmentConfigMap));
+      simpleFields.put(TableConfig.INSTANCE_ASSIGNMENT_CONFIG_MAP_KEY,
+          JsonUtils.objectToString(instanceAssignmentConfigMap));
     }
     List<FieldConfig> fieldConfigList = tableConfig.getFieldConfigList();
     if (fieldConfigList != null) {
@@ -263,11 +264,10 @@ public static ZNRecord toZNRecord(TableConfig tableConfig)
       simpleFields.put(TableConfig.INSTANCE_PARTITIONS_MAP_CONFIG_KEY,
           JsonUtils.objectToString(tableConfig.getInstancePartitionsMap()));
     }
-    Map<String, SegmentAssignmentConfig> segmentAssignmentConfigMap =
-        tableConfig.getSegmentAssignmentConfigMap();
+    Map<String, SegmentAssignmentConfig> segmentAssignmentConfigMap = tableConfig.getSegmentAssignmentConfigMap();
     if (segmentAssignmentConfigMap != null) {
-      simpleFields
-          .put(TableConfig.SEGMENT_ASSIGNMENT_CONFIG_MAP_KEY, JsonUtils.objectToString(segmentAssignmentConfigMap));
+      simpleFields.put(TableConfig.SEGMENT_ASSIGNMENT_CONFIG_MAP_KEY,
+          JsonUtils.objectToString(segmentAssignmentConfigMap));
     }
 
     ZNRecord znRecord = new ZNRecord(tableConfig.getTableName());
@@ -443,8 +443,8 @@ public static boolean hasPreConfiguredInstancePartitions(TableConfig tableConfig
    */
   public static boolean hasPreConfiguredInstancePartitions(TableConfig tableConfig,
       InstancePartitionsType instancePartitionsType) {
-    return hasPreConfiguredInstancePartitions(tableConfig)
-        && tableConfig.getInstancePartitionsMap().containsKey(instancePartitionsType);
+    return hasPreConfiguredInstancePartitions(tableConfig) && tableConfig.getInstancePartitionsMap()
+        .containsKey(instancePartitionsType);
   }
 
   /**

pinot-connectors/pinot-flink-connector/src/main/java/org/apache/pinot/connector/flink/sink/FlinkSegmentWriter.java

@@ -53,7 +53,6 @@
 import org.apache.pinot.spi.ingestion.batch.BatchConfigProperties;
 import org.apache.pinot.spi.ingestion.batch.spec.Constants;
 import org.apache.pinot.spi.ingestion.segment.writer.SegmentWriter;
-import org.apache.pinot.spi.recordenricher.RecordEnricherPipeline;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -78,7 +77,7 @@ public class FlinkSegmentWriter implements SegmentWriter {
   private String _outputDirURI;
   private Schema _schema;
   private Set<String> _fieldsToRead;
-  private RecordEnricherPipeline _recordEnricherPipeline;
+  private RecordTransformer _recordEnricherPipeline;
   private RecordTransformer _recordTransformer;
 
   private File _stagingDir;
@@ -139,7 +138,7 @@ public void init(TableConfig tableConfig, Schema schema, Map<String, String> bat
 
     _schema = schema;
     _fieldsToRead = _schema.getColumnNames();
-    _recordEnricherPipeline = RecordEnricherPipeline.fromTableConfig(_tableConfig);
+    _recordEnricherPipeline = RecordTransformer.fromTableConfig(_tableConfig);
     _recordTransformer = CompositeTransformer.getDefaultTransformer(_tableConfig, _schema);
     _avroSchema = SegmentProcessorAvroUtils.convertPinotSchemaToAvroSchema(_schema);
     _reusableRecord = new GenericData.Record(_avroSchema);

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

@@ -55,6 +55,7 @@
 import org.apache.pinot.segment.local.realtime.converter.ColumnIndicesForRealtimeTable;
 import org.apache.pinot.segment.local.realtime.converter.RealtimeSegmentConverter;
 import org.apache.pinot.segment.local.realtime.impl.RealtimeSegmentConfig;
+import org.apache.pinot.segment.local.recordtransformer.RecordTransformer;
 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;
@@ -76,7 +77,6 @@
 import org.apache.pinot.spi.data.readers.GenericRow;
 import org.apache.pinot.spi.metrics.PinotMeter;
 import org.apache.pinot.spi.plugin.PluginManager;
-import org.apache.pinot.spi.recordenricher.RecordEnricherPipeline;
 import org.apache.pinot.spi.stream.ConsumerPartitionState;
 import org.apache.pinot.spi.stream.LongMsgOffset;
 import org.apache.pinot.spi.stream.MessageBatch;
@@ -275,7 +275,7 @@ public void deleteSegmentFile() {
   private final int _partitionGroupId;
   private final PartitionGroupConsumptionStatus _partitionGroupConsumptionStatus;
   final String _clientId;
-  private final RecordEnricherPipeline _recordEnricherPipeline;
+  private final RecordTransformer _recordEnricherPipeline;
   private final TransformPipeline _transformPipeline;
   private PartitionGroupConsumer _partitionGroupConsumer = null;
   private StreamMetadataProvider _partitionMetadataProvider = null;
@@ -1514,10 +1514,10 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf
     }
 
     try {
-      _recordEnricherPipeline = RecordEnricherPipeline.fromTableConfig(tableConfig);
+      _recordEnricherPipeline = RecordTransformer.fromTableConfig(tableConfig);
     } catch (Exception e) {
       _realtimeTableDataManager.addSegmentError(_segmentNameStr,
-          new SegmentErrorInfo(now(), "Failed to initialize the RecordEnricherPipeline", e));
+          new SegmentErrorInfo(now(), "Failed to initialize the Record Transformer pipeline", e));
       throw e;
     }
     _transformPipeline = new TransformPipeline(tableConfig, schema);

pinot-core/src/main/java/org/apache/pinot/core/segment/processing/framework/SegmentProcessorFramework.java

@@ -43,7 +43,6 @@
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.data.readers.RecordReader;
 import org.apache.pinot.spi.data.readers.RecordReaderFileConfig;
-import org.apache.pinot.spi.recordenricher.RecordEnricherPipeline;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -290,7 +289,7 @@ private List<File> generateSegment(Map<String, GenericRowFileManager> partitionT
           GenericRowFileRecordReader recordReaderForRange = recordReader.getRecordReaderForRange(startRowId, endRowId);
           SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
           driver.init(generatorConfig, new RecordReaderSegmentCreationDataSource(recordReaderForRange),
-              RecordEnricherPipeline.getPassThroughPipeline(),
+              RecordTransformer.getPassThroughPipeline(),
               TransformPipeline.getPassThroughPipeline());
           driver.build();
           outputSegmentDirs.add(driver.getOutputDirectory());

pinot-core/src/main/java/org/apache/pinot/core/segment/processing/mapper/SegmentMapper.java

@@ -48,7 +48,6 @@
 import org.apache.pinot.spi.data.readers.GenericRow;
 import org.apache.pinot.spi.data.readers.RecordReader;
 import org.apache.pinot.spi.data.readers.RecordReaderFileConfig;
-import org.apache.pinot.spi.recordenricher.RecordEnricherPipeline;
 import org.apache.pinot.spi.utils.StringUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -69,7 +68,7 @@ public class SegmentMapper {
   private final List<FieldSpec> _fieldSpecs;
   private final boolean _includeNullFields;
   private final int _numSortFields;
-  private final RecordEnricherPipeline _recordEnricherPipeline;
+  private final RecordTransformer _recordEnricherPipeline;
   private final CompositeTransformer _recordTransformer;
   private final TimeHandler _timeHandler;
   private final Partitioner[] _partitioners;
@@ -94,7 +93,7 @@ public SegmentMapper(List<RecordReaderFileConfig> recordReaderFileConfigs,
     _fieldSpecs = pair.getLeft();
     _numSortFields = pair.getRight();
     _includeNullFields = tableConfig.getIndexingConfig().isNullHandlingEnabled();
-    _recordEnricherPipeline = RecordEnricherPipeline.fromTableConfig(tableConfig);
+    _recordEnricherPipeline = RecordTransformer.fromTableConfig(tableConfig);
     _recordTransformer = CompositeTransformer.composeAllTransformers(_customRecordTransformers, tableConfig, schema);
     _timeHandler = TimeHandlerFactory.getTimeHandler(processorConfig);
     List<PartitionerConfig> partitionerConfigs = processorConfig.getPartitionerConfigs();

pinot-spi/src/main/java/org/apache/pinot/spi/recordenricher/RecordEnricherRegistry.java → pinot-segment-local/src/main/java/org/apache/pinot/plugin/record/enricher/RecordEnricherRegistry.java

@@ -16,45 +16,29 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-package org.apache.pinot.spi.recordenricher;
+package org.apache.pinot.plugin.record.enricher;
 
-import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.ServiceLoader;
-import org.apache.pinot.spi.config.table.ingestion.EnrichmentConfig;
+import org.apache.pinot.segment.local.recordtransformer.RecordTransformer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 
 public class RecordEnricherRegistry {
   private static final Logger LOGGER = LoggerFactory.getLogger(RecordEnricherRegistry.class);
-  private static final Map<String, RecordEnricherFactory> RECORD_ENRICHER_FACTORY_MAP = new HashMap<>();
+  private static final Map<String, RecordTransformer> RECORD_ENRICHER_FACTORY_MAP = new HashMap<>();
 
   private RecordEnricherRegistry() {
   }
 
-  public static void validateEnrichmentConfig(EnrichmentConfig enrichmentConfig,
-      RecordEnricherValidationConfig config) {
-    if (!RECORD_ENRICHER_FACTORY_MAP.containsKey(enrichmentConfig.getEnricherType())) {
-      throw new IllegalArgumentException("No record enricher found for type: " + enrichmentConfig.getEnricherType());
-    }
-
-    RECORD_ENRICHER_FACTORY_MAP.get(enrichmentConfig.getEnricherType())
-        .validateEnrichmentConfig(enrichmentConfig.getProperties(), config);
-  }
-
-  public static RecordEnricher createRecordEnricher(EnrichmentConfig enrichmentConfig)
-      throws IOException {
-    if (!RECORD_ENRICHER_FACTORY_MAP.containsKey(enrichmentConfig.getEnricherType())) {
-      throw new IllegalArgumentException("No record enricher found for type: " + enrichmentConfig.getEnricherType());
-    }
-    return RECORD_ENRICHER_FACTORY_MAP.get(enrichmentConfig.getEnricherType())
-        .createEnricher(enrichmentConfig.getProperties());
+  public static Map<String, RecordTransformer> getRecordEnricherFactoryMap() {
+    return RECORD_ENRICHER_FACTORY_MAP;
   }
 
   static {
-    for (RecordEnricherFactory recordEnricherFactory : ServiceLoader.load(RecordEnricherFactory.class)) {
+    for (RecordTransformer recordEnricherFactory : ServiceLoader.load(RecordTransformer.class)) {
       LOGGER.info("Registered record enricher factory type: {}", recordEnricherFactory.getEnricherType());
       RECORD_ENRICHER_FACTORY_MAP.put(recordEnricherFactory.getEnricherType(), recordEnricherFactory);
     }

pinot-segment-local/src/main/java/org/apache/pinot/plugin/record/enricher/clp/CLPEncodingEnricher.java

@@ -24,8 +24,8 @@
 import com.yscope.clp.compressorfrontend.MessageEncoder;
 import java.io.IOException;
 import java.util.List;
+import org.apache.pinot.segment.local.recordtransformer.RecordTransformer;
 import org.apache.pinot.spi.data.readers.GenericRow;
-import org.apache.pinot.spi.recordenricher.RecordEnricher;
 import org.apache.pinot.spi.utils.JsonUtils;
 import org.apache.pinot.sql.parsers.rewriter.ClpRewriter;
 import org.slf4j.Logger;
@@ -39,11 +39,12 @@
  * 2. 'x_dictVars' - The dictionary variables of the encoded message
  * 3. 'x_encodedVars' - The encoded variables of the encoded message
  */
-public class CLPEncodingEnricher implements RecordEnricher {
+public class CLPEncodingEnricher implements RecordTransformer {
   private static final Logger LOGGER = LoggerFactory.getLogger(CLPEncodingEnricher.class);
   private final ClpEnricherConfig _config;
   private final EncodedMessage _clpEncodedMessage;
   private final MessageEncoder _clpMessageEncoder;
+  private static final String ENRICHER_TYPE = "clpEnricher";
 
   public CLPEncodingEnricher(JsonNode enricherProperties) throws IOException {
     _config = JsonUtils.jsonNodeToObject(enricherProperties, ClpEnricherConfig.class);
@@ -58,7 +59,7 @@ public List<String> getInputColumns() {
   }
 
   @Override
-  public void enrich(GenericRow record) {
+  public GenericRow transform(GenericRow record) {
     try {
       for (String field : _config.getFields()) {
         Object value = record.getValue(field);
@@ -69,6 +70,7 @@ public void enrich(GenericRow record) {
     } catch (Exception e) {
       LOGGER.error("Failed to enrich record: {}", record);
     }
+    return record;
   }
 
   private void enrichWithClpEncodedFields(String key, Object value, GenericRow to) {
@@ -97,4 +99,24 @@ private void enrichWithClpEncodedFields(String key, Object value, GenericRow to)
     to.putValue(key + ClpRewriter.DICTIONARY_VARS_COLUMN_SUFFIX, dictVars);
     to.putValue(key + ClpRewriter.ENCODED_VARS_COLUMN_SUFFIX, encodedVars);
   }
+
+  @Override
+  public String getEnricherType() {
+    return ENRICHER_TYPE;
+  }
+
+  @Override
+  public RecordTransformer createEnricher(JsonNode enricherProps)
+      throws IOException {
+    return new CLPEncodingEnricher(enricherProps);
+  }
+
+  @Override
+  public void validateEnrichmentConfig(JsonNode enricherProps, boolean validationConfig) {
+    try {
+      ClpEnricherConfig config = JsonUtils.jsonNodeToObject(enricherProps, ClpEnricherConfig.class);
+    } catch (IOException e) {
+      throw new IllegalArgumentException("Failed to parse clp enricher config", e);
+    }
+  }
 }