gortiz commented on code in PR #19552: URL: https://github.com/apache/pinot/pull/19552#discussion_r4280856950
########## 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: Good catch. The inferred array type now allows null when any operand does, so Calcite keeps the `COALESCE` fallback. `QueryCompilationTest` checks the type and the plan, and `GenerateArray.json` runs your query end to end. ########## 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: Fixed. The return type inference now maps each operand the way `RelToPlanNodeConverter.convertToColumnDataType` does, so scale zero decimals become INT or LONG and the unsigned types become INT or LONG. Then it calls the same `getElementType` that execution uses. The converter is in `pinot-query-planner`, which `pinot-common` cannot depend on, so I mirrored its numeric cases and left a comment that points to it. There are new `QueryCompilationTest` cases, and `GenerateArray.json` runs your `DECIMAL(10, 0)` query. ########## 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: Fixed in `TransformFunctionFactory`. It uses `GenerateArrayTransformFunction` only when every argument is a literal. Any other call falls through to the scalar function through `ScalarTransformFunctionWrapper`, which generates one array per row. Your example then hit a second problem. In the leaf stage, `intCol + 1` goes through the single-stage `plus`, which returns DOUBLE, so the call generates a `double[]` while the plan expects an INT_ARRAY. `TypeUtils.convert` already converts single values in this case, but for arrays it only asserted. It now converts between primitive numeric arrays the same way it converts single values. `ScalarTransformFunctionWrapperTest` and the new `TypeUtilsTest` cover both parts, and `GenerateArray.json` has a table-backed case with `generateArray(val, val + 2)`. ########## 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: Fixed. A non-finite increment is now rejected, and the tests cover double and float in both directions. ########## 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: Done. `ELEMENT_TYPE_NAMES` is gone, and the return type comes from `ColumnDataType.toType`. Planning and execution now both pick the element type through `getElementType`. ########## 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: Done, the error cases are in the first group now. They need an empty `outputs`, because the runner does not allow a group that mixes queries with and without expected outputs. ########## 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: I'd rather keep it here, if that's fine with you. Without it, folding `generateIntArray`, `generateLongArray` or `generateFloatArray` on literals fails at planning time, so the MSE story of this PR would still be half broken for the typed functions. The change is small and `testPrimitiveArrayFunctionFoldsIntoAnArrayLiteral` covers it. If you feel strongly about it, I'll move it to its own PR. -- 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]
