This is an automated email from the ASF dual-hosted git repository. xiangfu0 pushed a commit to branch xiangfu0/agg-rewrite/04-any-value in repository https://gitbox.apache.org/repos/asf/pinot.git
commit d9ccb8a790631274138aaffed8453c8d68a8f9d8 Author: Xiang Fu <[email protected]> AuthorDate: Tue Oct 20 18:32:01 2026 -0700 Return the input's logical type from ANY_VALUE Adds ANY_VALUE(expr, 'dataType'), whose result type is fixed at construction and whose values are read as that type, and rewrites ANY_VALUE(expr) to it in both query engines: AggregateFunctionRewriteOptimizer in the single-stage engine and a separate PinotAggregateTypeArgumentRewriteRule instance (AggregateAnyValueTypeRewrite) in the multi-stage engine. JSON maps to STRING in both. The inferAnyValueType=false query option keeps the form without the data type. The multi-stage engine rejects a data type argument that does not match the type of the value argument. getConvertedFinalResult now handles UUID, and the aggregation-only reducer skips formatting null results. --- .../common/utils/config/QueryOptionsUtils.java | 9 + .../function/AggregationFunctionUtils.java | 3 + .../function/AnyValueAggregationFunction.java | 96 +++++---- .../AggregateFunctionRewriteOptimizer.java | 67 +++--- .../query/reduce/AggregationDataTableReducer.java | 5 +- .../function/AggregationFunctionUtilsTest.java | 24 +++ .../function/AnyValueAggregationFunctionTest.java | 113 ++++++++++- .../AggregateFunctionRewriteOptimizerTest.java | 46 +++++ .../apache/pinot/queries/AnyValueQueriesTest.java | 226 +++++++++++++++++++++ .../PinotAggregateTypeArgumentRewriteRule.java | 75 +++---- .../calcite/rel/rules/PinotQueryRuleSets.java | 2 + .../org/apache/pinot/query/QueryEnvironment.java | 4 + .../src/test/resources/queries/AggregatePlans.json | 62 ++++++ .../src/test/resources/queries/UDFAggregates.json | 53 +++++ .../pinot/segment/spi/AggregationFunctionType.java | 40 +++- .../apache/pinot/spi/utils/CommonConstants.java | 6 + 16 files changed, 727 insertions(+), 104 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java index 936c968aaf0..d0728b02bb5 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java @@ -394,6 +394,15 @@ public class QueryOptionsUtils { return value == null || Boolean.parseBoolean(value); } + /// Returns whether ANY_VALUE(expr) is rewritten to take the data type of `expr`. Defaults to `true`. + public static boolean isInferAnyValueType(@Nullable Map<String, String> queryOptions) { + if (queryOptions == null) { + return true; + } + String value = queryOptions.get(QueryOptionKey.INFER_ANY_VALUE_TYPE); + return value == null || Boolean.parseBoolean(value); + } + public static boolean isSkipUpsertView(Map<String, String> queryOptions) { return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.SKIP_UPSERT_VIEW)); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index c52f772cb14..9dc4522ccd8 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -77,6 +77,7 @@ import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.exception.BadQueryRequestException; import org.apache.pinot.spi.query.QueryThreadContext; import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidUtils; /// The `AggregationFunctionUtils` class provides utility methods for aggregation function. @@ -341,6 +342,8 @@ public class AggregationFunctionUtils { return dataTable.getString(rowId, colId); case BYTES: return dataTable.getBytes(rowId, colId).getBytes(); + case UUID: + return UuidUtils.toUUID(dataTable.getBytes(rowId, colId)); case INT_ARRAY: return dataTable.getIntArray(rowId, colId); case LONG_ARRAY: diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java index c984363d2c5..023fc88199e 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.core.query.aggregation.function; +import com.google.common.base.Preconditions; import java.math.BigDecimal; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; @@ -47,6 +48,9 @@ import org.apache.pinot.spi.utils.ByteArray; /// only until it finds the first non-null value in the current batch for each group/key. This makes it O(n) over the /// input once per group until the first value is set, with early-exit fast paths when there are no nulls. /// +/// `ANY_VALUE(expr, 'dataType')` fixes the result type at construction and reads every block as that type. Without +/// the data type argument, the result type is the stored type of the first block, or STRING before any block. +/// /// **Example:** /// /// ``` @@ -58,12 +62,47 @@ import org.apache.pinot.spi.utils.ByteArray; /// ``` public class AnyValueAggregationFunction extends BaseSingleInputAggregationFunction<Object, Comparable<?>> { private static final DataType[] DATA_TYPE_VALUES = DataType.values(); - // Result type is determined at runtime based on input expression type + // Data type argument, or null when the result type is determined at runtime from the input blocks + @Nullable + private final DataType _dataType; private ColumnDataType _resultType; public AnyValueAggregationFunction(List<ExpressionContext> arguments, boolean nullHandlingEnabled) { - super(verifySingleArgument(arguments, "ANY_VALUE"), nullHandlingEnabled); - _resultType = null; + super(getValueArgument(arguments), nullHandlingEnabled); + if (arguments.size() == 2) { + ExpressionContext dataTypeArgument = arguments.get(1); + Preconditions.checkArgument(dataTypeArgument.getType() == ExpressionContext.Type.LITERAL, + "ANY_VALUE expects the 2nd argument to be literal, got: %s", dataTypeArgument.getType()); + _dataType = DataType.valueOf(dataTypeArgument.getLiteral().getStringValue().toUpperCase()); + _resultType = getResultType(_dataType); + } else { + _dataType = null; + } + } + + private static ExpressionContext getValueArgument(List<ExpressionContext> arguments) { + int numArguments = arguments.size(); + Preconditions.checkArgument(numArguments == 1 || numArguments == 2, "ANY_VALUE expects 1 or 2 arguments, got: %s", + numArguments); + return arguments.get(0); + } + + private static ColumnDataType getResultType(DataType dataType) { + switch (dataType) { + case INT: + case LONG: + case FLOAT: + case DOUBLE: + case BIG_DECIMAL: + case BOOLEAN: + case TIMESTAMP: + case STRING: + case BYTES: + case UUID: + return ColumnDataType.fromDataTypeSV(dataType); + default: + throw new IllegalArgumentException("Unsupported data type for ANY_VALUE: " + dataType); + } } @Override @@ -71,6 +110,14 @@ public class AnyValueAggregationFunction extends BaseSingleInputAggregationFunct return AggregationFunctionType.ANYVALUE; } + @Override + public String getResultColumnName() { + if (_dataType == null) { + return super.getResultColumnName(); + } + return getType().getName().toLowerCase() + "(" + _expression + ",'" + _dataType + "')"; + } + @Override public AggregationResultHolder createAggregationResultHolder() { return new ObjectAggregationResultHolder(); @@ -85,7 +132,7 @@ public class AnyValueAggregationFunction extends BaseSingleInputAggregationFunct public ColumnDataType getIntermediateResultColumnType() { // Default to STRING if result type is not yet determined // TODO: See if UNKNOWN can be used instead - return _resultType != null ? _resultType : ColumnDataType.STRING; + return _resultType != null ? _resultType.getStoredType() : ColumnDataType.STRING; } @Override @@ -184,13 +231,15 @@ public class AnyValueAggregationFunction extends BaseSingleInputAggregationFunct /// Generic helper for processing values with dictionary optimization for all supported data types private void aggregateHelper(int length, BlockValSet bvs, ValueProcessor<Object> processor) { + // Reading through the bound type keeps values consistent when segments store the column as different types + DataType storedType = _dataType != null ? _dataType.getStoredType() : bvs.getValueType().getStoredType(); // Use dictionary-based access for efficiency when available if (bvs.isDictionaryEncoded()) { final int[] dictIds = bvs.getDictionaryIdsSV(); final Dictionary dict = bvs.getDictionary(); forEachNotNull(length, bvs, (from, to) -> { for (int i = from; i < to; i++) { - Object value = getDictionaryValue(dict, dictIds[i], bvs.getValueType().getStoredType()); + Object value = getDictionaryValue(dict, dictIds[i], storedType); if (processor.process(i, value)) { break; } @@ -200,7 +249,7 @@ public class AnyValueAggregationFunction extends BaseSingleInputAggregationFunct // Fall back to direct value access based on type forEachNotNull(length, bvs, (from, to) -> { for (int i = from; i < to; i++) { - Object value = getDirectValue(bvs, i); + Object value = getDirectValue(bvs, i, storedType); if (value != null && processor.process(i, value)) { break; } @@ -232,8 +281,8 @@ public class AnyValueAggregationFunction extends BaseSingleInputAggregationFunct } /// Get value directly from BlockValSet based on data type - private Object getDirectValue(BlockValSet bvs, int index) { - switch (bvs.getValueType().getStoredType()) { + private Object getDirectValue(BlockValSet bvs, int index, DataType storedType) { + switch (storedType) { case INT: return bvs.getIntValuesSV()[index]; case LONG: @@ -249,7 +298,7 @@ public class AnyValueAggregationFunction extends BaseSingleInputAggregationFunct case BYTES: return new ByteArray(bvs.getBytesValuesSV()[index]); default: - throw new IllegalStateException("Unsupported direct access type: " + bvs.getValueType().getStoredType()); + throw new IllegalStateException("Unsupported direct access type: " + storedType); } } @@ -326,33 +375,8 @@ public class AnyValueAggregationFunction extends BaseSingleInputAggregationFunct } private void ensureResultType(BlockValSet bvs) { - if (_resultType != null) { - return; - } - switch (bvs.getValueType().getStoredType()) { - case INT: - _resultType = ColumnDataType.INT; - return; - case LONG: - _resultType = ColumnDataType.LONG; - return; - case FLOAT: - _resultType = ColumnDataType.FLOAT; - return; - case DOUBLE: - _resultType = ColumnDataType.DOUBLE; - return; - case STRING: - _resultType = ColumnDataType.STRING; - return; - case BIG_DECIMAL: - _resultType = ColumnDataType.BIG_DECIMAL; - return; - case BYTES: - _resultType = ColumnDataType.BYTES; - return; - default: - throw new IllegalStateException("ANY_VALUE unsupported type: " + bvs.getValueType()); + if (_resultType == null) { + _resultType = getResultType(bvs.getValueType().getStoredType()); } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java index 5bbc3274e05..a2cefda7258 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizer.java @@ -18,11 +18,14 @@ */ package org.apache.pinot.core.query.optimizer.statement; +import java.util.EnumSet; import java.util.List; +import java.util.Set; import javax.annotation.Nullable; import org.apache.pinot.common.request.Expression; import org.apache.pinot.common.request.Function; import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.spi.data.FieldSpec; @@ -36,6 +39,9 @@ import org.apache.pinot.spi.utils.CommonConstants; /// - FIRST_WITH_TIME(dataColumn, timeColumn) -> FIRST_WITH_TIME(dataColumn, timeColumn, 'dataType') /// - LAST_WITH_TIME(dataColumn, timeColumn) -> LAST_WITH_TIME(dataColumn, timeColumn, 'dataType') /// +/// Applied unless disabled through the `inferAnyValueType` query option: +/// - ANY_VALUE(column) -> ANY_VALUE(column, 'dataType') +/// /// Applied only when enabled through the `autoRewriteAggregationType` query option: /// - MIN(stringType) -> MINSTRING /// - MAX(stringType) -> MAXSTRING @@ -46,6 +52,14 @@ import org.apache.pinot.spi.utils.CommonConstants; /// /// Aggregations nested in other expressions, such as an alias, a FILTER clause or a transform, are rewritten as well. public class AggregateFunctionRewriteOptimizer implements StatementOptimizer { + // Data types of the explicit forms + private static final Set<FieldSpec.DataType> VALUE_WITH_TIME_DATA_TYPES = + EnumSet.of(FieldSpec.DataType.INT, FieldSpec.DataType.LONG, FieldSpec.DataType.FLOAT, FieldSpec.DataType.DOUBLE, + FieldSpec.DataType.STRING, FieldSpec.DataType.BOOLEAN, FieldSpec.DataType.TIMESTAMP); + private static final Set<FieldSpec.DataType> ANY_VALUE_DATA_TYPES = + EnumSet.of(FieldSpec.DataType.INT, FieldSpec.DataType.LONG, FieldSpec.DataType.FLOAT, FieldSpec.DataType.DOUBLE, + FieldSpec.DataType.BIG_DECIMAL, FieldSpec.DataType.STRING, FieldSpec.DataType.BOOLEAN, + FieldSpec.DataType.TIMESTAMP, FieldSpec.DataType.BYTES, FieldSpec.DataType.UUID); @Override public void optimize(PinotQuery pinotQuery, @Nullable Schema schema) { @@ -55,22 +69,25 @@ public class AggregateFunctionRewriteOptimizer implements StatementOptimizer { boolean rewriteToTypedVariants = pinotQuery.getQueryOptions() != null && Boolean.parseBoolean( pinotQuery.getQueryOptions().get(CommonConstants.Broker.Request.QueryOptionKey.AUTO_REWRITE_AGGREGATION_TYPE)); - rewrite(pinotQuery.getSelectList(), schema, rewriteToTypedVariants); - rewrite(pinotQuery.getGroupByList(), schema, rewriteToTypedVariants); - rewrite(pinotQuery.getOrderByList(), schema, rewriteToTypedVariants); - rewrite(pinotQuery.getFilterExpression(), schema, rewriteToTypedVariants); - rewrite(pinotQuery.getHavingExpression(), schema, rewriteToTypedVariants); + boolean inferAnyValueType = QueryOptionsUtils.isInferAnyValueType(pinotQuery.getQueryOptions()); + rewrite(pinotQuery.getSelectList(), schema, rewriteToTypedVariants, inferAnyValueType); + rewrite(pinotQuery.getGroupByList(), schema, rewriteToTypedVariants, inferAnyValueType); + rewrite(pinotQuery.getOrderByList(), schema, rewriteToTypedVariants, inferAnyValueType); + rewrite(pinotQuery.getFilterExpression(), schema, rewriteToTypedVariants, inferAnyValueType); + rewrite(pinotQuery.getHavingExpression(), schema, rewriteToTypedVariants, inferAnyValueType); } - private static void rewrite(@Nullable List<Expression> expressions, Schema schema, boolean rewriteToTypedVariants) { + private static void rewrite(@Nullable List<Expression> expressions, Schema schema, boolean rewriteToTypedVariants, + boolean inferAnyValueType) { if (expressions != null) { for (Expression expression : expressions) { - rewrite(expression, schema, rewriteToTypedVariants); + rewrite(expression, schema, rewriteToTypedVariants, inferAnyValueType); } } } - private static void rewrite(@Nullable Expression expression, Schema schema, boolean rewriteToTypedVariants) { + private static void rewrite(@Nullable Expression expression, Schema schema, boolean rewriteToTypedVariants, + boolean inferAnyValueType) { if (expression == null || !expression.isSetFunctionCall()) { return; } @@ -78,7 +95,7 @@ public class AggregateFunctionRewriteOptimizer implements StatementOptimizer { Function function = expression.getFunctionCall(); String functionName = function.getOperator(); if (!AggregationFunctionType.isAggregationFunction(functionName)) { - rewrite(function.getOperands(), schema, rewriteToTypedVariants); + rewrite(function.getOperands(), schema, rewriteToTypedVariants, inferAnyValueType); return; } @@ -94,7 +111,13 @@ public class AggregateFunctionRewriteOptimizer implements StatementOptimizer { String normalizedName = AggregationFunctionType.getNormalizedAggregationFunctionName(functionName); if ((normalizedName.equals(AggregationFunctionType.FIRSTWITHTIME.name()) || normalizedName.equals(AggregationFunctionType.LASTWITHTIME.name())) && function.getOperandsSize() == 2) { - addDataTypeArgument(function, fieldSpec); + addDataTypeArgument(function, fieldSpec, VALUE_WITH_TIME_DATA_TYPES); + return; + } + if (normalizedName.equals(AggregationFunctionType.ANYVALUE.name()) && function.getOperandsSize() == 1) { + if (inferAnyValueType) { + addDataTypeArgument(function, fieldSpec, ANY_VALUE_DATA_TYPES); + } return; } if (rewriteToTypedVariants) { @@ -103,26 +126,16 @@ public class AggregateFunctionRewriteOptimizer implements StatementOptimizer { } /// Appends the data column's logical type when the explicit form supports it. - private static void addDataTypeArgument(Function function, FieldSpec fieldSpec) { + private static void addDataTypeArgument(Function function, FieldSpec fieldSpec, + Set<FieldSpec.DataType> supportedDataTypes) { if (!fieldSpec.isSingleValueField()) { return; } - switch (fieldSpec.getDataType()) { - case INT: - case LONG: - case FLOAT: - case DOUBLE: - case STRING: - case BOOLEAN: - case TIMESTAMP: - function.addToOperands(RequestUtils.getLiteralExpression(fieldSpec.getDataType().name())); - break; - case JSON: - // Same as the multi-stage engine, which types a JSON column as VARCHAR. - function.addToOperands(RequestUtils.getLiteralExpression(FieldSpec.DataType.STRING.name())); - break; - default: - break; + // Same as the multi-stage engine, which types a JSON column as VARCHAR. + FieldSpec.DataType dataType = + fieldSpec.getDataType() == FieldSpec.DataType.JSON ? FieldSpec.DataType.STRING : fieldSpec.getDataType(); + if (supportedDataTypes.contains(dataType)) { + function.addToOperands(RequestUtils.getLiteralExpression(dataType.name())); } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/AggregationDataTableReducer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/AggregationDataTableReducer.java index 5c8143975f9..546d3ca7a8e 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/AggregationDataTableReducer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/AggregationDataTableReducer.java @@ -242,7 +242,10 @@ public class AggregationDataTableReducer implements DataTableReducer { int numColumns = columnDataTypes.length; for (Object[] rewrittenRow : rows) { for (int j = 0; j < numColumns; j++) { - rewrittenRow[j] = columnDataTypes[j].format(rewrittenRow[j]); + Object value = rewrittenRow[j]; + if (value != null) { + rewrittenRow[j] = columnDataTypes[j].format(value); + } } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java index 53cd082fc0c..9bce4e9bc5a 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java @@ -18,10 +18,18 @@ */ package org.apache.pinot.core.query.aggregation.function; +import java.util.UUID; +import org.apache.pinot.common.datatable.DataTable; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.core.common.datatable.DataTableBuilder; +import org.apache.pinot.core.common.datatable.DataTableBuilderFactory; import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.segment.spi.datasource.DataSource; import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.spi.exception.BadQueryRequestException; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidUtils; import org.testng.annotations.Test; import static org.mockito.Mockito.mock; @@ -114,4 +122,20 @@ public class AggregationFunctionUtilsTest { () -> AggregationFunctionUtils.getAggregationResult(mockFunction(AggregationFunctionType.MIN), null, 100, "TEST")); } + + @Test + public void testConvertedUuidFinalResult() + throws Exception { + // Reached when a single server returns the final result of e.g. ANY_VALUE(uuidCol, 'UUID'). + UUID uuid = UUID.fromString("123e4567-e89b-12d3-a456-426614174000"); + DataTableBuilder dataTableBuilder = DataTableBuilderFactory.getDataTableBuilder( + new DataSchema(new String[]{"uuidCol"}, new ColumnDataType[]{ColumnDataType.UUID})); + dataTableBuilder.startRow(); + dataTableBuilder.setColumn(0, new ByteArray(UuidUtils.toBytes(uuid))); + dataTableBuilder.finishRow(); + DataTable dataTable = dataTableBuilder.build(); + assertEquals(AggregationFunctionUtils.getConvertedFinalResult(dataTable, ColumnDataType.UUID, 0, 0), uuid); + assertEquals(AggregationFunctionUtils.getConvertedFinalResult(dataTable, ColumnDataType.UUID, 0, 0), + ColumnDataType.UUID.convert(AggregationFunctionUtils.getFinalResult(dataTable, ColumnDataType.UUID, 0, 0))); + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunctionTest.java index 8ae6327b6c6..eb1e1d372eb 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunctionTest.java @@ -18,11 +18,13 @@ */ package org.apache.pinot.core.query.aggregation.function; +import java.math.BigDecimal; import java.nio.ByteBuffer; import java.util.List; import java.util.Map; import org.apache.pinot.common.CustomObject; import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.core.common.BlockValSet; import org.apache.pinot.core.query.aggregation.AggregationResultHolder; import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; @@ -36,14 +38,18 @@ 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.assertNull; +import static org.testng.Assert.assertThrows; public class AnyValueAggregationFunctionTest extends AbstractAggregationFunctionTest { - @DataProvider(name = "bytesValueSources") - Object[][] bytesValueSources() { + private static final ExpressionContext EXPRESSION = ExpressionContext.forIdentifier("myField"); + + @DataProvider(name = "dictionaryEncoded") + Object[][] dictionaryEncoded() { return new Object[][]{{false}, {true}}; } - @Test(dataProvider = "bytesValueSources") + @Test(dataProvider = "dictionaryEncoded") void testBytesIntermediateResultRoundTrip(boolean dictionaryEncoded) { ExpressionContext expression = ExpressionContext.forIdentifier("myField"); AnyValueAggregationFunction function = new AnyValueAggregationFunction(List.of(expression), true); @@ -357,4 +363,105 @@ public class AnyValueAggregationFunctionTest extends AbstractAggregationFunction queryResult.thenResultIs(EXPECTED_COLUMN_TYPES, EXPECTED_NULL_RESULT); } } + + @DataProvider(name = "dataTypes") + Object[][] dataTypes() { + return new Object[][]{ + {"INT", ColumnDataType.INT, ColumnDataType.INT}, + {"LONG", ColumnDataType.LONG, ColumnDataType.LONG}, + {"FLOAT", ColumnDataType.FLOAT, ColumnDataType.FLOAT}, + {"DOUBLE", ColumnDataType.DOUBLE, ColumnDataType.DOUBLE}, + {"BIG_DECIMAL", ColumnDataType.BIG_DECIMAL, ColumnDataType.BIG_DECIMAL}, + {"BOOLEAN", ColumnDataType.INT, ColumnDataType.BOOLEAN}, + {"TIMESTAMP", ColumnDataType.LONG, ColumnDataType.TIMESTAMP}, + {"STRING", ColumnDataType.STRING, ColumnDataType.STRING}, + {"BYTES", ColumnDataType.BYTES, ColumnDataType.BYTES}, + {"uuid", ColumnDataType.BYTES, ColumnDataType.UUID} + }; + } + + @Test(dataProvider = "dataTypes") + void testDataTypeArgumentFixesResultTypes(String dataType, ColumnDataType intermediateResultType, + ColumnDataType finalResultType) { + // The types are known before any input, so an empty result is typed as well. + AnyValueAggregationFunction function = createFunction(dataType); + assertEquals(function.getIntermediateResultColumnType(), intermediateResultType); + assertEquals(function.getFinalResultColumnType(), finalResultType); + assertEquals(function.getResultColumnName(), "anyvalue(myField,'" + dataType.toUpperCase() + "')"); + assertEquals(function.getInputExpressions(), List.of(EXPRESSION)); + assertNull(function.extractAggregationResult(function.createAggregationResultHolder())); + } + + @Test + void testDataTypeArgumentRejectsInvalidArguments() { + assertThrows(IllegalArgumentException.class, () -> createFunction("JSON")); + assertThrows(IllegalArgumentException.class, () -> createFunction("NOT_A_TYPE")); + assertThrows(IllegalArgumentException.class, + () -> new AnyValueAggregationFunction(List.of(EXPRESSION, ExpressionContext.forIdentifier("INT")), false)); + assertThrows(IllegalArgumentException.class, () -> new AnyValueAggregationFunction( + List.of(EXPRESSION, ExpressionContext.forLiteral(DataType.STRING, "INT"), EXPRESSION), false)); + } + + @Test(dataProvider = "dictionaryEncoded") + void testDataTypeArgumentReadsBoundStoredType(boolean dictionaryEncoded) { + // A segment that still stores the column as INT after the schema widened it to LONG + AnyValueAggregationFunction function = createFunction("LONG"); + BlockValSet valueSet = mock(BlockValSet.class); + when(valueSet.getValueType()).thenReturn(DataType.INT); + if (dictionaryEncoded) { + Dictionary dictionary = mock(Dictionary.class); + when(dictionary.getLongValue(0)).thenReturn(5L); + when(valueSet.isDictionaryEncoded()).thenReturn(true); + when(valueSet.getDictionary()).thenReturn(dictionary); + when(valueSet.getDictionaryIdsSV()).thenReturn(new int[]{0}); + } else { + when(valueSet.getLongValuesSV()).thenReturn(new long[]{5L}); + } + Map<ExpressionContext, BlockValSet> values = Map.of(EXPRESSION, valueSet); + + AggregationResultHolder holder = function.createAggregationResultHolder(); + function.aggregate(1, holder, values); + assertEquals(function.extractAggregationResult(holder), 5L); + GroupByResultHolder groupHolder = function.createGroupByResultHolder(1, 1); + function.aggregateGroupBySV(1, new int[]{0}, groupHolder, values); + assertEquals(function.extractGroupByResult(groupHolder, 0), 5L); + assertEquals(function.getIntermediateResultColumnType(), ColumnDataType.LONG); + } + + @Test + void testDataTypeArgumentSerializedMerge() { + // One server found a value and the other did not. + AnyValueAggregationFunction function = createFunction("BOOLEAN"); + BlockValSet valueSet = mock(BlockValSet.class); + when(valueSet.getValueType()).thenReturn(DataType.BOOLEAN); + when(valueSet.getIntValuesSV()).thenReturn(new int[]{1}); + AggregationResultHolder holder = function.createAggregationResultHolder(); + function.aggregate(1, holder, Map.of(EXPRESSION, valueSet)); + CustomObject found = serialize(function, function.extractAggregationResult(holder)); + CustomObject empty = + serialize(function, function.extractAggregationResult(function.createAggregationResultHolder())); + + Object merged = AggregationFunctionUtils.mergeSerialized(function, null, empty); + assertNull(merged); + merged = AggregationFunctionUtils.mergeSerialized(function, merged, found); + merged = AggregationFunctionUtils.mergeSerialized(function, merged, empty); + assertEquals(merged, 1); + assertEquals(function.getFinalResultColumnType().convert(function.extractFinalResult(merged)), true); + + AnyValueAggregationFunction bigDecimalFunction = createFunction("BIG_DECIMAL"); + BigDecimal value = new BigDecimal("12345678901234567890.123"); + assertEquals(bigDecimalFunction.merge(bigDecimalFunction.deserializeIntermediateResult( + serialize(bigDecimalFunction, value)), new BigDecimal("1")), value); + } + + private static AnyValueAggregationFunction createFunction(String dataType) { + return new AnyValueAggregationFunction(List.of(EXPRESSION, ExpressionContext.forLiteral(DataType.STRING, dataType)), + true); + } + + private static CustomObject serialize(AnyValueAggregationFunction function, Object intermediateResult) { + AggregationFunction.SerializedIntermediateResult serialized = + function.serializeIntermediateResult(intermediateResult); + return new CustomObject(serialized.getType(), ByteBuffer.wrap(serialized.getBytes())); + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java index d44bc13a13a..8a84b017af9 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/optimizer/statement/AggregateFunctionRewriteOptimizerTest.java @@ -36,6 +36,7 @@ public class AggregateFunctionRewriteOptimizerTest { .addSingleValueDimension("longCol", DataType.LONG) .addSingleValueDimension("floatCol", DataType.FLOAT) .addSingleValueDimension("doubleCol", DataType.DOUBLE) + .addMetric("bigDecimalCol", DataType.BIG_DECIMAL) .addSingleValueDimension("stringCol", DataType.STRING) .addSingleValueDimension("boolCol", DataType.BOOLEAN) .addSingleValueDimension("timestampCol", DataType.TIMESTAMP) @@ -88,6 +89,51 @@ public class AggregateFunctionRewriteOptimizerTest { + "ORDER BY LAST_WITH_TIME(timestampCol, timeCol, 'TIMESTAMP')"); } + @DataProvider + public Object[][] anyValueDataTypes() { + return new Object[][]{ + {"intCol", "INT"}, {"longCol", "LONG"}, {"floatCol", "FLOAT"}, {"doubleCol", "DOUBLE"}, + {"bigDecimalCol", "BIG_DECIMAL"}, {"stringCol", "STRING"}, {"jsonCol", "STRING"}, {"boolCol", "BOOLEAN"}, + {"timestampCol", "TIMESTAMP"}, {"bytesCol", "BYTES"}, {"uuidCol", "UUID"} + }; + } + + @Test(dataProvider = "anyValueDataTypes") + public void testAnyValueInfersDataType(String column, String dataType) { + for (String function : new String[]{"ANY_VALUE", "anyValue"}) { + assertRewrite("SELECT " + function + "(" + column + ") FROM testTable", + "SELECT " + function + "(" + column + ", '" + dataType + "') FROM testTable"); + } + } + + @Test + public void testAnyValueKeepsUnsupportedArguments() { + for (String column : new String[]{"mvCol", "missingCol", "trim(stringCol)"}) { + assertUnchanged("SELECT ANY_VALUE(" + column + ") FROM testTable"); + } + assertUnchanged("SELECT ANY_VALUE(intCol, 'LONG') FROM testTable"); + } + + @Test + public void testRewritesNestedAnyValueConsistently() { + assertRewrite("SELECT intCol, ANY_VALUE(stringCol) AS v1, ANY_VALUE(doubleCol) FILTER (WHERE boolCol) " + + "FROM testTable GROUP BY intCol HAVING ANY_VALUE(longCol) > 1 ORDER BY ANY_VALUE(timestampCol)", + "SELECT intCol, ANY_VALUE(stringCol, 'STRING') AS v1, ANY_VALUE(doubleCol, 'DOUBLE') FILTER (WHERE boolCol) " + + "FROM testTable GROUP BY intCol HAVING ANY_VALUE(longCol, 'LONG') > 1 " + + "ORDER BY ANY_VALUE(timestampCol, 'TIMESTAMP')"); + } + + @Test + public void testAnyValueDataTypeCanBeDisabled() { + String query = "SELECT ANY_VALUE(intCol), FIRST_WITH_TIME(intCol, timeCol) FROM testTable"; + String expectedQuery = "SELECT ANY_VALUE(intCol, 'INT'), FIRST_WITH_TIME(intCol, timeCol, 'INT') FROM testTable"; + assertRewrite("SET inferAnyValueType=true; " + query, "SET inferAnyValueType=true; " + expectedQuery); + // The option only applies to ANY_VALUE. + assertRewrite("SET inferAnyValueType=false; " + query, + "SET inferAnyValueType=false; SELECT ANY_VALUE(intCol), FIRST_WITH_TIME(intCol, timeCol, 'INT') " + + "FROM testTable"); + } + @Test public void testTypedVariantsRequireQueryOption() { String query = "SELECT MIN(stringCol) AS m, MAX(longCol), SUM(intCol), SUM(longCol), MIN(doubleCol) " diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/AnyValueQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/AnyValueQueriesTest.java new file mode 100644 index 00000000000..5f7441cce0e --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/queries/AnyValueQueriesTest.java @@ -0,0 +1,226 @@ +/** + * 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.queries; + +import java.io.File; +import java.math.BigDecimal; +import java.sql.Timestamp; +import java.util.ArrayList; +import java.util.List; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +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.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +/// Queries test for ANY_VALUE, whose data type argument the optimizer supplies from the schema. All the rows with the +/// same key hold the same values, so that any value is the expected one. +public class AnyValueQueriesTest extends BaseQueriesTest { + private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "AnyValueQueriesTest"); + private static final String RAW_TABLE_NAME = "testTable"; + private static final String SEGMENT_NAME = "testSegment"; + private static final int NUM_RECORDS = 10; + + private static final Schema SCHEMA = new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME) + .addSingleValueDimension("keyCol", DataType.INT) + .addSingleValueDimension("intCol", DataType.INT) + .addSingleValueDimension("longCol", DataType.LONG) + .addSingleValueDimension("floatCol", DataType.FLOAT) + .addSingleValueDimension("doubleCol", DataType.DOUBLE) + .addMetric("bigDecimalCol", DataType.BIG_DECIMAL) + .addSingleValueDimension("boolCol", DataType.BOOLEAN) + .addSingleValueDimension("timestampCol", DataType.TIMESTAMP) + .addSingleValueDimension("stringCol", DataType.STRING) + .addSingleValueDimension("jsonCol", DataType.JSON) + .addSingleValueDimension("bytesCol", DataType.BYTES) + .addSingleValueDimension("uuidCol", DataType.UUID) + .build(); + private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) + .setNoDictionaryColumns(List.of("longCol", "doubleCol", "bytesCol")).build(); + private static final String ALL_COLUMNS = "ANY_VALUE(intCol), ANY_VALUE(longCol), ANY_VALUE(floatCol), " + + "ANY_VALUE(doubleCol), ANY_VALUE(bigDecimalCol), ANY_VALUE(boolCol), ANY_VALUE(timestampCol), " + + "ANY_VALUE(stringCol), ANY_VALUE(jsonCol), ANY_VALUE(bytesCol), ANY_VALUE(uuidCol)"; + + private IndexSegment _indexSegment; + private List<IndexSegment> _indexSegments; + + @Override + protected String getFilter() { + return ""; + } + + @Override + protected IndexSegment getIndexSegment() { + return _indexSegment; + } + + @Override + protected List<IndexSegment> getIndexSegments() { + return _indexSegments; + } + + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteDirectory(INDEX_DIR); + + List<GenericRow> records = new ArrayList<>(NUM_RECORDS); + for (int i = 0; i < NUM_RECORDS; i++) { + int key = i % 2; + GenericRow record = new GenericRow(); + record.putValue("keyCol", key); + record.putValue("intCol", 10 + key); + record.putValue("longCol", 10_000_000_000L + key); + record.putValue("floatCol", 0.5f + key); + record.putValue("doubleCol", 0.25 + key); + record.putValue("bigDecimalCol", new BigDecimal("12345678901234567890.5").add(BigDecimal.valueOf(key))); + record.putValue("boolCol", key == 0); + record.putValue("timestampCol", 1_600_000_000_000L + key); + record.putValue("stringCol", "s" + key); + record.putValue("jsonCol", "{\"k\":" + key + "}"); + record.putValue("bytesCol", new byte[]{(byte) key, 1}); + record.putValue("uuidCol", "123e4567-e89b-12d3-a456-42661417400" + key); + records.add(record); + } + + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA); + segmentGeneratorConfig.setTableName(RAW_TABLE_NAME); + segmentGeneratorConfig.setSegmentName(SEGMENT_NAME); + segmentGeneratorConfig.setOutDir(INDEX_DIR.getPath()); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records)); + driver.build(); + + ImmutableSegment immutableSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME), ReadMode.mmap); + _indexSegment = immutableSegment; + _indexSegments = List.of(immutableSegment, immutableSegment); + } + + @Test + public void testAggregationOnly() { + BrokerResponseNative brokerResponse = + getBrokerResponseForOptimizedQuery("SELECT " + ALL_COLUMNS + " FROM testTable WHERE keyCol = 0", SCHEMA); + DataSchema expectedDataSchema = new DataSchema(new String[]{ + "anyvalue(intCol,'INT')", "anyvalue(longCol,'LONG')", "anyvalue(floatCol,'FLOAT')", + "anyvalue(doubleCol,'DOUBLE')", "anyvalue(bigDecimalCol,'BIG_DECIMAL')", "anyvalue(boolCol,'BOOLEAN')", + "anyvalue(timestampCol,'TIMESTAMP')", "anyvalue(stringCol,'STRING')", "anyvalue(jsonCol,'STRING')", + "anyvalue(bytesCol,'BYTES')", "anyvalue(uuidCol,'UUID')" + }, new ColumnDataType[]{ + ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.FLOAT, ColumnDataType.DOUBLE, + ColumnDataType.BIG_DECIMAL, ColumnDataType.BOOLEAN, ColumnDataType.TIMESTAMP, ColumnDataType.STRING, + ColumnDataType.STRING, ColumnDataType.BYTES, ColumnDataType.UUID + }); + Object[] expectedRow = new Object[]{ + 10, 10_000_000_000L, 0.5f, 0.25, "12345678901234567890.5", true, + new Timestamp(1_600_000_000_000L).toString(), "s0", "{\"k\":0}", "0001", "123e4567-e89b-12d3-a456-426614174000" + }; + QueriesTestUtils.testInterSegmentsResult(brokerResponse, + new ResultTable(expectedDataSchema, List.<Object[]>of(expectedRow))); + } + + @Test + public void testGroupBy() { + ResultTable resultTable = getBrokerResponseForOptimizedQuery( + "SELECT keyCol, ANY_VALUE(intCol) AS v1, ANY_VALUE(boolCol) AS v2, ANY_VALUE(timestampCol) AS v3, " + + "ANY_VALUE(uuidCol) AS v4 FROM testTable GROUP BY keyCol ORDER BY ANY_VALUE(longCol) DESC", + SCHEMA).getResultTable(); + assertEquals(resultTable.getDataSchema(), new DataSchema(new String[]{"keyCol", "v1", "v2", "v3", "v4"}, + new ColumnDataType[]{ + ColumnDataType.INT, ColumnDataType.INT, ColumnDataType.BOOLEAN, ColumnDataType.TIMESTAMP, + ColumnDataType.UUID + })); + List<Object[]> rows = resultTable.getRows(); + assertEquals(rows.size(), 2); + assertEquals(rows.get(0), new Object[]{ + 1, 11, false, new Timestamp(1_600_000_000_001L).toString(), "123e4567-e89b-12d3-a456-426614174001" + }); + assertEquals(rows.get(1), new Object[]{ + 0, 10, true, new Timestamp(1_600_000_000_000L).toString(), "123e4567-e89b-12d3-a456-426614174000" + }); + } + + @Test + public void testEmptyInput() { + ResultTable resultTable = + getBrokerResponseForOptimizedQuery("SELECT " + ALL_COLUMNS + " FROM testTable WHERE keyCol < 0", SCHEMA) + .getResultTable(); + assertEquals(resultTable.getDataSchema().getColumnDataTypes(), new ColumnDataType[]{ + ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.FLOAT, ColumnDataType.DOUBLE, + ColumnDataType.BIG_DECIMAL, ColumnDataType.BOOLEAN, ColumnDataType.TIMESTAMP, ColumnDataType.STRING, + ColumnDataType.STRING, ColumnDataType.BYTES, ColumnDataType.UUID + }); + assertEquals(resultTable.getRows().size(), 1); + assertEquals(resultTable.getRows().get(0), new Object[11]); + + resultTable = getBrokerResponseForOptimizedQuery( + "SELECT keyCol, ANY_VALUE(boolCol), ANY_VALUE(timestampCol) FROM testTable WHERE keyCol < 0 GROUP BY keyCol", + SCHEMA).getResultTable(); + assertEquals(resultTable.getDataSchema().getColumnDataTypes(), + new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.BOOLEAN, ColumnDataType.TIMESTAMP}); + assertTrue(resultTable.getRows().isEmpty()); + } + + @Test + public void testExplicitDataType() { + ResultTable resultTable = getBrokerResponseForOptimizedQuery( + "SELECT ANY_VALUE(intCol, 'LONG'), ANY_VALUE(intCol, 'STRING'), ANY_VALUE(timestampCol, 'LONG') " + + "FROM testTable WHERE keyCol = 0", SCHEMA).getResultTable(); + assertEquals(resultTable.getDataSchema().getColumnDataTypes(), + new ColumnDataType[]{ColumnDataType.LONG, ColumnDataType.STRING, ColumnDataType.LONG}); + assertEquals(resultTable.getRows().get(0), new Object[]{10L, "10", 1_600_000_000_000L}); + } + + @Test + public void testWithoutInferredDataType() { + ResultTable resultTable = getBrokerResponseForOptimizedQuery( + "SET inferAnyValueType=false; SELECT ANY_VALUE(intCol), ANY_VALUE(boolCol), ANY_VALUE(timestampCol) " + + "FROM testTable WHERE keyCol = 0", SCHEMA).getResultTable(); + assertEquals(resultTable.getDataSchema(), + new DataSchema(new String[]{"anyvalue(intCol)", "anyvalue(boolCol)", "anyvalue(timestampCol)"}, + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING})); + assertEquals(resultTable.getRows().get(0), new Object[]{"10", "1", "1600000000000"}); + } + + @AfterClass + public void tearDown() + throws Exception { + _indexSegment.destroy(); + FileUtils.deleteDirectory(INDEX_DIR); + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java index dc4f8455363..846ed8b4251 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateTypeArgumentRewriteRule.java @@ -19,6 +19,7 @@ package org.apache.pinot.calcite.rel.rules; import java.util.ArrayList; +import java.util.EnumSet; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -35,30 +36,56 @@ import org.apache.calcite.rel.logical.LogicalAggregate; import org.apache.calcite.rel.logical.LogicalProject; import org.apache.calcite.rex.RexBuilder; import org.apache.calcite.rex.RexNode; -import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.tools.RelBuilderFactory; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.query.planner.logical.RelToPlanNodeConverter; import org.apache.pinot.segment.spi.AggregationFunctionType; /// Supplies the explicit data type argument from the type of the data argument: /// - FIRST_WITH_TIME(dataColumn, timeColumn) -> FIRST_WITH_TIME(dataColumn, timeColumn, 'dataType') /// - LAST_WITH_TIME(dataColumn, timeColumn) -> LAST_WITH_TIME(dataColumn, timeColumn, 'dataType') +/// - ANY_VALUE(expr) -> ANY_VALUE(expr, 'dataType'), by a separate instance so that it can be skipped on its own /// /// The result type of the call is unchanged. The literal is projected directly beneath the aggregate, appended to the /// existing project when there is one, because [PinotAggregateExchangeNodeInsertRule] and `AggregatePushdownRule` /// inline literal arguments only from the immediate project. public class PinotAggregateTypeArgumentRewriteRule extends RelOptRule { + private static final Set<String> VALUE_WITH_TIME_FUNCTIONS = + Set.of(AggregationFunctionType.FIRSTWITHTIME.name(), AggregationFunctionType.LASTWITHTIME.name()); + private static final Set<String> ANY_VALUE_FUNCTIONS = Set.of(AggregationFunctionType.ANYVALUE.name()); + // Data types of the explicit forms, by the type the multi-stage engine reads the data argument as + private static final Set<ColumnDataType> DATA_TYPES = + EnumSet.of(ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.FLOAT, ColumnDataType.DOUBLE, + ColumnDataType.BIG_DECIMAL, ColumnDataType.BOOLEAN, ColumnDataType.TIMESTAMP, ColumnDataType.STRING, + ColumnDataType.BYTES, ColumnDataType.UUID); + public static final PinotAggregateTypeArgumentRewriteRule INSTANCE = - new PinotAggregateTypeArgumentRewriteRule(PinotRuleUtils.PINOT_REL_FACTORY, null); + new PinotAggregateTypeArgumentRewriteRule(VALUE_WITH_TIME_FUNCTIONS, 2, PinotRuleUtils.PINOT_REL_FACTORY, null); + /// Rewrites FIRST_WITH_TIME and LAST_WITH_TIME. public static PinotAggregateTypeArgumentRewriteRule instanceWithDescription(String description) { - return new PinotAggregateTypeArgumentRewriteRule(PinotRuleUtils.PINOT_REL_FACTORY, description); + return new PinotAggregateTypeArgumentRewriteRule(VALUE_WITH_TIME_FUNCTIONS, 2, PinotRuleUtils.PINOT_REL_FACTORY, + description); + } + + /// Rewrites ANY_VALUE. + public static PinotAggregateTypeArgumentRewriteRule anyValueInstanceWithDescription(String description) { + return new PinotAggregateTypeArgumentRewriteRule(ANY_VALUE_FUNCTIONS, 1, PinotRuleUtils.PINOT_REL_FACTORY, + description); } + // Normalized names of the rewritten functions, and their number of arguments without the data type argument + private final Set<String> _functionNames; + private final int _numArguments; + // Matches with RelOptRule operands, like the other aggregate rewrite rules, instead of a RelRule config. @SuppressWarnings("deprecation") - private PinotAggregateTypeArgumentRewriteRule(RelBuilderFactory factory, @Nullable String description) { + private PinotAggregateTypeArgumentRewriteRule(Set<String> functionNames, int numArguments, RelBuilderFactory factory, + @Nullable String description) { super(operand(LogicalAggregate.class, any()), factory, description); + _functionNames = functionNames; + _numArguments = numArguments; } @Override @@ -126,39 +153,15 @@ public class PinotAggregateTypeArgumentRewriteRule extends RelOptRule { /// Returns the data type argument to append, or `null` when the call already has one or is not such a call. @Nullable - private static String getDataTypeArgument(AggregateCall aggCall, RelNode input) { - if (aggCall.getArgList().size() != 2) { + private String getDataTypeArgument(AggregateCall aggCall, RelNode input) { + if (aggCall.getArgList().size() != _numArguments || !_functionNames.contains( + AggregationFunctionType.getNormalizedAggregationFunctionName(aggCall.getAggregation().getName()))) { return null; } - String name = AggregationFunctionType.getNormalizedAggregationFunctionName(aggCall.getAggregation().getName()); - if (!name.equals(AggregationFunctionType.FIRSTWITHTIME.name()) - && !name.equals(AggregationFunctionType.LASTWITHTIME.name())) { - return null; - } - SqlTypeName dataType = - input.getRowType().getFieldList().get(aggCall.getArgList().get(0)).getType().getSqlTypeName(); - switch (dataType) { - case TINYINT: - case SMALLINT: - case INTEGER: - return "INT"; - case BIGINT: - return "LONG"; - case REAL: - case FLOAT: - return "FLOAT"; - case DOUBLE: - return "DOUBLE"; - case CHAR: - case VARCHAR: - return "STRING"; - case BOOLEAN: - return "BOOLEAN"; - case TIMESTAMP: - return "TIMESTAMP"; - default: - // Rejected during validation; leave the call as it is rather than guess a type. - return null; - } + ColumnDataType dataType = RelToPlanNodeConverter.convertToColumnDataType( + input.getRowType().getFieldList().get(aggCall.getArgList().get(0)).getType()); + // Validation limits FIRST_WITH_TIME / LAST_WITH_TIME to a subset of these types, and ANY_VALUE keeps its form + // without the data type for the others; leave the call as it is rather than guess a type. + return DATA_TYPES.contains(dataType) ? dataType.name() : null; } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java index 9393ed82cc8..6b50db44cc4 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java @@ -178,6 +178,8 @@ public class PinotQueryRuleSets { // Must stay ahead of PinotAggregateExchangeNodeInsertRule (POST_LOGICAL), which inlines the literal argument. PinotAggregateTypeArgumentRewriteRule .instanceWithDescription(PlannerRuleNames.AGGREGATE_TYPE_ARGUMENT_REWRITE), + PinotAggregateTypeArgumentRewriteRule + .anyValueInstanceWithDescription(PlannerRuleNames.AGGREGATE_ANY_VALUE_TYPE_REWRITE), // Must stay ahead of PinotAggregateExchangeNodeInsertRule (POST_LOGICAL), which fixes the leaf-to-final // intermediate result format from the function name. diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java index ce5d25c7558..33394a5f300 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java @@ -204,6 +204,10 @@ public class QueryEnvironment { } if (MapUtils.isNotEmpty(options)) { Set<String> skipRuleSet = QueryOptionsUtils.getSkipPlannerRules(options); + if (!QueryOptionsUtils.isInferAnyValueType(options)) { + skipRuleSet = new HashSet<>(skipRuleSet); + skipRuleSet.add(CommonConstants.Broker.PlannerRuleNames.AGGREGATE_ANY_VALUE_TYPE_REWRITE); + } if (!skipRuleSet.isEmpty() || !useRuleSet.isEmpty()) { // dynamically create optProgram according to rule options optProgram = getOptProgram(_envConfig.getRuleSet(), skipRuleSet, useRuleSet, _defaultDisabledPlannerRules); diff --git a/pinot-query-planner/src/test/resources/queries/AggregatePlans.json b/pinot-query-planner/src/test/resources/queries/AggregatePlans.json index c10f7185b79..c04a681dabd 100644 --- a/pinot-query-planner/src/test/resources/queries/AggregatePlans.json +++ b/pinot-query-planner/src/test/resources/queries/AggregatePlans.json @@ -184,6 +184,68 @@ "\n" ] }, + { + "description": "ANY_VALUE without a data type infers it from its argument", + "sql": "EXPLAIN PLAN FOR SELECT ANY_VALUE(a.col1), ANY_VALUE(a.col3), ANY_VALUE(a.col4), ANY_VALUE(a.col5), ANY_VALUE(a.col7), ANY_VALUE(a.ts_timestamp) FROM a", + "output": [ + "Execution Plan", + "\nPinotLogicalAggregate(group=[{}], agg#0=[ANYVALUE($0, _UTF-8'STRING')], agg#1=[ANYVALUE($1, _UTF-8'INT')], agg#2=[ANYVALUE($2, _UTF-8'BIG_DECIMAL')], agg#3=[ANYVALUE($3, _UTF-8'BOOLEAN')], agg#4=[ANYVALUE($4, _UTF-8'LONG')], agg#5=[ANYVALUE($5, _UTF-8'TIMESTAMP')], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash])", + "\n PinotLogicalAggregate(group=[{}], agg#0=[ANYVALUE($0, _UTF-8'STRING')], agg#1=[ANYVALUE($1, _UTF-8'INT')], agg#2=[ANYVALUE($2, _UTF-8'BIG_DECIMAL')], agg#3=[ANYVALUE($3, _UTF-8'BOOLEAN')], agg#4=[ANYVALUE($4, _UTF-8'LONG')], agg#5=[ANYVALUE($5, _UTF-8'TIMESTAMP')], aggType=[LEAF])", + "\n LogicalProject($f0=[$0], $f1=[$2], $f2=[$3], $f3=[$4], $f4=[$6], $f5=[$8], $f6=[_UTF-8'STRING'], $f7=[_UTF-8'INT'], $f8=[_UTF-8'BIG_DECIMAL'], $f9=[_UTF-8'BOOLEAN'], $f10=[_UTF-8'LONG'], $f11=[_UTF-8'TIMESTAMP'])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n" + ] + }, + { + "description": "Inferred ANY_VALUE data type for a transform and a filtered aggregate in a grouped aggregate", + "sql": "EXPLAIN PLAN FOR SELECT a.col2, ANY_VALUE(a.col3 + 1), ANY_VALUE(a.col5) FILTER (WHERE a.col3 > 0), ANY_VALUE(a.col1, 'STRING') FROM a GROUP BY a.col2", + "output": [ + "Execution Plan", + "\nPinotLogicalAggregate(group=[{0}], agg#0=[ANYVALUE($1, _UTF-8'INT')], agg#1=[ANYVALUE($2, _UTF-8'BOOLEAN')], agg#2=[ANYVALUE($3, _UTF-8'STRING')], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash[0]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[ANYVALUE($1, _UTF-8'INT')], agg#1=[ANYVALUE($2, _UTF-8'BOOLEAN') FILTER $3], agg#2=[ANYVALUE($4, _UTF-8'STRING')], aggType=[LEAF])", + "\n LogicalProject($f0=[$1], $f1=[+($2, 1)], $f2=[$4], $f3=[>($2, 0)], $f4=[$0], $f5=[_UTF-8'STRING'], $f6=[_UTF-8'INT'], $f7=[_UTF-8'BOOLEAN'])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n" + ] + }, + { + "description": "ANY_VALUE over a DECIMAL expression takes the type the expression is read as", + "sql": "EXPLAIN PLAN FOR SELECT ANY_VALUE(a.col3 * 1.5), ANY_VALUE(CAST(a.col3 AS DECIMAL(10, 0))) FROM a", + "output": [ + "Execution Plan", + "\nPinotLogicalAggregate(group=[{}], agg#0=[ANYVALUE($0, _UTF-8'DOUBLE')], agg#1=[ANYVALUE($1, _UTF-8'INT')], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash])", + "\n PinotLogicalAggregate(group=[{}], agg#0=[ANYVALUE($0, _UTF-8'DOUBLE')], agg#1=[ANYVALUE($1, _UTF-8'INT')], aggType=[LEAF])", + "\n LogicalProject($f0=[*($2, 1.5:DECIMAL(2, 1))], $f1=[CAST($2):DECIMAL(10, 0) NOT NULL], $f2=[_UTF-8'DOUBLE'], $f3=[_UTF-8'INT'])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n" + ] + }, + { + "description": "ANY_VALUE rejects a data type that does not match its argument", + "sql": "EXPLAIN PLAN FOR SELECT ANY_VALUE(a.col3, 'STRING') FROM a", + "expectedException": ".*Unsupported data type for ANYVALUE\\(INTEGER\\): STRING.*" + }, + { + "description": "ANY_VALUE rejects an unknown data type", + "sql": "EXPLAIN PLAN FOR SELECT ANY_VALUE(a.col1, 'NOT_A_TYPE') FROM a", + "expectedException": ".*Unsupported data type for ANYVALUE\\(VARCHAR\\): NOT_A_TYPE.*" + }, + { + "description": "ANY_VALUE keeps the form without a data type when inferAnyValueType is false", + "sql": "SET inferAnyValueType=false; EXPLAIN PLAN FOR SELECT ANY_VALUE(a.col5), FIRST_WITH_TIME(a.col5, a.ts) FROM a", + "output": [ + "Execution Plan", + "\nPinotLogicalAggregate(group=[{}], agg#0=[ANYVALUE($0)], agg#1=[FIRSTWITHTIME($1, $1, _UTF-8'BOOLEAN')], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash])", + "\n PinotLogicalAggregate(group=[{}], agg#0=[ANYVALUE($0)], agg#1=[FIRSTWITHTIME($0, $1, _UTF-8'BOOLEAN')], aggType=[LEAF])", + "\n LogicalProject($f0=[$4], $f1=[$7], $f2=[_UTF-8'BOOLEAN'])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n" + ] + }, { "description": "FIRST_WITH_TIME without a data type rejects an unsupported data argument", "sql": "EXPLAIN PLAN FOR SELECT FIRST_WITH_TIME(a.col4, a.ts) FROM a", diff --git a/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json b/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json index 01a19d7837c..89a48b44b5e 100644 --- a/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json +++ b/pinot-query-runtime/src/test/resources/queries/UDFAggregates.json @@ -77,6 +77,27 @@ "sql": "SELECT FIRST_WITH_TIME(timestamp_col, timestamp_col), LAST_WITH_TIME(timestamp_col, timestamp_col), FIRST_WITH_TIME(timestamp_col, timestamp_col, 'TIMESTAMP') FROM {tbl}", "outputs": [["2020-01-01 03:10:12.0", "2020-05-01 03:32:12.0", "2020-01-01 03:10:12.0"]] }, + { + "comment": "The data type argument of ANY_VALUE is inferred when omitted", + "sql": "SELECT ANY_VALUE(int_col), ANY_VALUE(long_col), ANY_VALUE(float_col), ANY_VALUE(double_col), ANY_VALUE(string_col), ANY_VALUE(bool_col), ANY_VALUE(timestamp_col) FROM {tbl} WHERE int_col = 175", + "outputs": [[175, 175, 1.75, 1.75, "c", true, "2020-01-02 03:32:12.0"]] + }, + { + "sql": "SELECT int_col, ANY_VALUE(string_col), ANY_VALUE(bool_col), ANY_VALUE(timestamp_col) FROM {tbl} WHERE int_col IN (3, 175) GROUP BY int_col", + "outputs": [[3, "b", false, "2020-01-01 08:32:12.0"], [175, "c", true, "2020-01-02 03:32:12.0"]] + }, + { + "sql": "SELECT /*+ aggOptions(is_skip_leaf_stage_aggregate='true') */ int_col, ANY_VALUE(string_col), ANY_VALUE(bool_col), ANY_VALUE(timestamp_col) FROM {tbl} WHERE int_col IN (3, 175) GROUP BY int_col", + "outputs": [[3, "b", false, "2020-01-01 08:32:12.0"], [175, "c", true, "2020-01-02 03:32:12.0"]] + }, + { + "sql": "SELECT ANY_VALUE(bool_col), ANY_VALUE(timestamp_col) FROM {tbl} WHERE int_col < 0", + "outputs": [[null, null]] + }, + { + "sql": "SET inferAnyValueType=false; SELECT int_col, ANY_VALUE(bool_col), ANY_VALUE(timestamp_col) FROM {tbl} WHERE int_col = 175 GROUP BY int_col", + "outputs": [[175, true, "2020-01-02 03:32:12.0"]] + }, { "sql": "SELECT SUMPRECISION(decimal_col) FROM {tbl}", "outputs": [["10000000000100000000110000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000.00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000 [...] @@ -127,6 +148,38 @@ } ] }, + "pinot_any_value_tests": { + "tables": { + "tbl": { + "schema": [ + {"name": "int_col", "type": "INT"}, + {"name": "decimal_col", "type": "BIG_DECIMAL"}, + {"name": "bytes_col", "type": "BYTES"}, + {"name": "uuid_col", "type": "UUID"} + ], + "inputs": [ + [1, "1.5", "DEADBEEF", "123e4567-e89b-12d3-a456-426614174000"], + [1, "1.5", "DEADBEEF", "123e4567-e89b-12d3-a456-426614174000"], + [2, "2.5", "CAFE", "123e4567-e89b-12d3-a456-426614174001"] + ] + } + }, + "queries": [ + { + "comment": "ANY_VALUE over the types that only ANY_VALUE infers a data type for; the test harness compares UUID values as strings", + "sql": "SELECT ANY_VALUE(decimal_col), ANY_VALUE(bytes_col), CAST(ANY_VALUE(uuid_col) AS VARCHAR), ANY_VALUE(int_col * 1.5) FROM {tbl} WHERE int_col = 1", + "outputs": [["1.5", "deadbeef", "123e4567-e89b-12d3-a456-426614174000", 1.5]] + }, + { + "sql": "SELECT int_col, ANY_VALUE(decimal_col), ANY_VALUE(bytes_col), CAST(ANY_VALUE(uuid_col) AS VARCHAR), ANY_VALUE(int_col * 1.5) FROM {tbl} GROUP BY int_col", + "outputs": [[1, "1.5", "deadbeef", "123e4567-e89b-12d3-a456-426614174000", 1.5], [2, "2.5", "cafe", "123e4567-e89b-12d3-a456-426614174001", 3.0]] + }, + { + "sql": "SELECT /*+ aggOptions(is_skip_leaf_stage_aggregate='true') */ int_col, ANY_VALUE(decimal_col), ANY_VALUE(bytes_col), CAST(ANY_VALUE(uuid_col) AS VARCHAR), ANY_VALUE(int_col * 1.5) FROM {tbl} GROUP BY int_col", + "outputs": [[1, "1.5", "deadbeef", "123e4567-e89b-12d3-a456-426614174000", 1.5], [2, "2.5", "cafe", "123e4567-e89b-12d3-a456-426614174001", 3.0]] + } + ] + }, "pinot_percentile_agg_udf_tests": { "tables": { "tbl": { diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java index bbd582c6cb0..a76e562fa18 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java @@ -20,12 +20,14 @@ package org.apache.pinot.segment.spi; import java.util.Arrays; import java.util.List; +import java.util.Map; import java.util.Set; import java.util.stream.Collectors; import java.util.stream.Stream; import javax.annotation.Nullable; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.sql.SqlCallBinding; import org.apache.calcite.sql.SqlKind; import org.apache.calcite.sql.SqlOperatorBinding; import org.apache.calcite.sql.type.OperandTypes; @@ -64,7 +66,9 @@ public enum AggregationFunctionType { SUMPRECISION("sumPrecision", ReturnTypes.explicit(SqlTypeName.DECIMAL), OperandTypes.ANY, SqlTypeName.OTHER), AVG("avg", SqlTypeName.OTHER, SqlTypeName.DOUBLE), MODE("mode", SqlTypeName.OTHER, SqlTypeName.DOUBLE), - ANYVALUE("anyValue", ReturnTypes.ARG0, OperandTypes.ANY, SqlTypeName.OTHER), + ANYVALUE("anyValue", new AnyValueReturnTypeInference(), + OperandTypes.or(OperandTypes.ANY, OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER)), + SqlTypeName.OTHER), FIRSTWITHTIME("firstWithTime", new ValueWithTimeReturnTypeInference(), OperandTypes.or( OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.ANY), OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER)), SqlTypeName.OTHER), @@ -447,6 +451,40 @@ public enum AggregationFunctionType { } } + /// Returns the type of the value argument, which a data type argument must match: the multi-stage engine types the + /// result after the value argument, so a value of another type would not fit the result. + private static class AnyValueReturnTypeInference implements SqlReturnTypeInference { + // DECIMAL values are read as INT, LONG, DOUBLE or BIG_DECIMAL depending on their precision and scale + private static final Map<String, Set<SqlTypeName>> SQL_TYPES_BY_DATA_TYPE = Map.of( + "INT", Set.of(SqlTypeName.TINYINT, SqlTypeName.SMALLINT, SqlTypeName.INTEGER, SqlTypeName.UTINYINT, + SqlTypeName.USMALLINT, SqlTypeName.DECIMAL), + "LONG", Set.of(SqlTypeName.BIGINT, SqlTypeName.UINTEGER, SqlTypeName.DECIMAL), + "FLOAT", Set.of(SqlTypeName.REAL, SqlTypeName.FLOAT), + "DOUBLE", Set.of(SqlTypeName.DOUBLE, SqlTypeName.DECIMAL), + "BIG_DECIMAL", Set.of(SqlTypeName.DECIMAL), + "BOOLEAN", Set.of(SqlTypeName.BOOLEAN), + "TIMESTAMP", Set.of(SqlTypeName.DATE, SqlTypeName.TIME, SqlTypeName.TIMESTAMP), + "STRING", Set.of(SqlTypeName.CHAR, SqlTypeName.VARCHAR), + "BYTES", Set.of(SqlTypeName.BINARY, SqlTypeName.VARBINARY), + "UUID", Set.of(SqlTypeName.UUID)); + + @Override + public RelDataType inferReturnType(SqlOperatorBinding opBinding) { + // Checked during validation, where the literal is visible; a rewritten call refers to it through its input. + if (opBinding instanceof SqlCallBinding && opBinding.getOperandCount() == 2 + && opBinding.isOperandLiteral(1, false)) { + String dataType = StringUtils.upperCase(opBinding.getOperandLiteralValue(1, String.class)); + SqlTypeName valueType = opBinding.getOperandType(0).getSqlTypeName(); + Set<SqlTypeName> sqlTypes = SQL_TYPES_BY_DATA_TYPE.get(dataType); + if (sqlTypes == null || !sqlTypes.contains(valueType)) { + throw new IllegalArgumentException(String.format("Unsupported data type for %s(%s): %s", + opBinding.getOperator().getName(), valueType, dataType)); + } + } + return ReturnTypes.ARG0.inferReturnType(opBinding); + } + } + private static class ArrayReturnTypeInference implements SqlReturnTypeInference { final SqlTypeName _sqlTypeName; diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 6fc50e1c84d..e5ee63f760d 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -1216,6 +1216,9 @@ public class CommonConstants { // MAX(stringCol) -> MAXSTRING(stringCol) // SUM(intCol) -> SUMINT(intCol) public static final String AUTO_REWRITE_AGGREGATION_TYPE = "autoRewriteAggregationType"; + /// Whether ANY_VALUE(expr) is rewritten to ANY_VALUE(expr, 'dataType') to return the type of `expr`. Defaults + /// to true; set to false to keep the legacy result type, which is STRING in the single-stage engine. + public static final String INFER_ANY_VALUE_TYPE = "inferAnyValueType"; // When enabled, allows multi cluster/federated queries to be executed. public static final String ENABLE_MULTI_CLUSTER_ROUTING = "enableMultiClusterRouting"; @@ -1293,6 +1296,9 @@ public class CommonConstants { public static final String AGGREGATE_FUNCTION_REWRITE = "AggregateFunctionRewrite"; /// Supplies the data type argument of FIRST_WITH_TIME / LAST_WITH_TIME(dataColumn, timeColumn). public static final String AGGREGATE_TYPE_ARGUMENT_REWRITE = "AggregateTypeArgumentRewrite"; + /// Supplies the data type argument of ANY_VALUE(expr), unless [Request.QueryOptionKey#INFER_ANY_VALUE_TYPE] is + /// false. + public static final String AGGREGATE_ANY_VALUE_TYPE_REWRITE = "AggregateAnyValueTypeRewrite"; /// Inert unless [CommonConstants.Broker#USE_APPROXIMATE_FUNCTION] is on; name it here to switch it off entirely. public static final String APPROXIMATE_AGGREGATE_REWRITE = "ApproximateAggregateRewrite"; public static final String AGGREGATE_CASE_TO_FILTER = "AggregateCaseToFilter"; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
