Jackie-Jiang commented on code in PR #19552: URL: https://github.com/apache/pinot/pull/19552#discussion_r4264346432
########## pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/GenerateArrayScalarFunction.java: ########## @@ -0,0 +1,187 @@ +/** + * 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.common.function.scalar.array; + +import java.lang.reflect.Method; +import java.util.EnumMap; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.sql.SqlOperatorBinding; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.SqlTypeFamily; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.pinot.common.function.FunctionInfo; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.scalar.ArrayFunctions; +import org.apache.pinot.common.function.sql.PinotSqlFunction; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.spi.annotations.FunctionVolatility; +import org.apache.pinot.spi.annotations.ScalarFunction; + + +/// `generateArray(start, end[, inc])` in the multi-stage engine, generating the sequence `start, start + inc, ...` +/// bounded by `end` inclusive. The single-stage engine resolves the same name through +/// `GenerateArrayTransformFunction`; both delegate to [ArrayFunctions] so the two engines agree on the sequence and +/// on which arguments are rejected. +/// +/// The element type is the widest of the argument types rather than a fixed one, so `generateArray(0, 4)` generates +/// `INT`s while `generateArray(1633078800000, 1633089600000, 1800000)` generates `LONG`s and `generateArray(0, 1, +/// 0.25)` generates `DOUBLE`s. This class exists because that dispatch is on argument types, which the method level +/// [ScalarFunction] annotation cannot express -- it dispatches on argument count alone. +/// +/// This class is stateless and thread safe. +@ScalarFunction +public class GenerateArrayScalarFunction implements PinotScalarFunction { + public static final String NAME = "generateArray"; + + /// The element types a sequence can be generated as, narrowest first. The widest of the argument types wins, so a + /// single fractional or wide argument decides the whole sequence. + private static final List<ColumnDataType> ELEMENT_TYPES = + List.of(ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.FLOAT, ColumnDataType.DOUBLE); + private static final List<SqlTypeName> ELEMENT_TYPE_NAMES = Review Comment: Optional cleanup: remove the duplicate SQL type mapping `ELEMENT_TYPE_NAMES` duplicates the mapping already provided by `ColumnDataType.toType()`. Return-type construction can use `ELEMENT_TYPES.get(width).toType(typeFactory)` instead, removing the parallel list and the requirement to keep its indices synchronized with `ELEMENT_TYPES`. While addressing the inference issue above, it would also help to consolidate argument normalization so planning and execution share the same policy. ########## pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/GenerateArrayScalarFunction.java: ########## @@ -0,0 +1,187 @@ +/** + * 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.common.function.scalar.array; + +import java.lang.reflect.Method; +import java.util.EnumMap; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.sql.SqlOperatorBinding; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.SqlTypeFamily; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.pinot.common.function.FunctionInfo; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.scalar.ArrayFunctions; +import org.apache.pinot.common.function.sql.PinotSqlFunction; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.spi.annotations.FunctionVolatility; +import org.apache.pinot.spi.annotations.ScalarFunction; + + +/// `generateArray(start, end[, inc])` in the multi-stage engine, generating the sequence `start, start + inc, ...` +/// bounded by `end` inclusive. The single-stage engine resolves the same name through +/// `GenerateArrayTransformFunction`; both delegate to [ArrayFunctions] so the two engines agree on the sequence and +/// on which arguments are rejected. +/// +/// The element type is the widest of the argument types rather than a fixed one, so `generateArray(0, 4)` generates +/// `INT`s while `generateArray(1633078800000, 1633089600000, 1800000)` generates `LONG`s and `generateArray(0, 1, +/// 0.25)` generates `DOUBLE`s. This class exists because that dispatch is on argument types, which the method level +/// [ScalarFunction] annotation cannot express -- it dispatches on argument count alone. +/// +/// This class is stateless and thread safe. +@ScalarFunction +public class GenerateArrayScalarFunction implements PinotScalarFunction { + public static final String NAME = "generateArray"; + + /// The element types a sequence can be generated as, narrowest first. The widest of the argument types wins, so a + /// single fractional or wide argument decides the whole sequence. + private static final List<ColumnDataType> ELEMENT_TYPES = + List.of(ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.FLOAT, ColumnDataType.DOUBLE); + private static final List<SqlTypeName> ELEMENT_TYPE_NAMES = + List.of(SqlTypeName.INTEGER, SqlTypeName.BIGINT, SqlTypeName.REAL, SqlTypeName.DOUBLE); + + /// Element type to the [ArrayFunctions] method generating it, by argument count. + private static final Map<ColumnDataType, Map<Integer, FunctionInfo>> FUNCTION_INFO_MAP = + new EnumMap<>(ColumnDataType.class); + + static { + register(ColumnDataType.INT, "generateIntArray", int.class); + register(ColumnDataType.LONG, "generateLongArray", long.class); + register(ColumnDataType.FLOAT, "generateFloatArray", float.class); + register(ColumnDataType.DOUBLE, "generateDoubleArray", double.class); + } + + private static void register(ColumnDataType elementType, String methodName, Class<?> argumentClass) { + try { + FUNCTION_INFO_MAP.put(elementType, Map.of(2, + functionInfo(ArrayFunctions.class.getMethod(methodName, argumentClass, argumentClass)), 3, + functionInfo(ArrayFunctions.class.getMethod(methodName, argumentClass, argumentClass, argumentClass)))); + } catch (NoSuchMethodException e) { + throw new IllegalStateException("Failed to find the " + methodName + " implementation of " + NAME, e); + } + } + + /// Marks the function as not deterministic so that neither engine folds a call on literal arguments into an array + /// literal at the broker. Folding would serialize the whole sequence, up to the maximum length [ArrayFunctions] + /// allows, into the query plan sent to every server, while evaluating it on the server only needs the three + /// arguments. The function is still [FunctionVolatility#IMMUTABLE]: the same arguments always generate the same + /// sequence, so it stays usable in ingestion transforms. + private static FunctionInfo functionInfo(Method method) { + return new FunctionInfo(method, ArrayFunctions.class, false, false, FunctionVolatility.IMMUTABLE); + } + + @Override + public String getName() { + return NAME; + } + + @Override + public PinotSqlFunction toPinotSqlFunction() { + // The third operand is the increment, which defaults to 1 (or -1 when counting down) when omitted. + return new PinotSqlFunction(NAME, GenerateArrayScalarFunction::inferReturnType, + OperandTypes.family(List.of(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC), i -> i == 2)); + } + + @Nullable + @Override + public FunctionInfo getFunctionInfo(ColumnDataType[] argumentTypes) { + ColumnDataType elementType = getElementType(argumentTypes); + return elementType != null ? getFunctionInfo(elementType, argumentTypes.length) : null; + } + + /// The element type of the sequence generated from arguments of these types: `INT`, `LONG`, `FLOAT` or `DOUBLE`, + /// whichever is the widest of the argument types. Returns `null` if a sequence cannot be generated from one of them. + /// + /// The single-stage `GenerateArrayTransformFunction` uses this too. It also evaluates the multi-stage leaf stage + /// projections, so it must generate the same element type as the type the multi-stage planner inferred. + @Nullable + public static ColumnDataType getElementType(ColumnDataType[] argumentTypes) { + int width = 0; + for (ColumnDataType argumentType : argumentTypes) { + int argumentWidth = width(argumentType); + if (argumentWidth < 0) { + return null; + } + width = Math.max(width, argumentWidth); + } + return ELEMENT_TYPES.get(width); + } + + @Nullable + @Override + public FunctionInfo getFunctionInfo(int numArguments) { + // Reached only when the argument types are unknown. LONG is the type the time grids this mainly generates use. + return getFunctionInfo(ColumnDataType.LONG, numArguments); + } + + @Nullable + private static FunctionInfo getFunctionInfo(ColumnDataType elementType, int numArguments) { + return FUNCTION_INFO_MAP.get(elementType).get(numArguments); + } + + private static RelDataType inferReturnType(SqlOperatorBinding binding) { + int width = 0; + for (RelDataType operandType : binding.collectOperandTypes()) { + width = Math.max(width, width(operandType.getSqlTypeName())); + } + RelDataTypeFactory typeFactory = binding.getTypeFactory(); + return typeFactory.createArrayType(typeFactory.createSqlType(ELEMENT_TYPE_NAMES.get(width)), -1); Review Comment: [P2] Preserve operand nullability in the inferred array type `createArrayType()` produces a nonnullable array, but this function's `FunctionInfo` is null-intolerant, so execution returns null when any argument is null. For example, `SELECT g.x FROM UNNEST(COALESCE(generateArray(CAST(NULL AS INTEGER), 4), ARRAY[7])) AS g(x)` loses the fallback during planning: Calcite's `RelBuilder` simplifies `COALESCE` by dropping operands after the first nonnullable one. The resulting UNNEST emits no row instead of `7`. Please propagate operand nullability to the inferred array type and cover this through the query pipeline. ########## pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java: ########## @@ -362,50 +368,130 @@ public static Object arrayValueConstructor(Object... arr) { return arr; } + /// Generates the sequence `start, start + inc, ...`, stopping at the last value that does not pass `end`. Both + /// bounds are inclusive when `end` lands exactly on a step, and `inc` may be negative to count down. + /// + /// The main use is building a dense time grid to `UNNEST` and left join a sparse time series onto, which is how gap + /// filling is expressed in the multi-stage engine: + /// ```sql + /// SELECT ts FROM UNNEST(generateArray(1633078800000, 1633089600000, 1800000)) AS grid(ts) + /// ``` + /// `generateArray` selects the variant to call from its argument types; see + /// [org.apache.pinot.common.function.scalar.array.GenerateArrayScalarFunction]. + /// + /// @throws IllegalArgumentException if `inc` is zero, if its sign does not lead from `start` to `end`, or if the + /// sequence would hold more than [#MAX_GENERATED_ARRAY_LENGTH] elements @ScalarFunction - public static int[] generateIntArray(int start, int end, int inc) { - int size = (end - start) / inc + 1; - int[] arr = new int[size]; - - for (int i = 0, value = start; i < size; i++, value += inc) { - arr[i] = value; + public static long[] generateLongArray(long start, long end, long inc) { + long[] arr = new long[generatedArrayLength(start, end, inc)]; + for (int i = 0; i < arr.length; i++) { + arr[i] = start + i * inc; } return arr; } + /// Same as [#generateLongArray(long, long, long)] with a step of `1`, or `-1` when `end` is below `start`. @ScalarFunction - public static long[] generateLongArray(long start, long end, long inc) { - int size = (int) ((end - start) / inc + 1); - long[] arr = new long[size]; + public static long[] generateLongArray(long start, long end) { + return generateLongArray(start, end, end < start ? -1 : 1); + } - for (int i = 0; i < size; i++, start += inc) { - arr[i] = start; + @ScalarFunction + public static int[] generateIntArray(int start, int end, int inc) { + int[] arr = new int[generatedArrayLength(start, end, inc)]; + for (int i = 0; i < arr.length; i++) { + arr[i] = start + i * inc; } return arr; } @ScalarFunction - public static float[] generateFloatArray(float start, float end, float inc) { - int size = (int) ((end - start) / inc + 1); - float[] arr = new float[size]; + public static int[] generateIntArray(int start, int end) { + return generateIntArray(start, end, end < start ? -1 : 1); + } - for (int i = 0; i < size; i++, start += inc) { - arr[i] = start; + @ScalarFunction + public static float[] generateFloatArray(float start, float end, float inc) { + float[] arr = new float[generatedArrayLength(start, end, inc)]; + for (int i = 0; i < arr.length; i++) { + arr[i] = start + i * inc; } return arr; } @ScalarFunction - public static double[] generateDoubleArray(double start, double end, double inc) { - int size = (int) ((end - start) / inc + 1); - double[] arr = new double[size]; + public static float[] generateFloatArray(float start, float end) { + return generateFloatArray(start, end, end < start ? -1 : 1); + } - for (int i = 0; i < size; i++, start += inc) { - arr[i] = start; + @ScalarFunction + public static double[] generateDoubleArray(double start, double end, double inc) { + double[] arr = new double[generatedArrayLength(start, end, inc)]; + for (int i = 0; i < arr.length; i++) { + arr[i] = start + i * inc; } return arr; } + @ScalarFunction + public static double[] generateDoubleArray(double start, double end) { + return generateDoubleArray(start, end, end < start ? -1 : 1); + } + + /// Length of the integral sequence `start, start + inc, ...` bounded by `end`, inclusive. + private static int generatedArrayLength(long start, long end, long inc) { + Preconditions.checkArgument(inc != 0, "Increment must not be zero"); + long span; + try { + span = Math.subtractExact(end, start); + } catch (ArithmeticException e) { + // The span does not fit in a long, so the sequence is far past the length limit whatever the increment is. + throw new IllegalArgumentException( + String.format("Range from %s to %s is too wide to generate an array from", start, end)); + } + checkIncrementLeadsToEnd(Long.signum(span), Long.signum(inc), start, end, inc); + return checkGeneratedArrayLength(span / inc); + } + + /// Length of the `float` sequence `start, start + inc, ...` bounded by `end`, inclusive. The step count is computed + /// in `float` arithmetic, as the single-stage engine always did. + private static int generatedArrayLength(float start, float end, float inc) { + checkFloatingPointArguments(start, end, inc); + // TODO: decimal steps such as 0.1 are not exact in binary, so this can drop the end; see apache/pinot#19814 + return checkGeneratedArrayLength((long) ((end - start) / inc + 1) - 1); + } + + /// Length of the `double` sequence `start, start + inc, ...` bounded by `end`, inclusive. + private static int generatedArrayLength(double start, double end, double inc) { + checkFloatingPointArguments(start, end, inc); + // Adding 1 before truncating matches the length the single-stage engine always computed. It absorbs part of the + // rounding error of decimal steps, so 0.1 to 0.3 by 0.1 keeps its end. + // TODO: it does not absorb all of it, so 0.1 to 0.7 by 0.1 still drops its end; see apache/pinot#19814 + // A span too wide for a long saturates the cast at Long.MAX_VALUE, which the length check below rejects. + return checkGeneratedArrayLength((long) ((end - start) / inc + 1) - 1); + } + + private static void checkFloatingPointArguments(double start, double end, double inc) { + Preconditions.checkArgument(inc != 0 && !Double.isNaN(inc), "Increment must not be zero or NaN, got: %s", inc); Review Comment: [P2] Reject infinite increments before using index multiplication This accepts infinite increments. `generateDoubleArray(0, 1, Double.POSITIVE_INFINITY)` therefore allocates one element, but the new `start + i * inc` loop computes `0 * Infinity` at index zero and returns `[NaN]`; the previous implementation returned `[0.0]`. FLOAT and descending ranges with negative infinity have the same regression. Please reject nonfinite increments, or preserve the initial value without multiplying infinity by zero, and cover both floating-point variants. ########## pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/GenerateArrayScalarFunction.java: ########## @@ -0,0 +1,187 @@ +/** + * 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.common.function.scalar.array; + +import java.lang.reflect.Method; +import java.util.EnumMap; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.sql.SqlOperatorBinding; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.SqlTypeFamily; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.pinot.common.function.FunctionInfo; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.scalar.ArrayFunctions; +import org.apache.pinot.common.function.sql.PinotSqlFunction; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.spi.annotations.FunctionVolatility; +import org.apache.pinot.spi.annotations.ScalarFunction; + + +/// `generateArray(start, end[, inc])` in the multi-stage engine, generating the sequence `start, start + inc, ...` +/// bounded by `end` inclusive. The single-stage engine resolves the same name through +/// `GenerateArrayTransformFunction`; both delegate to [ArrayFunctions] so the two engines agree on the sequence and +/// on which arguments are rejected. +/// +/// The element type is the widest of the argument types rather than a fixed one, so `generateArray(0, 4)` generates +/// `INT`s while `generateArray(1633078800000, 1633089600000, 1800000)` generates `LONG`s and `generateArray(0, 1, +/// 0.25)` generates `DOUBLE`s. This class exists because that dispatch is on argument types, which the method level +/// [ScalarFunction] annotation cannot express -- it dispatches on argument count alone. +/// +/// This class is stateless and thread safe. +@ScalarFunction +public class GenerateArrayScalarFunction implements PinotScalarFunction { + public static final String NAME = "generateArray"; + + /// The element types a sequence can be generated as, narrowest first. The widest of the argument types wins, so a + /// single fractional or wide argument decides the whole sequence. + private static final List<ColumnDataType> ELEMENT_TYPES = + List.of(ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.FLOAT, ColumnDataType.DOUBLE); + private static final List<SqlTypeName> ELEMENT_TYPE_NAMES = + List.of(SqlTypeName.INTEGER, SqlTypeName.BIGINT, SqlTypeName.REAL, SqlTypeName.DOUBLE); + + /// Element type to the [ArrayFunctions] method generating it, by argument count. + private static final Map<ColumnDataType, Map<Integer, FunctionInfo>> FUNCTION_INFO_MAP = + new EnumMap<>(ColumnDataType.class); + + static { + register(ColumnDataType.INT, "generateIntArray", int.class); + register(ColumnDataType.LONG, "generateLongArray", long.class); + register(ColumnDataType.FLOAT, "generateFloatArray", float.class); + register(ColumnDataType.DOUBLE, "generateDoubleArray", double.class); + } + + private static void register(ColumnDataType elementType, String methodName, Class<?> argumentClass) { + try { + FUNCTION_INFO_MAP.put(elementType, Map.of(2, + functionInfo(ArrayFunctions.class.getMethod(methodName, argumentClass, argumentClass)), 3, + functionInfo(ArrayFunctions.class.getMethod(methodName, argumentClass, argumentClass, argumentClass)))); + } catch (NoSuchMethodException e) { + throw new IllegalStateException("Failed to find the " + methodName + " implementation of " + NAME, e); + } + } + + /// Marks the function as not deterministic so that neither engine folds a call on literal arguments into an array + /// literal at the broker. Folding would serialize the whole sequence, up to the maximum length [ArrayFunctions] + /// allows, into the query plan sent to every server, while evaluating it on the server only needs the three + /// arguments. The function is still [FunctionVolatility#IMMUTABLE]: the same arguments always generate the same + /// sequence, so it stays usable in ingestion transforms. + private static FunctionInfo functionInfo(Method method) { + return new FunctionInfo(method, ArrayFunctions.class, false, false, FunctionVolatility.IMMUTABLE); + } + + @Override + public String getName() { + return NAME; + } + + @Override + public PinotSqlFunction toPinotSqlFunction() { + // The third operand is the increment, which defaults to 1 (or -1 when counting down) when omitted. + return new PinotSqlFunction(NAME, GenerateArrayScalarFunction::inferReturnType, + OperandTypes.family(List.of(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC), i -> i == 2)); + } + + @Nullable + @Override + public FunctionInfo getFunctionInfo(ColumnDataType[] argumentTypes) { + ColumnDataType elementType = getElementType(argumentTypes); + return elementType != null ? getFunctionInfo(elementType, argumentTypes.length) : null; + } + + /// The element type of the sequence generated from arguments of these types: `INT`, `LONG`, `FLOAT` or `DOUBLE`, + /// whichever is the widest of the argument types. Returns `null` if a sequence cannot be generated from one of them. + /// + /// The single-stage `GenerateArrayTransformFunction` uses this too. It also evaluates the multi-stage leaf stage + /// projections, so it must generate the same element type as the type the multi-stage planner inferred. + @Nullable + public static ColumnDataType getElementType(ColumnDataType[] argumentTypes) { + int width = 0; + for (ColumnDataType argumentType : argumentTypes) { + int argumentWidth = width(argumentType); + if (argumentWidth < 0) { + return null; + } + width = Math.max(width, argumentWidth); + } + return ELEMENT_TYPES.get(width); + } + + @Nullable + @Override + public FunctionInfo getFunctionInfo(int numArguments) { + // Reached only when the argument types are unknown. LONG is the type the time grids this mainly generates use. + return getFunctionInfo(ColumnDataType.LONG, numArguments); + } + + @Nullable + private static FunctionInfo getFunctionInfo(ColumnDataType elementType, int numArguments) { + return FUNCTION_INFO_MAP.get(elementType).get(numArguments); + } + + private static RelDataType inferReturnType(SqlOperatorBinding binding) { + int width = 0; + for (RelDataType operandType : binding.collectOperandTypes()) { + width = Math.max(width, width(operandType.getSqlTypeName())); Review Comment: [P2] Align return-type inference with normalized argument types `width(SqlTypeName)` treats every DECIMAL (and unsigned integer type) as DOUBLE, whereas `RelToPlanNodeConverter` normalizes scale-zero decimals and supported unsigned integers to INT/LONG. For `SELECT generateArray(CAST(0 AS DECIMAL(10,0)), 4)`, the outer call retains its inferred DOUBLE_ARRAY type after its argument becomes INT, so runtime dispatch generates `int[]`. `TypeUtils.convert()` does not widen primitive numeric arrays, and `DataBlockBuilder` then casts the value to `double[]`. Please infer from the full `RelDataType` using normalization consistent with runtime dispatch, and add a query regression case. ########## pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java: ########## @@ -230,17 +231,13 @@ private static RexNode evaluateLiteralOnlyFunction(RexCall rexCall, RexBuilder r } try { if (rexNodeType instanceof ArraySqlType) { - List<Object> resultValues = new ArrayList<>(); - - // SQL FLOAT and DOUBLE literals are represented as Java double - if (resultValue instanceof double[]) { - for (double value: (double[]) resultValue) { - resultValues.add(convertResultValue(value, rexNodeType.getComponentType())); - } - } else { - for (Object value : (Object[]) resultValue) { - resultValues.add(convertResultValue(value, rexNodeType.getComponentType())); - } + // Scalar functions may return a primitive array (e.g. generateArray returns long[]), which cannot be cast to Review Comment: Optional scope cleanup: consider splitting out the primitive-array folding fix This is a useful fix for the existing typed generators, but the new `generateArray` explicitly disables constant folding, so its MSE support does not depend on this change. The reflective array handling and `testPrimitiveArrayFunctionFoldsIntoAnArrayLiteral` could move together into a separate focused PR. This is a scope suggestion, not a request to drop the fix. ########## pinot-query-runtime/src/test/resources/queries/GenerateArray.json: ########## @@ -0,0 +1,267 @@ +{ + "generate_array_test": { + "tables": { + "sparse": { + "schema": [ + { + "name": "eventTs", + "type": "LONG" + }, + { + "name": "series", + "type": "STRING" + }, + { + "name": "val", + "type": "INT" + } + ], + "inputs": [ + [ + 0, + "a", + 10 + ], + [ + 3, + "a", + 40 + ], + [ + 0, + "b", + 100 + ] + ] + } + }, + "queries": [ + { + "description": "generateArray unnests into a dense grid, with the increment defaulting to 1", + "sql": "SELECT grid.bucket FROM {sparse} CROSS JOIN UNNEST(generateArray(0, 4)) AS grid(bucket) WHERE series = 'b' ORDER BY grid.bucket", + "keepOutputRowOrder": true, + "outputs": [ + [ + 0 + ], + [ + 1 + ], + [ + 2 + ], + [ + 3 + ], + [ + 4 + ] + ], + "ignoreLiteMode": true, + "ignoreV2Optimizer": true + }, + { + "description": "explicit increment", + "sql": "SELECT grid.bucket FROM {sparse} CROSS JOIN UNNEST(generateArray(0, 10, 5)) AS grid(bucket) WHERE series = 'b' ORDER BY grid.bucket", + "keepOutputRowOrder": true, + "outputs": [ + [ + 0 + ], + [ + 5 + ], + [ + 10 + ] + ], + "ignoreLiteMode": true, + "ignoreV2Optimizer": true + }, + { + "description": "negative increment counts down, and the end is not passed when it is not on a step", + "sql": "SELECT grid.bucket FROM {sparse} CROSS JOIN UNNEST(generateArray(10, 1, -4)) AS grid(bucket) WHERE series = 'b' ORDER BY grid.bucket DESC", + "keepOutputRowOrder": true, + "outputs": [ + [ + 10 + ], + [ + 6 + ], + [ + 2 + ] + ], + "ignoreLiteMode": true, + "ignoreV2Optimizer": true + }, + { + "description": "epoch millis bounds generate a LONG sequence rather than overflowing an INT one", + "sql": "SELECT grid.bucket FROM {sparse} CROSS JOIN UNNEST(generateArray(1633078800000, 1633080600000, 1800000)) AS grid(bucket) WHERE series = 'b' ORDER BY grid.bucket", + "keepOutputRowOrder": true, + "outputs": [ + [ + 1633078800000 + ], + [ + 1633080600000 + ] + ], + "ignoreLiteMode": true, + "ignoreV2Optimizer": true + }, + { + "description": "a fractional increment makes the whole sequence DOUBLE", + "sql": "SELECT grid.bucket FROM {sparse} CROSS JOIN UNNEST(generateArray(0, 1, 0.5)) AS grid(bucket) WHERE series = 'b' ORDER BY grid.bucket", + "keepOutputRowOrder": true, + "outputs": [ + [ + 0.0 + ], + [ + 0.5 + ], + [ + 1.0 + ] + ], + "ignoreLiteMode": true, + "ignoreV2Optimizer": true + }, + { + "description": "WITH ORDINALITY numbers the generated buckets", + "sql": "SELECT grid.bucket, grid.pos FROM {sparse} CROSS JOIN UNNEST(generateArray(100, 300, 100)) WITH ORDINALITY AS grid(bucket, pos) WHERE series = 'b' ORDER BY grid.bucket", + "keepOutputRowOrder": true, + "outputs": [ + [ + 100, + 1 + ], + [ + 200, + 2 + ], + [ + 300, + 3 + ] + ], + "ignoreLiteMode": true, + "ignoreV2Optimizer": true + }, + { + "description": "gap fill: a generated grid crossed with the series keys, left joined to the sparse data and carried forward", + "sql": "SELECT grid.bucket, grid.series, COALESCE(LAST_VALUE(s.val) IGNORE NULLS OVER (PARTITION BY grid.series ORDER BY grid.bucket ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW), 0) FROM (SELECT g.bucket, k.series FROM UNNEST(generateArray(0, 4)) AS g(bucket) CROSS JOIN (SELECT DISTINCT series FROM {sparse}) AS k) AS grid LEFT JOIN {sparse} s ON s.eventTs = grid.bucket AND s.series = grid.series ORDER BY grid.series, grid.bucket", + "keepOutputRowOrder": true, + "outputs": [ + [ + 0, + "a", + 10 + ], + [ + 1, + "a", + 10 + ], + [ + 2, + "a", + 10 + ], + [ + 3, + "a", + 40 + ], + [ + 4, + "a", + 40 + ], + [ + 0, + "b", + 100 + ], + [ + 1, + "b", + 100 + ], + [ + 2, + "b", + 100 + ], + [ + 3, + "b", + 100 + ], + [ + 4, + "b", + 100 + ] + ], + "ignoreLiteMode": true, + "ignoreV2Optimizer": true + } + ] + }, + "generate_array_unusable_arguments": { Review Comment: Optional cleanup: reuse the existing fixture for the error cases This group repeats the first group's entire `sparse` schema and three input rows. Since `expectedException` is configured per query, the two error cases can go in the first group's `queries` array and this duplicate fixture can be removed without losing coverage. ########## pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/GenerateArrayScalarFunction.java: ########## @@ -0,0 +1,187 @@ +/** + * 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.common.function.scalar.array; + +import java.lang.reflect.Method; +import java.util.EnumMap; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.sql.SqlOperatorBinding; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.SqlTypeFamily; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.pinot.common.function.FunctionInfo; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.scalar.ArrayFunctions; +import org.apache.pinot.common.function.sql.PinotSqlFunction; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.spi.annotations.FunctionVolatility; +import org.apache.pinot.spi.annotations.ScalarFunction; + + +/// `generateArray(start, end[, inc])` in the multi-stage engine, generating the sequence `start, start + inc, ...` +/// bounded by `end` inclusive. The single-stage engine resolves the same name through +/// `GenerateArrayTransformFunction`; both delegate to [ArrayFunctions] so the two engines agree on the sequence and +/// on which arguments are rejected. +/// +/// The element type is the widest of the argument types rather than a fixed one, so `generateArray(0, 4)` generates +/// `INT`s while `generateArray(1633078800000, 1633089600000, 1800000)` generates `LONG`s and `generateArray(0, 1, +/// 0.25)` generates `DOUBLE`s. This class exists because that dispatch is on argument types, which the method level +/// [ScalarFunction] annotation cannot express -- it dispatches on argument count alone. +/// +/// This class is stateless and thread safe. +@ScalarFunction +public class GenerateArrayScalarFunction implements PinotScalarFunction { + public static final String NAME = "generateArray"; + + /// The element types a sequence can be generated as, narrowest first. The widest of the argument types wins, so a + /// single fractional or wide argument decides the whole sequence. + private static final List<ColumnDataType> ELEMENT_TYPES = + List.of(ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.FLOAT, ColumnDataType.DOUBLE); + private static final List<SqlTypeName> ELEMENT_TYPE_NAMES = + List.of(SqlTypeName.INTEGER, SqlTypeName.BIGINT, SqlTypeName.REAL, SqlTypeName.DOUBLE); + + /// Element type to the [ArrayFunctions] method generating it, by argument count. + private static final Map<ColumnDataType, Map<Integer, FunctionInfo>> FUNCTION_INFO_MAP = + new EnumMap<>(ColumnDataType.class); + + static { + register(ColumnDataType.INT, "generateIntArray", int.class); + register(ColumnDataType.LONG, "generateLongArray", long.class); + register(ColumnDataType.FLOAT, "generateFloatArray", float.class); + register(ColumnDataType.DOUBLE, "generateDoubleArray", double.class); + } + + private static void register(ColumnDataType elementType, String methodName, Class<?> argumentClass) { + try { + FUNCTION_INFO_MAP.put(elementType, Map.of(2, + functionInfo(ArrayFunctions.class.getMethod(methodName, argumentClass, argumentClass)), 3, + functionInfo(ArrayFunctions.class.getMethod(methodName, argumentClass, argumentClass, argumentClass)))); + } catch (NoSuchMethodException e) { + throw new IllegalStateException("Failed to find the " + methodName + " implementation of " + NAME, e); + } + } + + /// Marks the function as not deterministic so that neither engine folds a call on literal arguments into an array + /// literal at the broker. Folding would serialize the whole sequence, up to the maximum length [ArrayFunctions] + /// allows, into the query plan sent to every server, while evaluating it on the server only needs the three + /// arguments. The function is still [FunctionVolatility#IMMUTABLE]: the same arguments always generate the same + /// sequence, so it stays usable in ingestion transforms. + private static FunctionInfo functionInfo(Method method) { + return new FunctionInfo(method, ArrayFunctions.class, false, false, FunctionVolatility.IMMUTABLE); + } + + @Override + public String getName() { + return NAME; + } + + @Override + public PinotSqlFunction toPinotSqlFunction() { + // The third operand is the increment, which defaults to 1 (or -1 when counting down) when omitted. + return new PinotSqlFunction(NAME, GenerateArrayScalarFunction::inferReturnType, + OperandTypes.family(List.of(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC), i -> i == 2)); Review Comment: [P2] Handle nonliteral arguments in leaf-stage execution The new signature accepts numeric expressions, but `SELECT generateArray(intCol, intCol + 1) FROM t` is pushed into the SSE leaf projection. `TransformFunctionFactory` unconditionally selects `GenerateArrayTransformFunction` for this name, whose constructor rejects column/function arguments because it only accepts literals. This makes a valid scalar call fail depending on its placement in the plan. Please route nonliteral calls through the scalar implementation, extend the transform to evaluate them, or prevent their leaf pushdown; a table-backed query would cover the gap. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
