You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xi...@apache.org on 2026/10/21 01:50:53 UTC
(pinot) branch xiangfu0/agg-rewrite/04-any-value created (now d9ccb8a7906)
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a change to branch xiangfu0/agg-rewrite/04-any-value
in repository https://gitbox.apache.org/repos/asf/pinot.git
at d9ccb8a7906 Return the input's logical type from ANY_VALUE
This branch includes the following new commits:
new d9ccb8a7906 Return the input's logical type from ANY_VALUE
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
(pinot) 01/01: Return the input's logical type from ANY_VALUE
Posted by xi...@apache.org.
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 <xi...@apache.org>
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: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org