You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by go...@apache.org on 2026/10/15 13:05:39 UTC
(pinot) branch master updated: Support generateArray in the multi-stage engine (#19552)
This is an automated email from the ASF dual-hosted git repository.
gortiz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 19fdcc0b9da Support generateArray in the multi-stage engine (#19552)
19fdcc0b9da is described below
commit 19fdcc0b9da15c9ea35217b1f6807d9edb226b23
Author: Gonzalo Ortiz Jaureguizar <go...@users.noreply.github.com>
AuthorDate: Thu Oct 15 15:05:32 2026 +0200
Support generateArray in the multi-stage engine (#19552)
---
.../common/function/scalar/ArrayFunctions.java | 130 ++++++++--
.../scalar/array/GenerateArrayScalarFunction.java | 206 +++++++++++++++
.../scalar/GenerateArrayFunctionsTest.java | 156 ++++++++++++
.../function/GenerateArrayTransformFunction.java | 118 +++------
.../function/TransformFunctionFactory.java | 6 +-
.../GenerateArrayTransformFunctionTest.java | 28 +-
.../ScalarTransformFunctionWrapperTest.java | 26 ++
.../resources/udf-test-results/all-functions.yaml | 16 +-
.../rel/rules/PinotEvaluateLiteralRule.java | 19 +-
.../apache/pinot/query/QueryCompilationTest.java | 63 +++++
.../query/runtime/operator/utils/TypeUtils.java | 54 ++++
.../runtime/operator/utils/TypeUtilsTest.java | 21 ++
.../src/test/resources/queries/GenerateArray.json | 282 +++++++++++++++++++++
13 files changed, 1007 insertions(+), 118 deletions(-)
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java
index a6db5b7906f..67643d17dea 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.common.function.scalar;
+import com.google.common.base.Preconditions;
import it.unimi.dsi.fastutil.ints.IntArrayList;
import it.unimi.dsi.fastutil.ints.IntLinkedOpenHashSet;
import it.unimi.dsi.fastutil.ints.IntSet;
@@ -32,6 +33,11 @@ import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder;
/// Inbuilt array scalar functions. See [ArrayUtils] for details.
public class ArrayFunctions {
+ /// Maximum number of elements the `generate*Array` functions may produce. Calls on literal arguments are folded
+ /// into the query plan as an array literal, so an unbounded range would exhaust broker memory rather than fail
+ /// cleanly. 100000 covers a minute resolution grid over 69 days, which is well past any practical time grid.
+ private static final int MAX_GENERATED_ARRAY_LENGTH = 100_000;
+
private ArrayFunctions() {
}
@@ -362,50 +368,132 @@ public class ArrayFunctions {
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) {
+ // An infinite increment would turn the first element, start + 0 * inc, into NaN.
+ Preconditions.checkArgument(inc != 0 && Double.isFinite(inc), "Increment must be finite and not zero, got: %s",
+ inc);
+ Preconditions.checkArgument(Double.isFinite(start) && Double.isFinite(end),
+ "Range from %s to %s must be finite", start, end);
+ checkIncrementLeadsToEnd((int) Math.signum(end - start), (int) Math.signum(inc), start, end, inc);
+ }
+
+ private static void checkIncrementLeadsToEnd(int spanSignum, int incSignum, Number start, Number end, Number inc) {
+ // A zero span yields the single element sequence [start], so any non-zero increment is fine there.
+ Preconditions.checkArgument(spanSignum == 0 || spanSignum == incSignum,
+ "Increment: %s does not lead from start: %s to end: %s", inc, start, end);
+ }
+
+ /// Checks the length of a sequence taking `steps` increments to reach its end, and returns it. The check is on the
+ /// step count rather than on the length so that a span of `Long.MAX_VALUE` cannot overflow into a small length.
+ private static int checkGeneratedArrayLength(long steps) {
+ Preconditions.checkArgument(steps < MAX_GENERATED_ARRAY_LENGTH,
+ "Generating more than %s elements exceeds the maximum of %s", steps, MAX_GENERATED_ARRAY_LENGTH);
+ return (int) steps + 1;
+ }
+
@ScalarFunction
public static String arrayToString(String[] values, String delimiter) {
return String.join(delimiter, values);
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/GenerateArrayScalarFunction.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/GenerateArrayScalarFunction.java
new file mode 100644
index 00000000000..a70e2f4a9a6
--- /dev/null
+++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/GenerateArrayScalarFunction.java
@@ -0,0 +1,206 @@
+/**
+ * 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.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);
+
+ /// 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) {
+ List<RelDataType> operandTypes = binding.collectOperandTypes();
+ ColumnDataType[] argumentTypes = new ColumnDataType[operandTypes.size()];
+ boolean nullable = false;
+ for (int i = 0; i < argumentTypes.length; i++) {
+ RelDataType operandType = operandTypes.get(i);
+ argumentTypes[i] = toArgumentType(operandType);
+ nullable |= operandType.isNullable();
+ }
+ ColumnDataType elementType = getElementType(argumentTypes);
+ assert elementType != null : "The operand type checker accepts only numeric operands";
+ RelDataTypeFactory typeFactory = binding.getTypeFactory();
+ RelDataType arrayType = typeFactory.createArrayType(elementType.toType(typeFactory), -1);
+ // The function returns null when any argument is null. Calcite must know it, or it can drop the fallbacks of a
+ // COALESCE that follow this call.
+ return typeFactory.createTypeWithNullability(arrayType, nullable);
+ }
+
+ /// The type an operand of this type has at execution time. It must match the numeric cases of
+ /// `RelToPlanNodeConverter.convertToColumnDataType`, which converts the operands for execution, so that the
+ /// planner infers the same element type that [#getFunctionInfo(ColumnDataType[])] picks at execution time. That
+ /// converter is in `pinot-query-planner`, which this module cannot depend on.
+ private static ColumnDataType toArgumentType(RelDataType operandType) {
+ switch (operandType.getSqlTypeName()) {
+ case TINYINT:
+ case SMALLINT:
+ case INTEGER:
+ case UTINYINT:
+ case USMALLINT:
+ return ColumnDataType.INT;
+ case BIGINT:
+ case UINTEGER:
+ return ColumnDataType.LONG;
+ case DECIMAL:
+ if (operandType.getScale() == 0 && operandType.getPrecision() <= 10) {
+ return ColumnDataType.INT;
+ }
+ if (operandType.getScale() == 0 && operandType.getPrecision() <= 38) {
+ return ColumnDataType.LONG;
+ }
+ return ColumnDataType.DOUBLE;
+ case REAL:
+ case FLOAT:
+ return ColumnDataType.FLOAT;
+ default:
+ // DOUBLE. The converter rejects UBIGINT, so its element type does not matter.
+ return ColumnDataType.DOUBLE;
+ }
+ }
+
+ /// Index into [#ELEMENT_TYPES] of the narrowest sequence this argument type can be generated as, or `-1` if no
+ /// sequence can be generated from it.
+ private static int width(ColumnDataType argumentType) {
+ switch (argumentType) {
+ case INT:
+ return 0;
+ case LONG:
+ return 1;
+ case FLOAT:
+ return 2;
+ case DOUBLE:
+ case BIG_DECIMAL:
+ return 3;
+ default:
+ return -1;
+ }
+ }
+}
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/GenerateArrayFunctionsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/GenerateArrayFunctionsTest.java
new file mode 100644
index 00000000000..aac73b25684
--- /dev/null
+++ b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/GenerateArrayFunctionsTest.java
@@ -0,0 +1,156 @@
+/**
+ * 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;
+
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
+
+
+/// Tests the `generate*Array` sequence generators, in particular the bounds that keep a literal call from producing
+/// an unusable array on the broker.
+public class GenerateArrayFunctionsTest {
+
+ @Test
+ public void testAscendingSequence() {
+ assertEquals(ArrayFunctions.generateLongArray(0, 10, 2), new long[]{0, 2, 4, 6, 8, 10});
+ assertEquals(ArrayFunctions.generateLongArray(0, 4), new long[]{0, 1, 2, 3, 4});
+ }
+
+ @Test
+ public void testDescendingSequence() {
+ assertEquals(ArrayFunctions.generateLongArray(10, 0, -5), new long[]{10, 5, 0});
+ assertEquals(ArrayFunctions.generateLongArray(2, -2), new long[]{2, 1, 0, -1, -2});
+ }
+
+ @Test
+ public void testEndNotOnAStepIsNotPassed() {
+ assertEquals(ArrayFunctions.generateLongArray(0, 9, 4), new long[]{0, 4, 8});
+ assertEquals(ArrayFunctions.generateLongArray(0, -9, -4), new long[]{0, -4, -8});
+ }
+
+ @Test
+ public void testSingleElementSequence() {
+ // A zero span is a legitimate one bucket grid rather than an error.
+ assertEquals(ArrayFunctions.generateLongArray(7, 7, 1), new long[]{7});
+ assertEquals(ArrayFunctions.generateLongArray(7, 7, -1), new long[]{7});
+ assertEquals(ArrayFunctions.generateLongArray(7, 7), new long[]{7});
+ }
+
+ @Test
+ public void testTimeGrid() {
+ // The motivating case: a 30 minute grid over 3 hours, as used to gap fill a sparse time series.
+ long[] grid = ArrayFunctions.generateLongArray(1633078800000L, 1633089600000L, 1800000L);
+ assertEquals(grid.length, 7);
+ assertEquals(grid[0], 1633078800000L);
+ assertEquals(grid[6], 1633089600000L);
+ }
+
+ @Test
+ public void testTimestampsBeyondDoublePrecision() {
+ // Nanosecond timestamps exceed 2^53, so the length must not be computed in floating point.
+ long start = 1633078800000000000L;
+ assertEquals(ArrayFunctions.generateLongArray(start, start + 4, 2), new long[]{start, start + 2, start + 4});
+ }
+
+ @Test
+ public void testZeroIncrementRejected() {
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateLongArray(0, 10, 0));
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateIntArray(0, 10, 0));
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateFloatArray(0, 10, 0));
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateDoubleArray(0, 10, 0));
+ }
+
+ @Test
+ public void testIncrementPointingAwayFromEndRejected() {
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateLongArray(10, 0, 1));
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateLongArray(0, 10, -1));
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateIntArray(10, 0, 1));
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateFloatArray(10, 0, 1));
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateDoubleArray(10, 0, 1));
+ }
+
+ @Test
+ public void testOversizedSequenceRejected() {
+ assertEquals(ArrayFunctions.generateLongArray(1, 100000, 1).length, 100000);
+ // One element past the limit, so the rejection is the limit and not an incidental overflow.
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateLongArray(0, 100000, 1));
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateDoubleArray(0, 1, 0.000001));
+ }
+
+ @Test
+ public void testRangeWiderThanLongRejected() {
+ // The span overflows a long. It used to be truncated to an int, silently yielding an empty array.
+ assertThrows(IllegalArgumentException.class,
+ () -> ArrayFunctions.generateLongArray(Long.MIN_VALUE, Long.MAX_VALUE, 1));
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateLongArray(0, Long.MAX_VALUE, 1));
+ }
+
+ @Test
+ public void testNonFiniteBoundsRejected() {
+ assertThrows(IllegalArgumentException.class,
+ () -> ArrayFunctions.generateDoubleArray(0, Double.POSITIVE_INFINITY, 1));
+ assertThrows(IllegalArgumentException.class, () -> ArrayFunctions.generateDoubleArray(0, 10, Double.NaN));
+ }
+
+ @Test
+ public void testInfiniteIncrementRejected() {
+ // The first element is start + 0 * inc, which is NaN for an infinite increment.
+ assertThrows(IllegalArgumentException.class,
+ () -> ArrayFunctions.generateDoubleArray(0, 1, Double.POSITIVE_INFINITY));
+ assertThrows(IllegalArgumentException.class,
+ () -> ArrayFunctions.generateDoubleArray(1, 0, Double.NEGATIVE_INFINITY));
+ assertThrows(IllegalArgumentException.class,
+ () -> ArrayFunctions.generateFloatArray(0f, 1f, Float.POSITIVE_INFINITY));
+ assertThrows(IllegalArgumentException.class,
+ () -> ArrayFunctions.generateFloatArray(1f, 0f, Float.NEGATIVE_INFINITY));
+ }
+
+ @Test
+ public void testIntAndFloatSequences() {
+ assertEquals(ArrayFunctions.generateIntArray(0, 10, 5), new int[]{0, 5, 10});
+ assertEquals(ArrayFunctions.generateFloatArray(0f, 1f, 0.5f), new float[]{0f, 0.5f, 1f});
+ }
+
+ @Test
+ public void testFloatingPointStepDoesNotDrift() {
+ // Accumulating the step would leave the last element short of 1.0.
+ double[] values = ArrayFunctions.generateDoubleArray(0, 1, 0.1);
+ assertEquals(values.length, 11);
+ assertEquals(values[10], 1.0, 1e-12);
+ }
+
+ @Test
+ public void testFloatingPointStepKeepsInexactEnd() {
+ // (0.3 - 0.1) / 0.1 is 1.9999999999999998, so flooring the step count would drop the end. The single-stage
+ // engine has always kept it, and so does this.
+ assertEquals(ArrayFunctions.generateDoubleArray(0.1, 0.3, 0.1), new double[]{0.1, 0.1 + 0.1, 0.1 + 0.1 * 2});
+ assertEquals(ArrayFunctions.generateDoubleArray(0.3, 0.1, -0.1).length, 3);
+ assertEquals(ArrayFunctions.generateFloatArray(0.1f, 0.3f, 0.1f).length, 3);
+ assertEquals(ArrayFunctions.generateDoubleArray(1e6 + 0.1, 1e6 + 0.3, 0.1).length, 3);
+ }
+
+ @Test
+ public void testFloatingPointEndBetweenStepsIsExcluded() {
+ assertEquals(ArrayFunctions.generateDoubleArray(0, 0.95, 0.1).length, 10);
+ assertEquals(ArrayFunctions.generateDoubleArray(0, 0.2999, 0.1).length, 3);
+ assertEquals(ArrayFunctions.generateFloatArray(0f, 0.95f, 0.1f).length, 10);
+ }
+}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/GenerateArrayTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/GenerateArrayTransformFunction.java
index 393432b875e..5343a22efad 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/GenerateArrayTransformFunction.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/GenerateArrayTransformFunction.java
@@ -24,7 +24,11 @@ import java.util.Arrays;
import java.util.List;
import java.util.Map;
import javax.annotation.Nullable;
+import org.apache.pinot.common.function.scalar.ArrayFunctions;
+import org.apache.pinot.common.function.scalar.array.GenerateArrayScalarFunction;
import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.request.context.LiteralContext;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.core.operator.ColumnContext;
import org.apache.pinot.core.operator.blocks.ValueBlock;
import org.apache.pinot.core.operator.transform.TransformResultMetadata;
@@ -66,95 +70,51 @@ public class GenerateArrayTransformFunction implements TransformFunction {
Preconditions.checkState(literalContext.getType() == ExpressionContext.Type.LITERAL,
"GenerateArrayTransformFunction only takes literals as arguments, found: %s", literalContext);
}
- // Get the type of the first member in the literalContext and generate an array
- _dataType = literalContexts.get(0).getLiteral().getType();
-
- switch (_dataType) {
+ // Generate the widest of the argument types, as the multi-stage engine does. A leaf stage projection is evaluated
+ // here, and its result must have the type the multi-stage planner inferred.
+ ColumnDataType[] argumentTypes = new ColumnDataType[literalContexts.size()];
+ for (int i = 0; i < argumentTypes.length; i++) {
+ argumentTypes[i] = ColumnDataType.fromDataTypeSV(literalContexts.get(i).getLiteral().getType());
+ }
+ ColumnDataType elementType = GenerateArrayScalarFunction.getElementType(argumentTypes);
+ Preconditions.checkState(elementType != null, "Illegal argument types for GenerateArrayTransformFunction: %s",
+ literalContexts);
+ _dataType = elementType.toDataType();
+ LiteralContext start = literalContexts.get(0).getLiteral();
+ LiteralContext end = literalContexts.get(1).getLiteral();
+ LiteralContext inc = literalContexts.size() == 3 ? literalContexts.get(2).getLiteral() : null;
+ int[] intArray = null;
+ long[] longArray = null;
+ float[] floatArray = null;
+ double[] doubleArray = null;
+ switch (elementType) {
case INT:
- int startInt = literalContexts.get(0).getLiteral().getIntValue();
- int endInt = literalContexts.get(1).getLiteral().getIntValue();
- int incInt;
- if (literalContexts.size() == 3) {
- incInt = literalContexts.get(2).getLiteral().getIntValue();
- } else {
- incInt = 1;
- }
- Preconditions.checkState((endInt > startInt && incInt > 0) || (startInt > endInt
- && incInt < 0), "Incorrect Step value.");
- int size = (endInt - startInt) / incInt + 1;
- _intArrayLiteral = new int[size];
- for (int i = 0, value = startInt; i < size; i++, value += incInt) {
- _intArrayLiteral[i] = value;
- }
- _longArrayLiteral = null;
- _floatArrayLiteral = null;
- _doubleArrayLiteral = null;
+ intArray = inc != null
+ ? ArrayFunctions.generateIntArray(start.getIntValue(), end.getIntValue(), inc.getIntValue())
+ : ArrayFunctions.generateIntArray(start.getIntValue(), end.getIntValue());
break;
case LONG:
- long startLong = Long.parseLong(literalContexts.get(0).getLiteral().getStringValue());
- long endLong = Long.parseLong(literalContexts.get(1).getLiteral().getStringValue());
- long incLong;
- if (literalContexts.size() == 3) {
- incLong = Long.parseLong(literalContexts.get(2).getLiteral().getStringValue());
- } else {
- incLong = 1L;
- }
- Preconditions.checkState((endLong > startLong && incLong > 0) || (startLong > endLong
- && incLong < 0), "Incorrect Step value.");
- size = (int) ((endLong - startLong) / incLong + 1);
- _longArrayLiteral = new long[size];
- for (int i = 0; i < size; i++, startLong += incLong) {
- _longArrayLiteral[i] = startLong;
- }
- _intArrayLiteral = null;
- _floatArrayLiteral = null;
- _doubleArrayLiteral = null;
+ longArray = inc != null
+ ? ArrayFunctions.generateLongArray(start.getLongValue(), end.getLongValue(), inc.getLongValue())
+ : ArrayFunctions.generateLongArray(start.getLongValue(), end.getLongValue());
break;
case FLOAT:
- float startFloat = Float.parseFloat(literalContexts.get(0).getLiteral().getStringValue());
- float endFloat = Float.parseFloat(literalContexts.get(1).getLiteral().getStringValue());
- float incFloat;
- if (literalContexts.size() == 3) {
- incFloat = Float.parseFloat(literalContexts.get(2).getLiteral().getStringValue());
- } else {
- incFloat = 1;
- }
- Preconditions.checkState((endFloat > startFloat && incFloat > 0) || (startFloat > endFloat
- && incFloat < 0), "Incorrect Step value.");
- size = (int) ((endFloat - startFloat) / incFloat + 1);
- _floatArrayLiteral = new float[size];
- for (int i = 0; i < size; i++, startFloat += incFloat) {
- _floatArrayLiteral[i] = startFloat;
- }
- _intArrayLiteral = null;
- _longArrayLiteral = null;
- _doubleArrayLiteral = null;
+ floatArray = inc != null
+ ? ArrayFunctions.generateFloatArray(start.getFloatValue(), end.getFloatValue(), inc.getFloatValue())
+ : ArrayFunctions.generateFloatArray(start.getFloatValue(), end.getFloatValue());
break;
case DOUBLE:
- double startDouble = Double.parseDouble(literalContexts.get(0).getLiteral().getStringValue());
- double endDouble = Double.parseDouble(literalContexts.get(1).getLiteral().getStringValue());
- double incDouble;
- if (literalContexts.size() == 3) {
- incDouble = Double.parseDouble(literalContexts.get(2).getLiteral().getStringValue());
- } else {
- incDouble = 1.0;
- }
- Preconditions.checkState((endDouble > startDouble && incDouble > 0) || (startDouble > endDouble
- && incDouble < 0), "Incorrect Step value.");
- size = (int) ((endDouble - startDouble) / incDouble + 1);
- _doubleArrayLiteral = new double[size];
- for (int i = 0; i < size; i++, startDouble += incDouble) {
- _doubleArrayLiteral[i] = startDouble;
- }
- _intArrayLiteral = null;
- _longArrayLiteral = null;
- _floatArrayLiteral = null;
+ doubleArray = inc != null
+ ? ArrayFunctions.generateDoubleArray(start.getDoubleValue(), end.getDoubleValue(), inc.getDoubleValue())
+ : ArrayFunctions.generateDoubleArray(start.getDoubleValue(), end.getDoubleValue());
break;
default:
- throw new IllegalStateException(
- "Illegal data type for GenerateArrayTransformFunction: " + _dataType + ", literal contexts: "
- + Arrays.toString(literalContexts.toArray()));
+ throw new IllegalStateException("Unexpected element type for GenerateArrayTransformFunction: " + elementType);
}
+ _intArrayLiteral = intArray;
+ _longArrayLiteral = longArray;
+ _floatArrayLiteral = floatArray;
+ _doubleArrayLiteral = doubleArray;
}
public int[] getIntArrayLiteral() {
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java
index 8b8842374d5..53baf464ed7 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java
@@ -320,8 +320,10 @@ public class TransformFunctionFactory {
expression.getFunction().getArguments(), ArrayLiteralTransformFunction::new);
}
- // Check if the function is GenerateArray transform function
- if (functionName.equalsIgnoreCase(GenerateArrayTransformFunction.FUNCTION_NAME)) {
+ // Check if the function is GenerateArray transform function. It generates the array once, so it takes only
+ // literals. Other arguments go through the scalar function below, which generates an array per row.
+ if (functionName.equalsIgnoreCase(GenerateArrayTransformFunction.FUNCTION_NAME) && function.getArguments()
+ .stream().allMatch(argument -> argument.getType() == ExpressionContext.Type.LITERAL)) {
return queryContext.getOrComputeSharedValue(GenerateArrayTransformFunction.class,
expression.getFunction().getArguments(),
GenerateArrayTransformFunction::new);
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/GenerateArrayTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/GenerateArrayTransformFunctionTest.java
index 72bd48f7ff1..507941e19ae 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/GenerateArrayTransformFunctionTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/GenerateArrayTransformFunctionTest.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.core.operator.transform.function;
+import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.List;
import org.apache.pinot.common.request.context.ExpressionContext;
@@ -150,10 +151,35 @@ public class GenerateArrayTransformFunctionTest {
arrayExpressions.add(ExpressionContext.forLiteral(DataType.INT, j));
}
+ // Argument validation now lives in ArrayFunctions, shared with the multi-stage engine, and rejects a bad step
+ // with IllegalArgumentException rather than IllegalStateException.
try {
GenerateArrayTransformFunction intArray = new GenerateArrayTransformFunction(arrayExpressions);
Assert.fail();
- } catch (IllegalStateException ignored) {
+ } catch (IllegalArgumentException ignored) {
}
}
+
+ /// The element type is the widest of the argument types, as in the multi-stage engine, whose leaf stage projections
+ /// are evaluated here and must produce the type the multi-stage planner inferred.
+ @Test
+ public void testElementTypeIsTheWidestArgumentType() {
+ GenerateArrayTransformFunction longArray = new GenerateArrayTransformFunction(
+ List.of(ExpressionContext.forLiteral(DataType.INT, 0), ExpressionContext.forLiteral(DataType.LONG, 4L),
+ ExpressionContext.forLiteral(DataType.INT, 2)));
+ Assert.assertEquals(longArray.getResultMetadata().getDataType(), DataType.LONG);
+ Assert.assertEquals(longArray.getLongArrayLiteral(), new long[]{0L, 2L, 4L});
+
+ GenerateArrayTransformFunction doubleArray = new GenerateArrayTransformFunction(
+ List.of(ExpressionContext.forLiteral(DataType.INT, 0), ExpressionContext.forLiteral(DataType.INT, 1),
+ ExpressionContext.forLiteral(DataType.DOUBLE, 0.5)));
+ Assert.assertEquals(doubleArray.getResultMetadata().getDataType(), DataType.DOUBLE);
+ Assert.assertEquals(doubleArray.getDoubleArrayLiteral(), new double[]{0.0, 0.5, 1.0});
+
+ GenerateArrayTransformFunction bigDecimalArray = new GenerateArrayTransformFunction(
+ List.of(ExpressionContext.forLiteral(DataType.INT, 0), ExpressionContext.forLiteral(DataType.INT, 1),
+ ExpressionContext.forLiteral(DataType.BIG_DECIMAL, new BigDecimal("0.5"))));
+ Assert.assertEquals(bigDecimalArray.getResultMetadata().getDataType(), DataType.DOUBLE);
+ Assert.assertEquals(bigDecimalArray.getDoubleArrayLiteral(), new double[]{0.0, 0.5, 1.0});
+ }
}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java
index 50e9775d257..90ab19eea37 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ScalarTransformFunctionWrapperTest.java
@@ -73,6 +73,32 @@ public class ScalarTransformFunctionWrapperTest extends BaseTransformFunctionTes
testTransformFunction(transformFunction, expectedValues);
}
+ @Test
+ public void testGenerateArrayOnColumnsTransformFunction() {
+ // GenerateArrayTransformFunction takes only literals, so a call on columns must use the scalar function.
+ ExpressionContext expression =
+ RequestContextUtils.getExpression(String.format("generateArray(%s, %s, 1)", INT_SV_COLUMN, INT_SV_COLUMN));
+ TransformFunction transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap);
+ assertTrue(transformFunction instanceof ScalarTransformFunctionWrapper);
+ assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.INT);
+ int[][] expectedIntValues = new int[NUM_ROWS][];
+ for (int i = 0; i < NUM_ROWS; i++) {
+ expectedIntValues[i] = new int[]{_intSVValues[i]};
+ }
+ testTransformFunctionMV(transformFunction, expectedIntValues);
+
+ expression =
+ RequestContextUtils.getExpression(String.format("generateArray(%s, %s)", LONG_SV_COLUMN, LONG_SV_COLUMN));
+ transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap);
+ assertTrue(transformFunction instanceof ScalarTransformFunctionWrapper);
+ assertEquals(transformFunction.getResultMetadata().getDataType(), DataType.LONG);
+ long[][] expectedLongValues = new long[NUM_ROWS][];
+ for (int i = 0; i < NUM_ROWS; i++) {
+ expectedLongValues[i] = new long[]{_longSVValues[i]};
+ }
+ testTransformFunctionMV(transformFunction, expectedLongValues);
+ }
+
@Test
public void testStringUpperTransformFunction() {
ExpressionContext expression =
diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml
index 5fb9cfcb4b6..87028ae77b7 100644
--- a/pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml
+++ b/pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml
@@ -774,20 +774,28 @@ gcd:
scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.gcd}"
transform: null
udf: null
+generatearray:
+ scalar: "org.apache.pinot.common.function.scalar.array.GenerateArrayScalarFunction"
+ transform: null
+ udf: null
generatedoublearray:
- scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.generateDoubleArray}"
+ scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.ArrayFunctions.generateDoubleArray,\
+ \ 3: org.apache.pinot.common.function.scalar.ArrayFunctions.generateDoubleArray]}"
transform: null
udf: null
generatefloatarray:
- scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.generateFloatArray}"
+ scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.ArrayFunctions.generateFloatArray,\
+ \ 3: org.apache.pinot.common.function.scalar.ArrayFunctions.generateFloatArray]}"
transform: null
udf: null
generateintarray:
- scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.generateIntArray}"
+ scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.ArrayFunctions.generateIntArray,\
+ \ 3: org.apache.pinot.common.function.scalar.ArrayFunctions.generateIntArray]}"
transform: null
udf: null
generatelongarray:
- scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.generateLongArray}"
+ scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.ArrayFunctions.generateLongArray,\
+ \ 3: org.apache.pinot.common.function.scalar.ArrayFunctions.generateLongArray]}"
transform: null
udf: null
geotoh3:
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java
index 2363cb54794..1c721f203ae 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java
@@ -19,6 +19,7 @@
package org.apache.pinot.calcite.rel.rules;
import com.google.common.base.Preconditions;
+import java.lang.reflect.Array;
import java.math.BigDecimal;
import java.sql.Timestamp;
import java.util.ArrayList;
@@ -230,17 +231,13 @@ public class PinotEvaluateLiteralRule {
}
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
+ // Object[]. Read the elements reflectively instead, which boxes them the same way either kind of array does.
+ int length = Array.getLength(resultValue);
+ RelDataType componentType = rexNodeType.getComponentType();
+ List<Object> resultValues = new ArrayList<>(length);
+ for (int i = 0; i < length; i++) {
+ resultValues.add(convertResultValue(Array.get(resultValue, i), componentType));
}
return rexBuilder.makeLiteral(resultValues, rexNodeType, false);
}
diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
index 714d3983944..baa47e43635 100644
--- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
+++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
@@ -1776,4 +1776,67 @@ public class QueryCompilationTest extends QueryEnvironmentTestBase {
"usePlannerRules=JoinToEnrichedJoin must be a no-op, got:\n" + explain);
assertTrue(explain.contains("LogicalJoin"), "expected an ordinary LogicalJoin, got:\n" + explain);
}
+
+ /// `generateArray` on literal arguments must not fold into an array literal. The array would be serialized into the
+ /// plan of every worker, while the workers only need the three arguments to generate it themselves.
+ @Test
+ public void testGenerateArrayIsNotFolded() {
+ String explain = _queryEnvironment.explainQuery(
+ "EXPLAIN PLAN FOR SELECT grid.ts FROM UNNEST(generateArray(0, 6, 2)) AS grid(ts)",
+ RANDOM_REQUEST_ID_GEN.nextLong());
+ assertTrue(explain.contains("GENERATEARRAY(0, 6, 2)"), "expected an unfolded generateArray call, got:\n" + explain);
+ assertFalse(explain.contains("ARRAY(0, 2, 4, 6)"), "expected no folded array literal, got:\n" + explain);
+ }
+
+ /// Folding a scalar function returning a primitive array used to fail in
+ /// [org.apache.pinot.calcite.rel.rules.PinotEvaluateLiteralRule], which read the elements by casting to `Object[]`
+ /// -- a cast a `long[]` cannot satisfy.
+ @Test
+ public void testPrimitiveArrayFunctionFoldsIntoAnArrayLiteral() {
+ String explain = _queryEnvironment.explainQuery(
+ "EXPLAIN PLAN FOR SELECT grid.ts FROM UNNEST(generateLongArray(0, 6, 2)) AS grid(ts)",
+ RANDOM_REQUEST_ID_GEN.nextLong());
+ assertTrue(explain.contains("ARRAY(0:BIGINT, 2:BIGINT, 4:BIGINT, 6:BIGINT)"),
+ "expected a folded array literal, got:\n" + explain);
+ }
+
+ /// The element type is the widest of the argument types, so a time grid in epoch millis stays a LONG sequence
+ /// instead of overflowing an INT, and a fractional increment makes the whole sequence DOUBLE.
+ @Test
+ public void testGenerateArrayElementTypeFollowsItsArguments() {
+ assertEquals(generateArrayElementType("0, 6, 2"), SqlTypeName.INTEGER);
+ assertEquals(generateArrayElementType("0, 6"), SqlTypeName.INTEGER);
+ assertEquals(generateArrayElementType("1633078800000, 1633080600000, 1800000"), SqlTypeName.BIGINT);
+ assertEquals(generateArrayElementType("0, 1633080600000, 1800000"), SqlTypeName.BIGINT);
+ assertEquals(generateArrayElementType("CAST(0 AS FLOAT), CAST(1 AS FLOAT), CAST(0.5 AS FLOAT)"),
+ SqlTypeName.REAL);
+ assertEquals(generateArrayElementType("0, 1, 0.5"), SqlTypeName.DOUBLE);
+ // The plan converts these decimals to INT, LONG and DOUBLE, so execution picks those element types.
+ assertEquals(generateArrayElementType("CAST(0 AS DECIMAL(10, 0)), 4"), SqlTypeName.INTEGER);
+ assertEquals(generateArrayElementType("CAST(0 AS DECIMAL(12, 0)), 4"), SqlTypeName.BIGINT);
+ assertEquals(generateArrayElementType("CAST(0 AS DECIMAL(5, 2)), 4"), SqlTypeName.DOUBLE);
+ }
+
+ /// The function returns null when an argument is null. If the inferred type did not allow null, Calcite would drop
+ /// the fallback of a `COALESCE` that follows the call.
+ @Test
+ public void testGenerateArrayTypeAllowsNullWhenAnArgumentDoes() {
+ assertTrue(generateArrayType("CAST(NULL AS INTEGER), 4").isNullable());
+ assertFalse(generateArrayType("0, 4").isNullable());
+ String explain = _queryEnvironment.explainQuery("EXPLAIN PLAN FOR SELECT g.x FROM "
+ + "UNNEST(COALESCE(generateArray(CAST(NULL AS INTEGER), 4), ARRAY[7])) AS g(x)",
+ RANDOM_REQUEST_ID_GEN.nextLong());
+ assertTrue(explain.contains("ARRAY(7)"), "expected the COALESCE fallback, got:\n" + explain);
+ }
+
+ private SqlTypeName generateArrayElementType(String arguments) {
+ String query = "SELECT grid.ts FROM UNNEST(generateArray(" + arguments + ")) AS grid(ts)";
+ return _queryEnvironment.compile(query).getRelRoot().validatedRowType.getFieldList().get(0).getType()
+ .getSqlTypeName();
+ }
+
+ private RelDataType generateArrayType(String arguments) {
+ String query = "SELECT generateArray(" + arguments + ") AS grid FROM a";
+ return _queryEnvironment.compile(query).getRelRoot().validatedRowType.getFieldList().get(0).getType();
+ }
}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/TypeUtils.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/TypeUtils.java
index f7a5ebe9c03..2d4972e3dd5 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/TypeUtils.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/TypeUtils.java
@@ -23,6 +23,7 @@ import it.unimi.dsi.fastutil.floats.FloatArrayList;
import it.unimi.dsi.fastutil.ints.IntArrayList;
import it.unimi.dsi.fastutil.longs.LongArrayList;
import it.unimi.dsi.fastutil.objects.ObjectArrayList;
+import java.lang.reflect.Array;
import java.math.BigDecimal;
import org.apache.pinot.common.utils.ArrayListUtils;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
@@ -59,6 +60,9 @@ public class TypeUtils {
// For ArrayAggregationFunction
return ArrayListUtils.toIntArray((IntArrayList) value);
}
+ if (isOtherNumericArray(value, int[].class)) {
+ return convertNumericArray(value, storedType);
+ }
assert value instanceof int[];
return value;
case LONG_ARRAY:
@@ -66,6 +70,9 @@ public class TypeUtils {
// For FunnelCountAggregationFunction and ArrayAggregationFunction
return ArrayListUtils.toLongArray((LongArrayList) value);
}
+ if (isOtherNumericArray(value, long[].class)) {
+ return convertNumericArray(value, storedType);
+ }
assert value instanceof long[];
return value;
case FLOAT_ARRAY:
@@ -83,6 +90,9 @@ public class TypeUtils {
}
return floatArray;
}
+ if (isOtherNumericArray(value, float[].class)) {
+ return convertNumericArray(value, storedType);
+ }
assert value instanceof float[];
return value;
case DOUBLE_ARRAY:
@@ -111,6 +121,9 @@ public class TypeUtils {
}
return doubleArray;
}
+ if (isOtherNumericArray(value, double[].class)) {
+ return convertNumericArray(value, storedType);
+ }
assert value instanceof double[];
return value;
case BIG_DECIMAL_ARRAY:
@@ -140,6 +153,47 @@ public class TypeUtils {
}
}
+ /// Returns whether the value is a primitive numeric array with an element type other than the expected one.
+ private static boolean isOtherNumericArray(Object value, Class<?> expectedClass) {
+ return value.getClass() != expectedClass && (value instanceof int[] || value instanceof long[]
+ || value instanceof float[] || value instanceof double[]);
+ }
+
+ /// Converts the elements of a primitive numeric array the way [#convert] converts a single value. The single-stage
+ /// engine can compute another element type than the one the multi-stage planner inferred. For example, its `plus`
+ /// returns DOUBLE, so `generateArray(intCol, intCol + 1)` generates a `double[]` where the plan has an INT_ARRAY.
+ private static Object convertNumericArray(Object value, ColumnDataType storedType) {
+ int length = Array.getLength(value);
+ switch (storedType) {
+ case INT_ARRAY:
+ int[] intArray = new int[length];
+ for (int i = 0; i < length; i++) {
+ intArray[i] = ((Number) Array.get(value, i)).intValue();
+ }
+ return intArray;
+ case LONG_ARRAY:
+ long[] longArray = new long[length];
+ for (int i = 0; i < length; i++) {
+ longArray[i] = ((Number) Array.get(value, i)).longValue();
+ }
+ return longArray;
+ case FLOAT_ARRAY:
+ float[] floatArray = new float[length];
+ for (int i = 0; i < length; i++) {
+ floatArray[i] = ((Number) Array.get(value, i)).floatValue();
+ }
+ return floatArray;
+ case DOUBLE_ARRAY:
+ double[] doubleArray = new double[length];
+ for (int i = 0; i < length; i++) {
+ doubleArray[i] = ((Number) Array.get(value, i)).doubleValue();
+ }
+ return doubleArray;
+ default:
+ throw new IllegalStateException("Not a numeric array type: " + storedType);
+ }
+ }
+
/// Converts row to the desired stored [ColumnDataType]s in-place. This is used to convert rows generated from
/// single-stage engine to be used in multi-stage engine.
public static void convertRow(Object[] row, ColumnDataType[] outputStoredTypes) {
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/utils/TypeUtilsTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/utils/TypeUtilsTest.java
index 653f08c699a..2b14f6c8e01 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/utils/TypeUtilsTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/utils/TypeUtilsTest.java
@@ -84,4 +84,25 @@ public class TypeUtilsTest {
assertThrows(NullPointerException.class, () -> TypeUtils.convert(null, type));
}
}
+
+ @Test
+ public void testNumericArrayKeepsItsType() {
+ int[] intArray = {1, 2};
+ assertSame(TypeUtils.convert(intArray, ColumnDataType.INT_ARRAY), intArray);
+ double[] doubleArray = {1.5, 2.5};
+ assertSame(TypeUtils.convert(doubleArray, ColumnDataType.DOUBLE_ARRAY), doubleArray);
+ }
+
+ /// The single-stage engine can compute another element type than the one the multi-stage planner inferred, for
+ /// example a `double[]` from `generateArray(intCol, intCol + 1)`, whose `plus` returns DOUBLE.
+ @Test
+ public void testNumericArrayConvertsToTheStoredType() {
+ assertEquals(TypeUtils.convert(new double[]{1, 2}, ColumnDataType.INT_ARRAY), new int[]{1, 2});
+ assertEquals(TypeUtils.convert(new double[]{1, 2}, ColumnDataType.LONG_ARRAY), new long[]{1, 2});
+ assertEquals(TypeUtils.convert(new int[]{1, 2}, ColumnDataType.LONG_ARRAY), new long[]{1, 2});
+ assertEquals(TypeUtils.convert(new long[]{1, 2}, ColumnDataType.INT_ARRAY), new int[]{1, 2});
+ assertEquals(TypeUtils.convert(new int[]{1, 2}, ColumnDataType.FLOAT_ARRAY), new float[]{1, 2});
+ assertEquals(TypeUtils.convert(new float[]{1.5f}, ColumnDataType.DOUBLE_ARRAY), new double[]{1.5});
+ assertEquals(TypeUtils.convert(new long[]{1L << 40}, ColumnDataType.DOUBLE_ARRAY), new double[]{1L << 40});
+ }
}
diff --git a/pinot-query-runtime/src/test/resources/queries/GenerateArray.json b/pinot-query-runtime/src/test/resources/queries/GenerateArray.json
new file mode 100644
index 00000000000..2e44c13fbc2
--- /dev/null
+++ b/pinot-query-runtime/src/test/resources/queries/GenerateArray.json
@@ -0,0 +1,282 @@
+{
+ "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
+ },
+ {
+ "description": "column arguments generate one array per row in the leaf stage",
+ "sql": "SELECT val, generateArray(val, val + 2) FROM {sparse} WHERE series = 'a' ORDER BY val",
+ "keepOutputRowOrder": true,
+ "outputs": [
+ [
+ 10,
+ [
+ 10,
+ 11,
+ 12
+ ]
+ ],
+ [
+ 40,
+ [
+ 40,
+ 41,
+ 42
+ ]
+ ]
+ ],
+ "ignoreLiteMode": true,
+ "ignoreV2Optimizer": true
+ },
+ {
+ "description": "a scale zero DECIMAL argument generates the INT sequence the plan converts it to",
+ "sql": "SELECT generateArray(CAST(0 AS DECIMAL(10, 0)), 2) FROM {sparse} WHERE series = 'b'",
+ "keepOutputRowOrder": true,
+ "outputs": [
+ [
+ [
+ 0,
+ 1,
+ 2
+ ]
+ ]
+ ],
+ "ignoreLiteMode": true,
+ "ignoreV2Optimizer": true
+ },
+ {
+ "description": "a null argument yields null, so COALESCE keeps its fallback",
+ "sql": "SELECT g.x FROM UNNEST(COALESCE(generateArray(CAST(NULL AS INTEGER), 4), ARRAY[7])) AS g(x)",
+ "outputs": [
+ [
+ 7
+ ]
+ ],
+ "ignoreLiteMode": true,
+ "ignoreV2Optimizer": true
+ },
+ {
+ "description": "a zero increment is rejected instead of dividing by zero",
+ "sql": "SELECT grid.bucket FROM {sparse} CROSS JOIN UNNEST(generateArray(0, 4, 0)) AS grid(bucket)",
+ "expectedException": ".*Increment must not be zero.*",
+ "outputs": [],
+ "ignoreLiteMode": true,
+ "ignoreV2Optimizer": true
+ },
+ {
+ "description": "a sequence too long to materialize is rejected",
+ "sql": "SELECT grid.bucket FROM {sparse} CROSS JOIN UNNEST(generateArray(0, 1000000)) AS grid(bucket)",
+ "expectedException": ".*exceeds the maximum of 100000.*",
+ "outputs": [],
+ "ignoreLiteMode": true,
+ "ignoreV2Optimizer": true
+ }
+ ]
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org