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 <[email protected]>
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: [email protected]
For additional commands, e-mail: [email protected]
