diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java index a978221ef696..80592de5af31 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java @@ -79,6 +79,8 @@ import org.apache.pinot.spi.config.table.SegmentZKPropsConfig; import org.apache.pinot.spi.config.table.StarTreeIndexConfig; import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.ingestion.IngestionConfig; +import org.apache.pinot.spi.config.table.ingestion.TransformConfig; import org.apache.pinot.spi.data.ComplexFieldSpec; import org.apache.pinot.spi.data.DateTimeFieldSpec; import org.apache.pinot.spi.data.DateTimeFormatSpec; @@ -549,7 +551,8 @@ protected void writeMetadata() ForwardIndexConfig fwdConfig = columnIndexCreators.getIndexConfigs().getConfig(StandardIndexes.forward()); addColumnMetadataInfo(properties, column, columnStatistics, _totalDocs, _schema.getFieldSpecFor(column), - hasDictionary, dictionaryElementSize, fwdConfig.getEncodingType(), false); + hasDictionary, dictionaryElementSize, fwdConfig.getEncodingType(), false, + getTransformFunctionForColumn(column)); } SegmentZKPropsConfig segmentZKPropsConfig = _config.getSegmentZKPropsConfig(); @@ -564,6 +567,15 @@ protected void writeMetadata() public static void addColumnMetadataInfo(PropertiesConfiguration properties, String column, ColumnStatistics columnStatistics, int totalDocs, FieldSpec fieldSpec, boolean hasDictionary, int dictionaryElementSize, FieldConfig.EncodingType forwardIndexEncoding, boolean autoGenerated) { + addColumnMetadataInfo(properties, column, columnStatistics, totalDocs, fieldSpec, hasDictionary, + dictionaryElementSize, forwardIndexEncoding, autoGenerated, null); + } + + /// Adds column metadata information to the properties configuration. + public static void addColumnMetadataInfo(PropertiesConfiguration properties, String column, + ColumnStatistics columnStatistics, int totalDocs, FieldSpec fieldSpec, boolean hasDictionary, + int dictionaryElementSize, FieldConfig.EncodingType forwardIndexEncoding, boolean autoGenerated, + @Nullable String transformFunction) { addFieldSpec(properties, column, fieldSpec); properties.setProperty(getKeyFor(column, TOTAL_DOCS), String.valueOf(totalDocs)); int cardinality = columnStatistics.getCardinality(); @@ -606,6 +618,7 @@ public static void addColumnMetadataInfo(PropertiesConfiguration properties, Str if (autoGenerated) { properties.setProperty(getKeyFor(column, IS_AUTO_GENERATED), "true"); } + addTransformFunction(properties, column, transformFunction); // Min/max value if (fieldSpec.getFieldType() != FieldType.COMPLEX) { @@ -639,6 +652,34 @@ public static void addColumnMetadataInfo(PropertiesConfiguration properties, Str } } + private static void addTransformFunction(PropertiesConfiguration properties, String column, + @Nullable String transformFunction) { + if (transformFunction != null) { + String validTransformFunction = + CommonsConfigurationUtils.replaceSpecialCharacterInPropertyValue(transformFunction); + if (validTransformFunction != null) { + properties.setProperty(getKeyFor(column, TRANSFORM_FUNCTION), validTransformFunction); + } + } + } + + @Nullable + @SuppressWarnings("deprecation") + private String getTransformFunctionForColumn(String column) { + TableConfig tableConfig = _config.getTableConfig(); + IngestionConfig ingestionConfig = tableConfig != null ? tableConfig.getIngestionConfig() : null; + if (ingestionConfig != null && ingestionConfig.getTransformConfigs() != null) { + for (TransformConfig transformConfig : ingestionConfig.getTransformConfigs()) { + if (column.equals(transformConfig.getColumnName())) { + return transformConfig.getTransformFunction(); + } + } + } + FieldSpec fieldSpec = _schema.getFieldSpecFor(column); + // Keep the schema-level transform fallback for segments built from legacy configs. + return fieldSpec != null ? fieldSpec.getTransformFunction() : null; + } + /** * In order to persist complex field metadata, we need to recursively add child field specs * So, each complex field spec will have a property for its child field names and each child field will have its diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java index d88b3b90c60b..1480dccff4f5 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java @@ -30,6 +30,8 @@ import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.Objects; +import javax.annotation.Nullable; import org.apache.commons.configuration2.PropertiesConfiguration; import org.apache.pinot.common.evaluator.FunctionEvaluatorFactory; import org.apache.pinot.common.function.FunctionUtils; @@ -110,13 +112,17 @@ protected enum DefaultColumnAction { UPDATE_DIMENSION_DATA_TYPE, UPDATE_DIMENSION_DEFAULT_VALUE, UPDATE_DIMENSION_NUMBER_OF_VALUES, + UPDATE_DIMENSION_TRANSFORM_FUNCTION, UPDATE_METRIC_DATA_TYPE, UPDATE_METRIC_DEFAULT_VALUE, UPDATE_METRIC_NUMBER_OF_VALUES, + UPDATE_METRIC_TRANSFORM_FUNCTION, UPDATE_DATE_TIME_DATA_TYPE, UPDATE_DATE_TIME_DEFAULT_VALUE, + UPDATE_DATE_TIME_TRANSFORM_FUNCTION, UPDATE_COMPLEX_DATA_TYPE, - UPDATE_COMPLEX_DEFAULT_VALUE; + UPDATE_COMPLEX_DEFAULT_VALUE, + UPDATE_COMPLEX_TRANSFORM_FUNCTION; boolean isAddAction() { return this == ADD_DIMENSION || this == ADD_METRIC || this == ADD_DATE_TIME || this == ADD_COMPLEX; @@ -293,6 +299,8 @@ Map computeDefaultColumnActionMap() { defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_DIMENSION_DEFAULT_VALUE); } else if (isSingleValueInMetadata != isSingleValueInSchema) { defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_DIMENSION_NUMBER_OF_VALUES); + } else if (isTransformFunctionChanged(column, columnMetadata)) { + defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_DIMENSION_TRANSFORM_FUNCTION); } } else if (fieldTypeInMetadata == METRIC) { if (dataTypeInMetadata != dataTypeInSchema) { @@ -301,18 +309,24 @@ Map computeDefaultColumnActionMap() { defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_METRIC_DEFAULT_VALUE); } else if (isSingleValueInMetadata != isSingleValueInSchema) { defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_METRIC_NUMBER_OF_VALUES); + } else if (isTransformFunctionChanged(column, columnMetadata)) { + defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_METRIC_TRANSFORM_FUNCTION); } } else if (fieldTypeInMetadata == DATE_TIME) { if (dataTypeInMetadata != dataTypeInSchema) { defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_DATE_TIME_DATA_TYPE); } else if (!defaultValueInSchema.equals(defaultValueInMetadata)) { defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_DATE_TIME_DEFAULT_VALUE); + } else if (isTransformFunctionChanged(column, columnMetadata)) { + defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_DATE_TIME_TRANSFORM_FUNCTION); } } else if (fieldTypeInMetadata == COMPLEX) { if (dataTypeInMetadata != dataTypeInSchema) { defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_COMPLEX_DATA_TYPE); } else if (!defaultValueInSchema.equals(defaultValueInMetadata)) { defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_COMPLEX_DEFAULT_VALUE); + } else if (isTransformFunctionChanged(column, columnMetadata)) { + defaultColumnActionMap.put(column, DefaultColumnAction.UPDATE_COMPLEX_TRANSFORM_FUNCTION); } } } else { @@ -359,6 +373,11 @@ Map computeDefaultColumnActionMap() { return defaultColumnActionMap; } + private boolean isTransformFunctionChanged(String column, ColumnMetadata columnMetadata) { + String transformFunctionInMetadata = columnMetadata.getTransformFunction(); + return !Objects.equals(transformFunctionInMetadata, getTransformFunctionForColumn(column)); + } + /** * Helper method to update default column indices, returns {@code true} if the update succeeds, {@code false} * otherwise. @@ -390,65 +409,74 @@ protected void removeColumnIndices(String column) { protected boolean createColumnV1Indices(String column) throws Exception { boolean errorOnFailure = _indexLoadingConfig.isErrorOnColumnBuildFailure(); - IngestionConfig ingestionConfig = _tableConfig.getIngestionConfig(); - if (ingestionConfig != null && ingestionConfig.getTransformConfigs() != null) { - List transformConfigs = ingestionConfig.getTransformConfigs(); - for (TransformConfig transformConfig : transformConfigs) { - if (transformConfig.getColumnName().equals(column)) { - String transformFunction = transformConfig.getTransformFunction(); - FunctionEvaluator functionEvaluator = FunctionEvaluatorFactory.getExpressionEvaluator(transformFunction); - - // Check if all arguments exist in the segment - // TODO: Support chained derived column - List arguments = functionEvaluator.getArguments(); - List argumentsMetadata = new ArrayList<>(arguments.size()); - for (String argument : arguments) { - ColumnMetadata columnMetadata = _segmentMetadata.getColumnMetadataFor(argument); - if (columnMetadata == null) { - LOGGER.warn("Assigning default value to derived column: {} because argument: {} does not exist in the " - + "segment", column, argument); - createDefaultValueColumnV1Indices(column); - return true; - } - // TODO: Support creation of derived columns from forward index disabled columns - if (!_segmentWriter.hasIndexFor(argument, StandardIndexes.forward())) { - throw new UnsupportedOperationException(String.format("Operation not supported! Cannot create a derived " - + "column %s because argument: %s does not have a forward index. Enable forward index and " - + "refresh/backfill the segments to create a derived column from source column", column, - argument)); - } - argumentsMetadata.add(columnMetadata); - } + String transformFunction = getTransformFunctionForColumn(column); + if (transformFunction != null) { + FunctionEvaluator functionEvaluator = FunctionEvaluatorFactory.getExpressionEvaluator(transformFunction); + + // Check if all arguments exist in the segment + // TODO: Support chained derived column + List arguments = functionEvaluator.getArguments(); + List argumentsMetadata = new ArrayList<>(arguments.size()); + for (String argument : arguments) { + ColumnMetadata columnMetadata = _segmentMetadata.getColumnMetadataFor(argument); + if (columnMetadata == null) { + LOGGER.warn("Assigning default value to derived column: {} because argument: {} does not exist in the " + + "segment", column, argument); + createDefaultValueColumnV1Indices(column, transformFunction); + return true; + } + // TODO: Support creation of derived columns from forward index disabled columns + if (!_segmentWriter.hasIndexFor(argument, StandardIndexes.forward())) { + throw new UnsupportedOperationException(String.format("Operation not supported! Cannot create a derived " + + "column %s because argument: %s does not have a forward index. Enable forward index and " + + "refresh/backfill the segments to create a derived column from source column", column, + argument)); + } + argumentsMetadata.add(columnMetadata); + } - // TODO: Support forward index disabled derived column - if (isForwardIndexDisabled(column)) { - LOGGER.warn("Skip creating forward index disabled derived column: {}", column); - if (errorOnFailure) { - throw new UnsupportedOperationException( - String.format("Failed to create forward index disabled derived column: %s", column)); - } - return false; - } + // TODO: Support forward index disabled derived column + if (isForwardIndexDisabled(column)) { + LOGGER.warn("Skip creating forward index disabled derived column: {}", column); + if (errorOnFailure) { + throw new UnsupportedOperationException( + String.format("Failed to create forward index disabled derived column: %s", column)); + } + return false; + } - try { - createDerivedColumnV1Indices(column, functionEvaluator, argumentsMetadata, errorOnFailure); - return true; - } catch (Exception e) { - LOGGER.error("Caught exception while creating derived column: {} with transform function: {}", column, - transformFunction, e); - if (errorOnFailure) { - throw e; - } - return false; - } + try { + createDerivedColumnV1Indices(column, transformFunction, functionEvaluator, argumentsMetadata, errorOnFailure); + return true; + } catch (Exception e) { + LOGGER.error("Caught exception while creating derived column: {} with transform function: {}", column, + transformFunction, e); + if (errorOnFailure) { + throw e; } + return false; } } - createDefaultValueColumnV1Indices(column); + createDefaultValueColumnV1Indices(column, null); return true; } + @SuppressWarnings("deprecation") + private String getTransformFunctionForColumn(String column) { + IngestionConfig ingestionConfig = _tableConfig.getIngestionConfig(); + if (ingestionConfig != null && ingestionConfig.getTransformConfigs() != null) { + for (TransformConfig transformConfig : ingestionConfig.getTransformConfigs()) { + if (column.equals(transformConfig.getColumnName())) { + return transformConfig.getTransformFunction(); + } + } + } + FieldSpec fieldSpec = _schema.getFieldSpecFor(column); + // Keep the schema-level transform fallback for legacy configs. + return fieldSpec != null ? fieldSpec.getTransformFunction() : null; + } + /** * Check and return whether the forward index is disabled for a given column */ @@ -460,7 +488,7 @@ protected boolean isForwardIndexDisabled(String column) { /** * Helper method to create the V1 indices (dictionary and forward index) for a column with default values. */ - private void createDefaultValueColumnV1Indices(String column) + private void createDefaultValueColumnV1Indices(String column, @Nullable String transformFunction) throws Exception { FieldSpec fieldSpec = _schema.getFieldSpecFor(column); @@ -528,7 +556,7 @@ private void createDefaultValueColumnV1Indices(String column) // Add the column metadata information to the metadata properties. BaseSegmentCreator.addColumnMetadataInfo(_segmentProperties, column, columnStatistics, totalDocs, fieldSpec, true, - dictionaryElementSize, FieldConfig.EncodingType.DICTIONARY, true); + dictionaryElementSize, FieldConfig.EncodingType.DICTIONARY, true, transformFunction); } private boolean isNullable(FieldSpec fieldSpec) { @@ -545,8 +573,8 @@ private boolean isNullable(FieldSpec fieldSpec) { * - Support chained derived column * - Support forward index disabled derived column */ - private void createDerivedColumnV1Indices(String column, FunctionEvaluator functionEvaluator, - List argumentsMetadata, boolean errorOnFailure) + private void createDerivedColumnV1Indices(String column, String transformFunction, + FunctionEvaluator functionEvaluator, List argumentsMetadata, boolean errorOnFailure) throws Exception { // Initialize value readers for all arguments int numArguments = argumentsMetadata.size(); @@ -766,10 +794,11 @@ private void createDerivedColumnV1Indices(String column, FunctionEvaluator funct } if (createDictionary) { - createDerivedColumnForwardIndexWithDictionary(column, fieldSpec, outputValues, statsCollector, - useVarLengthDictionary); + createDerivedColumnForwardIndexWithDictionary(column, transformFunction, fieldSpec, outputValues, + statsCollector, useVarLengthDictionary); } else { - createDerivedColumnForwardIndexWithoutDictionary(column, fieldSpec, outputValues, statsCollector); + createDerivedColumnForwardIndexWithoutDictionary(column, transformFunction, fieldSpec, outputValues, + statsCollector); } } finally { for (ValueReader valueReader : valueReaders) { @@ -1083,8 +1112,8 @@ private Object getBytesOutputValue(Object outputValue, boolean isSingleValue, Pi /** * Helper method to create the dictionary and forward indices for a column with derived values. */ - private void createDerivedColumnForwardIndexWithDictionary(String column, FieldSpec fieldSpec, Object[] outputValues, - ColumnStatistics columnStatistics, boolean useVarLengthDictionary) + private void createDerivedColumnForwardIndexWithDictionary(String column, String transformFunction, + FieldSpec fieldSpec, Object[] outputValues, ColumnStatistics columnStatistics, boolean useVarLengthDictionary) throws Exception { // Create dictionary @@ -1111,7 +1140,7 @@ private void createDerivedColumnForwardIndexWithDictionary(String column, FieldS } // Add the column metadata BaseSegmentCreator.addColumnMetadataInfo(_segmentProperties, column, columnStatistics, numDocs, fieldSpec, true, - dictionaryCreator.getNumBytesPerEntry(), FieldConfig.EncodingType.DICTIONARY, true); + dictionaryCreator.getNumBytesPerEntry(), FieldConfig.EncodingType.DICTIONARY, true, transformFunction); } } } @@ -1119,8 +1148,8 @@ private void createDerivedColumnForwardIndexWithDictionary(String column, FieldS /** * Helper method to create a forward index for a raw encoded column with derived values. */ - private void createDerivedColumnForwardIndexWithoutDictionary(String column, FieldSpec fieldSpec, - Object[] outputValues, ColumnStatistics columnStatistics) + private void createDerivedColumnForwardIndexWithoutDictionary(String column, String transformFunction, + FieldSpec fieldSpec, Object[] outputValues, ColumnStatistics columnStatistics) throws Exception { // Create forward index @@ -1196,7 +1225,7 @@ private void createDerivedColumnForwardIndexWithoutDictionary(String column, Fie // Add the column metadata BaseSegmentCreator.addColumnMetadataInfo(_segmentProperties, column, columnStatistics, numDocs, fieldSpec, false, - 0, FieldConfig.EncodingType.RAW, true); + 0, FieldConfig.EncodingType.RAW, true, transformFunction); } private ForwardIndexCreator getForwardIndexCreator(ColumnStatistics columnStatistics, diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/map/SimpleColumnMetadata.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/map/SimpleColumnMetadata.java index d62d1fd305e0..00b3e9deb84e 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/map/SimpleColumnMetadata.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/map/SimpleColumnMetadata.java @@ -132,6 +132,12 @@ public boolean isAutoGenerated() { return false; } + @Nullable + @Override + public String getTransformFunction() { + return null; + } + @Override public int getNumIndexes() { return 0; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/DefaultColumnHandlerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/DefaultColumnHandlerTest.java index 78a64eef7f49..e1950c8e67bf 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/DefaultColumnHandlerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/DefaultColumnHandlerTest.java @@ -20,17 +20,26 @@ import java.io.File; import java.net.URL; +import java.util.List; import java.util.Map; +import java.util.TreeMap; +import org.apache.commons.configuration2.PropertiesConfiguration; import org.apache.commons.io.FileUtils; import org.apache.pinot.segment.local.segment.creator.SegmentTestUtils; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; import org.apache.pinot.segment.local.segment.index.loader.defaultcolumn.BaseDefaultColumnHandler.DefaultColumnAction; import org.apache.pinot.segment.local.segment.store.SegmentLocalFSDirectory; +import org.apache.pinot.segment.spi.ColumnMetadata; +import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.apache.pinot.segment.spi.store.SegmentDirectory; +import org.apache.pinot.segment.spi.utils.SegmentMetadataUtils; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.config.table.ingestion.IngestionConfig; +import org.apache.pinot.spi.config.table.ingestion.TransformConfig; import org.apache.pinot.spi.data.DateTimeFieldSpec; import org.apache.pinot.spi.data.DimensionFieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; @@ -42,8 +51,11 @@ import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; public class DefaultColumnHandlerTest { @@ -53,6 +65,9 @@ public class DefaultColumnHandlerTest { new File(FileUtils.getTempDirectory(), DefaultColumnHandlerTest.class.getSimpleName()); private static final File INDEX_DIR = new File(TEMP_DIR, SEGMENT_NAME); private static final String AVRO_DATA = "data/test_data-mv.avro"; + private static final String DERIVED_COLUMN = "derivedColumn"; + private static final String ORIGINAL_TRANSFORM_FUNCTION = "plus(column1, 1)"; + private static final String UPDATED_TRANSFORM_FUNCTION = "plus(column1, 2)"; private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build(); @@ -125,10 +140,171 @@ public void testComputeDefaultColumnActionMap() } } + @Test + public void testSegmentGenerationPersistsTransformFunction() + throws Exception { + File segmentGenerationTempDir = new File(TEMP_DIR, "segmentGenerationWithTransform"); + File indexDir = buildSegmentWithDerivedColumn(segmentGenerationTempDir, ORIGINAL_TRANSFORM_FUNCTION); + + SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); + assertEquals(segmentMetadata.getColumnMetadataFor(DERIVED_COLUMN).getTransformFunction(), + ORIGINAL_TRANSFORM_FUNCTION); + } + + @Test + public void testLegacyDerivedColumnWithoutStoredTransformFunctionIsRegenerated() + throws Exception { + File segmentGenerationTempDir = new File(TEMP_DIR, "legacyDerivedTransform"); + File indexDir = + buildSegment(segmentGenerationTempDir, TABLE_CONFIG, Schema.fromString(_schema.toSingleLineJsonString())); + Schema schema = getSchemaWithDerivedColumn(); + + try (SegmentDirectory segmentDirectory = new SegmentLocalFSDirectory(indexDir, ReadMode.mmap); + SegmentDirectory.Writer writer = segmentDirectory.createWriter()) { + DefaultColumnHandler defaultColumnHandler = + new V3DefaultColumnHandler(indexDir, segmentDirectory.getSegmentMetadata(), + new IndexLoadingConfig(getTableConfigWithTransformFunction(ORIGINAL_TRANSFORM_FUNCTION), schema), writer); + defaultColumnHandler.updateDefaultColumns(); + } + + SegmentMetadataImpl originalSegmentMetadata = new SegmentMetadataImpl(indexDir); + assertEquals(originalSegmentMetadata.getColumnMetadataFor(DERIVED_COLUMN).getTransformFunction(), + ORIGINAL_TRANSFORM_FUNCTION); + removeTransformFunctionFromMetadata(indexDir, DERIVED_COLUMN); + + SegmentMetadataImpl legacySegmentMetadata = new SegmentMetadataImpl(indexDir); + assertNull(legacySegmentMetadata.getColumnMetadataFor(DERIVED_COLUMN).getTransformFunction()); + ColumnMetadata sourceColumnMetadata = legacySegmentMetadata.getColumnMetadataFor("column1"); + + try (SegmentDirectory segmentDirectory = new SegmentLocalFSDirectory(indexDir, ReadMode.mmap); + SegmentDirectory.Writer writer = segmentDirectory.createWriter()) { + DefaultColumnHandler defaultColumnHandler = + new V3DefaultColumnHandler(indexDir, segmentDirectory.getSegmentMetadata(), + new IndexLoadingConfig(getTableConfigWithTransformFunction(UPDATED_TRANSFORM_FUNCTION), schema), writer); + defaultColumnHandler.updateDefaultColumns(); + } + + SegmentMetadataImpl updatedSegmentMetadata = new SegmentMetadataImpl(indexDir); + ColumnMetadata derivedColumnMetadata = updatedSegmentMetadata.getColumnMetadataFor(DERIVED_COLUMN); + assertEquals(derivedColumnMetadata.getTransformFunction(), UPDATED_TRANSFORM_FUNCTION); + assertEquals(derivedColumnMetadata.getMinValue(), (Integer) sourceColumnMetadata.getMinValue() + 2); + assertEquals(derivedColumnMetadata.getMaxValue(), (Integer) sourceColumnMetadata.getMaxValue() + 2); + } + + @Test + public void testMissingDerivedColumnArgumentPersistsTransformFunction() + throws Exception { + File segmentGenerationTempDir = new File(TEMP_DIR, "missingArgumentDerivedTransform"); + FileUtils.deleteQuietly(segmentGenerationTempDir); + File indexDir = + buildSegment(segmentGenerationTempDir, TABLE_CONFIG, Schema.fromString(_schema.toSingleLineJsonString())); + String transformFunction = "plus(missingColumn, 1)"; + + Schema schema = getSchemaWithDerivedColumn(); + try (SegmentDirectory segmentDirectory = new SegmentLocalFSDirectory(indexDir, ReadMode.mmap); + SegmentDirectory.Writer writer = segmentDirectory.createWriter()) { + DefaultColumnHandler defaultColumnHandler = + new V3DefaultColumnHandler(indexDir, segmentDirectory.getSegmentMetadata(), + new IndexLoadingConfig(getTableConfigWithTransformFunction(transformFunction), schema), writer); + defaultColumnHandler.updateDefaultColumns(); + } + + SegmentMetadataImpl updatedSegmentMetadata = new SegmentMetadataImpl(indexDir); + assertEquals(updatedSegmentMetadata.getColumnMetadataFor(DERIVED_COLUMN).getTransformFunction(), + transformFunction); + } + + @Test + public void testComputeDefaultColumnActionMapForTransformFunction() { + assertEquals(computeDefaultColumnActionMap(ORIGINAL_TRANSFORM_FUNCTION, UPDATED_TRANSFORM_FUNCTION), + Map.of(DERIVED_COLUMN, DefaultColumnAction.UPDATE_DIMENSION_TRANSFORM_FUNCTION)); + assertEquals(computeDefaultColumnActionMap(ORIGINAL_TRANSFORM_FUNCTION, ORIGINAL_TRANSFORM_FUNCTION), Map.of()); + assertEquals(computeDefaultColumnActionMap(null, UPDATED_TRANSFORM_FUNCTION), + Map.of(DERIVED_COLUMN, DefaultColumnAction.UPDATE_DIMENSION_TRANSFORM_FUNCTION)); + assertEquals(computeDefaultColumnActionMap(ORIGINAL_TRANSFORM_FUNCTION, null), + Map.of(DERIVED_COLUMN, DefaultColumnAction.UPDATE_DIMENSION_TRANSFORM_FUNCTION)); + assertEquals(computeDefaultColumnActionMap(null, null), Map.of()); + } + private void testComputeDefaultColumnActionMap(Map expected) { BaseDefaultColumnHandler defaultColumnHandler = new V3DefaultColumnHandler(INDEX_DIR, _segmentDirectory.getSegmentMetadata(), new IndexLoadingConfig(TABLE_CONFIG, _schema), _writer); assertEquals(defaultColumnHandler.computeDefaultColumnActionMap(), expected); } + + private static Map computeDefaultColumnActionMap(String transformFunctionInMetadata, + String transformFunctionInTableConfig) { + Schema schema = new Schema.SchemaBuilder() + .setSchemaName(RAW_TABLE_NAME) + .addSingleValueDimension(DERIVED_COLUMN, DataType.INT) + .build(); + ColumnMetadata columnMetadata = mock(ColumnMetadata.class); + when(columnMetadata.getColumnName()).thenReturn(DERIVED_COLUMN); + when(columnMetadata.isAutoGenerated()).thenReturn(true); + when(columnMetadata.getFieldSpec()).thenReturn(schema.getFieldSpecFor(DERIVED_COLUMN)); + when(columnMetadata.getTransformFunction()).thenReturn(transformFunctionInMetadata); + TreeMap columnMetadataMap = new TreeMap<>(); + columnMetadataMap.put(DERIVED_COLUMN, columnMetadata); + + SegmentMetadataImpl segmentMetadata = mock(SegmentMetadataImpl.class); + when(segmentMetadata.getColumnMetadataFor(DERIVED_COLUMN)).thenReturn(columnMetadata); + when(segmentMetadata.getColumnMetadataMap()).thenReturn(columnMetadataMap); + BaseDefaultColumnHandler defaultColumnHandler = + new V3DefaultColumnHandler(new File("."), segmentMetadata, + new IndexLoadingConfig(getTableConfigWithTransformFunction(transformFunctionInTableConfig), schema), + mock(SegmentDirectory.Writer.class)); + return defaultColumnHandler.computeDefaultColumnActionMap(); + } + + private static TableConfig getTableConfigWithTransformFunction(String transformFunction) { + IngestionConfig ingestionConfig = new IngestionConfig(); + if (transformFunction != null) { + ingestionConfig.setTransformConfigs(List.of(new TransformConfig(DERIVED_COLUMN, transformFunction))); + } else { + ingestionConfig.setTransformConfigs(List.of()); + } + return new TableConfigBuilder(TableType.OFFLINE) + .setTableName(RAW_TABLE_NAME) + .setIngestionConfig(ingestionConfig) + .build(); + } + + private Schema getSchemaWithDerivedColumn() + throws Exception { + Schema schema = Schema.fromString(_schema.toSingleLineJsonString()); + schema.addField(new DimensionFieldSpec(DERIVED_COLUMN, DataType.INT, true)); + return schema; + } + + private File buildSegmentWithDerivedColumn(File segmentGenerationTempDir, String transformFunction) + throws Exception { + return buildSegment(segmentGenerationTempDir, getTableConfigWithTransformFunction(transformFunction), + getSchemaWithDerivedColumn()); + } + + private File buildSegment(File segmentGenerationTempDir, TableConfig tableConfig, Schema schema) + throws Exception { + FileUtils.deleteQuietly(segmentGenerationTempDir); + + URL resourceUrl = getClass().getClassLoader().getResource(AVRO_DATA); + assertNotNull(resourceUrl); + File avroFile = new File(resourceUrl.getFile()); + SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig, schema); + config.setInputFilePath(avroFile.getAbsolutePath()); + config.setOutDir(segmentGenerationTempDir.getAbsolutePath()); + config.setSegmentName(SEGMENT_NAME); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(config); + driver.build(); + return new File(segmentGenerationTempDir, SEGMENT_NAME); + } + + private static void removeTransformFunctionFromMetadata(File indexDir, String column) + throws Exception { + PropertiesConfiguration segmentProperties = SegmentMetadataUtils.getPropertiesConfiguration(indexDir); + segmentProperties.clearProperty( + V1Constants.MetadataKeys.Column.getKeyFor(column, V1Constants.MetadataKeys.Column.TRANSFORM_FUNCTION)); + SegmentMetadataUtils.savePropertiesConfiguration(segmentProperties, indexDir); + } } diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/ColumnMetadata.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/ColumnMetadata.java index f493db66dc5e..ba28fb9c65ba 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/ColumnMetadata.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/ColumnMetadata.java @@ -20,6 +20,7 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonProperty; +import javax.annotation.Nullable; import org.apache.pinot.segment.spi.index.IndexType; import org.apache.pinot.spi.annotations.InterfaceAudience; import org.apache.pinot.spi.config.table.FieldConfig.EncodingType; @@ -46,6 +47,12 @@ public interface ColumnMetadata extends ColumnShape { /// Returns `true` when the column is auto-generated by the default column handler. boolean isAutoGenerated(); + /// Returns the transform function expression used to generate the column, if persisted in the segment metadata. + @Nullable + default String getTransformFunction() { + return null; + } + /// Returns the total number of indexes for the column. @JsonIgnore int getNumIndexes(); diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/V1Constants.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/V1Constants.java index d611f2d26207..c1ca50e17913 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/V1Constants.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/V1Constants.java @@ -190,6 +190,8 @@ public static class Column { public static final String MAX_ROW_LENGTH_IN_BYTES = "maxRowLengthInBytes"; // Optional, default false public static final String IS_AUTO_GENERATED = "isAutoGenerated"; + // Optional. Transform function expression used to generate the column. + public static final String TRANSFORM_FUNCTION = "transformFunction"; // Optional. The parent column name for columns materialized from an OPEN_STRUCT index. // E.g. for column metrics$cpu with parent metrics: diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java index a8b6f3dab395..0778975abf37 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java @@ -77,6 +77,8 @@ public class ColumnMetadataImpl implements ColumnMetadata { private final Set _partitions; private final boolean _autoGenerated; @Nullable + private final String _transformFunction; + @Nullable private final String _parentColumn; /// List of longs, each encodes: @@ -91,7 +93,7 @@ private ColumnMetadataImpl(FieldSpec fieldSpec, int totalDocs, int cardinality, boolean minMaxValueInvalid, int lengthOfShortestElement, int lengthOfLongestElement, boolean isAscii, int totalNumberOfEntries, int maxNumberOfMultiValues, int maxRowLengthInBytes, int bitsPerElement, @Nullable PartitionFunction partitionFunction, @Nullable Set partitions, boolean autoGenerated, - @Nullable String parentColumn) { + @Nullable String transformFunction, @Nullable String parentColumn) { _fieldSpec = fieldSpec; _totalDocs = totalDocs; _cardinality = cardinality; @@ -111,6 +113,7 @@ private ColumnMetadataImpl(FieldSpec fieldSpec, int totalDocs, int cardinality, _partitionFunction = partitionFunction; _partitions = partitions; _autoGenerated = autoGenerated; + _transformFunction = transformFunction; _parentColumn = parentColumn; } @@ -213,6 +216,12 @@ public boolean isAutoGenerated() { return _autoGenerated; } + @Nullable + @Override + public String getTransformFunction() { + return _transformFunction; + } + /// Returns {@code true} if this column is a materialized column produced from an OPEN_STRUCT parent column. public boolean isMaterializedChild() { return _parentColumn != null; @@ -296,6 +305,7 @@ public boolean equals(Object o) { && Objects.equals(_maxValue, that._maxValue) && Objects.equals(_partitionFunction, that._partitionFunction) && Objects.equals(_partitions, that._partitions) + && Objects.equals(_transformFunction, that._transformFunction) && Objects.equals(_parentColumn, that._parentColumn) && Objects.equals(_indexTypeSizeList, that._indexTypeSizeList); } @@ -305,7 +315,7 @@ public int hashCode() { return Objects.hash(_fieldSpec, _totalDocs, _cardinality, _hasDictionary, _forwardIndexEncoding, _sorted, _minValue, _maxValue, _minMaxValueInvalid, _lengthOfShortestElement, _lengthOfLongestElement, _isAscii, _totalNumberOfEntries, _maxNumberOfMultiValues, _maxRowLengthInBytes, _bitsPerElement, _partitionFunction, - _partitions, _autoGenerated, _parentColumn, _indexTypeSizeList); + _partitions, _autoGenerated, _transformFunction, _parentColumn, _indexTypeSizeList); } @Override @@ -330,6 +340,7 @@ public String toString() { + ", _partitionFunction=" + _partitionFunction + ", _partitions=" + _partitions + ", _autoGenerated=" + _autoGenerated + + ", _transformFunction=" + _transformFunction + ", _parentColumn=" + _parentColumn + ", _indexTypeSizeList=" + _indexTypeSizeList + '}'; @@ -359,6 +370,7 @@ public static ColumnMetadataImpl fromPropertiesConfiguration(PropertiesConfigura config.getInt(Column.getKeyFor(column, Column.MAX_ROW_LENGTH_IN_BYTES), UNAVAILABLE)) .setBitsPerElement(config.getInt(Column.getKeyFor(column, Column.BITS_PER_ELEMENT), UNAVAILABLE)) .setAutoGenerated(config.getBoolean(Column.getKeyFor(column, Column.IS_AUTO_GENERATED), false)) + .setTransformFunction(extractTransformFunction(column, config)) .setParentColumn(config.getString(Column.getKeyFor(column, Column.PARENT_COLUMN), null)); // Set min/max value @@ -395,6 +407,15 @@ public static ColumnMetadataImpl fromPropertiesConfiguration(PropertiesConfigura return builder.build(); } + @Nullable + private static String extractTransformFunction(String column, PropertiesConfiguration config) { + Object transformFunctionProperty = config.getProperty(Column.getKeyFor(column, Column.TRANSFORM_FUNCTION)); + String transformFunction = transformFunctionProperty != null ? transformFunctionProperty.toString() : null; + return transformFunction != null + ? CommonsConfigurationUtils.recoverSpecialCharacterInPropertyValue(transformFunction) + : null; + } + public static FieldSpec extractFieldSpec(String column, PropertiesConfiguration config) { String fieldName = config.getString(Column.getKeyFor(column, Column.COLUMN_NAME), column); FieldType fieldType = config.getEnum(Column.getKeyFor(column, Column.COLUMN_TYPE), FieldType.class); @@ -525,6 +546,7 @@ public static class Builder { private PartitionFunction _partitionFunction; private Set _partitions; private boolean _autoGenerated; + private String _transformFunction; private String _parentColumn; public Builder setFieldSpec(FieldSpec fieldSpec) { @@ -627,6 +649,11 @@ public Builder setAutoGenerated(boolean autoGenerated) { return this; } + public Builder setTransformFunction(String transformFunction) { + _transformFunction = transformFunction; + return this; + } + public Builder setParentColumn(String parentColumn) { _parentColumn = parentColumn; return this; @@ -671,7 +698,7 @@ public ColumnMetadataImpl build() { return new ColumnMetadataImpl(_fieldSpec, _totalDocs, _cardinality, _hasDictionary, _forwardIndexEncoding, _sorted, _minValue, _maxValue, _minMaxValueInvalid, _lengthOfShortestElement, _lengthOfLongestElement, _isAscii, _totalNumberOfEntries, _maxNumberOfMultiValues, _maxRowLengthInBytes, _bitsPerElement, - _partitionFunction, _partitions, _autoGenerated, _parentColumn); + _partitionFunction, _partitions, _autoGenerated, _transformFunction, _parentColumn); } } } diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/EmptyColumnMetadata.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/EmptyColumnMetadata.java index 241c9e22649a..a4457010b4b8 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/EmptyColumnMetadata.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/EmptyColumnMetadata.java @@ -72,6 +72,12 @@ public boolean isAutoGenerated() { return false; } + @Nullable + @Override + public String getTransformFunction() { + return null; + } + @Override public int getNumIndexes() { return 0; diff --git a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImplTest.java b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImplTest.java index ffa5ee2a981d..a8feaa5e5013 100644 --- a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImplTest.java +++ b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImplTest.java @@ -24,10 +24,13 @@ import org.apache.pinot.spi.data.DimensionFieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.FieldSpec.FieldType; +import org.apache.pinot.spi.env.CommonsConfigurationUtils; import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; @@ -132,6 +135,40 @@ public void parentColumnReadFromPropertiesConfig() { assertTrue(metadata.isMaterializedChild()); } + @Test + public void transformFunctionRoundtrip() { + String transformFunction = "Groovy({x + ',' + y}, x, y)"; + ColumnMetadataImpl meta = ColumnMetadataImpl.builder() + .setFieldSpec(new DimensionFieldSpec("col", DataType.STRING, true)) + .setTransformFunction(transformFunction) + .build(); + + assertEquals(meta.getTransformFunction(), transformFunction); + } + + @Test + public void transformFunctionReadFromPropertiesConfig() { + String transformFunction = "Groovy({x + ',' + y}, x, y)"; + String escapedTransformFunction = + CommonsConfigurationUtils.replaceSpecialCharacterInPropertyValue(transformFunction); + assertNotNull(escapedTransformFunction); + PropertiesConfiguration config = baseConfig("col"); + config.setProperty(Column.getKeyFor("col", Column.TRANSFORM_FUNCTION), escapedTransformFunction); + + ColumnMetadataImpl metadata = ColumnMetadataImpl.fromPropertiesConfiguration(config, 1, "col"); + + assertEquals(metadata.getTransformFunction(), transformFunction); + } + + @Test + public void missingTransformFunctionIsBackwardCompatible() { + PropertiesConfiguration config = baseConfig("col"); + + ColumnMetadataImpl metadata = ColumnMetadataImpl.fromPropertiesConfiguration(config, 1, "col"); + + assertNull(metadata.getTransformFunction()); + } + private static PropertiesConfiguration baseConfig(String column) { PropertiesConfiguration config = new PropertiesConfiguration(); config.setProperty(Column.getKeyFor(column, Column.COLUMN_NAME), column);