This is an automated email from the ASF dual-hosted git repository.
xiong pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/calcite.git
The following commit(s) were added to refs/heads/main by this push:
new 955fd9c1cf [CALCITE-5830] Add ARRAY_INSERT function(enabled in Spark
library)
955fd9c1cf is described below
commit 955fd9c1cf7601d88305c3206890480e2dc076e8
Author: Ran Tao <[email protected]>
AuthorDate: Tue Jul 11 10:24:41 2023 +0800
[CALCITE-5830] Add ARRAY_INSERT function(enabled in Spark library)
---
.../calcite/adapter/enumerable/RexImpTable.java | 2 +
.../org/apache/calcite/runtime/SqlFunctions.java | 94 +++++++++++++++++++
.../main/java/org/apache/calcite/sql/SqlKind.java | 3 +
.../calcite/sql/fun/SqlLibraryOperators.java | 23 +++++
.../sql/type/ArrayInsertOperandTypeChecker.java | 100 +++++++++++++++++++++
.../org/apache/calcite/sql/type/OperandTypes.java | 3 +
.../org/apache/calcite/util/BuiltInMethod.java | 1 +
site/_docs/reference.md | 1 +
.../org/apache/calcite/test/SqlOperatorTest.java | 75 ++++++++++++++++
9 files changed, 302 insertions(+)
diff --git
a/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImpTable.java
b/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImpTable.java
index 7614eb784a..c23c2f8678 100644
--- a/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImpTable.java
+++ b/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImpTable.java
@@ -127,6 +127,7 @@ import static
org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY_CONCAT_AGG;
import static org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY_CONTAINS;
import static org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY_DISTINCT;
import static org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY_EXCEPT;
+import static org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY_INSERT;
import static org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY_INTERSECT;
import static org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY_JOIN;
import static org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY_LENGTH;
@@ -757,6 +758,7 @@ public class RexImpTable {
defineMethod(ARRAY_DISTINCT, BuiltInMethod.ARRAY_DISTINCT.method,
NullPolicy.STRICT);
defineMethod(ARRAY_EXCEPT, BuiltInMethod.ARRAY_EXCEPT.method,
NullPolicy.ANY);
defineMethod(ARRAY_JOIN, "arrayToString", NullPolicy.STRICT);
+ defineMethod(ARRAY_INSERT, BuiltInMethod.ARRAY_INSERT.method,
NullPolicy.NONE);
defineMethod(ARRAY_INTERSECT, BuiltInMethod.ARRAY_INTERSECT.method,
NullPolicy.ANY);
defineMethod(ARRAY_LENGTH, BuiltInMethod.COLLECTION_SIZE.method,
NullPolicy.STRICT);
defineMethod(ARRAY_MAX, BuiltInMethod.ARRAY_MAX.method,
NullPolicy.STRICT);
diff --git a/core/src/main/java/org/apache/calcite/runtime/SqlFunctions.java
b/core/src/main/java/org/apache/calcite/runtime/SqlFunctions.java
index 15fbc472fa..9374b31054 100644
--- a/core/src/main/java/org/apache/calcite/runtime/SqlFunctions.java
+++ b/core/src/main/java/org/apache/calcite/runtime/SqlFunctions.java
@@ -155,6 +155,12 @@ public class SqlFunctions {
private static final Base32 BASE_32 = new Base32();
+ // Some JVMs can't allocate arrays of length Integer.MAX_VALUE; actual max
is somewhat smaller.
+ // Be conservative and lower this value a little.
+ // @see
http://hg.openjdk.java.net/jdk8/jdk8/jdk/file/tip/src/share/classes/java/util/ArrayList.java#l229
+ // Note: this variable handling is inspired by Apache Spark
+ private static final int MAX_ARRAY_LENGTH = Integer.MAX_VALUE - 15;
+
private static final Function1<List<Object>, Enumerable<Object>>
LIST_AS_ENUMERABLE =
a0 -> a0 == null ? Linq4j.emptyEnumerable() : Linq4j.asEnumerable(a0);
@@ -4336,6 +4342,94 @@ public class SqlFunctions {
return new ArrayList<>(result);
}
+ /** Support the ARRAY_INSERT function. */
+ public static @Nullable List arrayInsert(List baselist, Object pos, Object
val) {
+ if (baselist == null || pos == null) {
+ return null;
+ }
+ int posInt = (int) pos;
+ Object[] baseArray = baselist.toArray();
+ if (posInt == 0 || posInt >= MAX_ARRAY_LENGTH || posInt <=
-MAX_ARRAY_LENGTH) {
+ throw new IllegalArgumentException("The index 0 is invalid. "
+ + "An index shall be either < 0 or > 0 (the first element has index
1) "
+ + "and not exceeds the allowed limit.");
+ }
+
+ boolean usePositivePos = posInt > 0;
+
+ if (usePositivePos) {
+ int newArrayLength = Math.max(baseArray.length + 1, posInt);
+
+ if (newArrayLength > MAX_ARRAY_LENGTH) {
+ throw new IndexOutOfBoundsException(
+ String.format(Locale.ROOT, "The new array length %s exceeds the
allowed limit.",
+ newArrayLength));
+ }
+
+ Object[] newArray = new Object[newArrayLength];
+
+ int posIndex = posInt - 1;
+ if (posIndex < baseArray.length) {
+ System.arraycopy(baseArray, 0, newArray, 0, posIndex);
+ newArray[posIndex] = val;
+ System.arraycopy(baseArray, posIndex, newArray, posIndex + 1,
baseArray.length - posIndex);
+ } else {
+ System.arraycopy(baseArray, 0, newArray, 0, baseArray.length);
+ newArray[posIndex] = val;
+ }
+
+ return Arrays.asList(newArray);
+ } else {
+ int posIndex = posInt;
+
+ boolean newPosExtendsArrayLeft = baseArray.length + posIndex < 0;
+
+ if (newPosExtendsArrayLeft) {
+ // special case, if the new position is negative but larger than the
current array size
+ // place the new item at start of array, place the current array
contents at the end
+ // and fill the newly created array elements in middle with a null
+ int newArrayLength = -posIndex + 1;
+
+ if (newArrayLength > MAX_ARRAY_LENGTH) {
+ throw new IndexOutOfBoundsException(
+ String.format(Locale.ROOT, "The new array length %s exceeds the
allowed limit.",
+ newArrayLength));
+ }
+
+ Object[] newArray = new Object[newArrayLength];
+ System.arraycopy(baseArray, 0, newArray, Math.abs(posIndex +
baseArray.length) + 1,
+ baseArray.length);
+ newArray[0] = val;
+
+ return Arrays.asList(newArray);
+ } else {
+ posIndex = posIndex + baseArray.length;
+
+ int newArrayLength = Math.max(baseArray.length + 1, posIndex + 1);
+
+ if (newArrayLength > MAX_ARRAY_LENGTH) {
+ throw new IndexOutOfBoundsException(
+ String.format(Locale.ROOT, "The new array length %s exceeds the
allowed limit.",
+ newArrayLength));
+ }
+
+ Object[] newArray = new Object[newArrayLength];
+
+ if (posIndex < baseArray.length) {
+ System.arraycopy(baseArray, 0, newArray, 0, posIndex);
+ newArray[posIndex] = val;
+ System.arraycopy(baseArray, posIndex, newArray, posIndex + 1,
+ baseArray.length - posIndex);
+ } else {
+ System.arraycopy(baseArray, 0, newArray, 0, baseArray.length);
+ newArray[posIndex] = val;
+ }
+
+ return Arrays.asList(newArray);
+ }
+ }
+ }
+
/** Support the ARRAY_INTERSECT function. */
public static List arrayIntersect(List list1, List list2) {
final Set result = new LinkedHashSet<>(list1);
diff --git a/core/src/main/java/org/apache/calcite/sql/SqlKind.java
b/core/src/main/java/org/apache/calcite/sql/SqlKind.java
index 0eb3a165f8..e6c368fa0d 100644
--- a/core/src/main/java/org/apache/calcite/sql/SqlKind.java
+++ b/core/src/main/java/org/apache/calcite/sql/SqlKind.java
@@ -704,6 +704,9 @@ public enum SqlKind {
/** {@code ARRAY_EXCEPT} function (Spark semantics). */
ARRAY_EXCEPT,
+ /** {@code ARRAY_INSERT} function (Spark semantics). */
+ ARRAY_INSERT,
+
/** {@code ARRAY_INTERSECT} function (Spark semantics). */
ARRAY_INTERSECT,
diff --git
a/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java
b/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java
index 6159d1162c..6976d52cf8 100644
--- a/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java
+++ b/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java
@@ -1055,6 +1055,29 @@ public abstract class SqlLibraryOperators {
OperandTypes.SAME_SAME,
OperandTypes.family(SqlTypeFamily.ARRAY, SqlTypeFamily.ARRAY)));
+ @SuppressWarnings("argument.type.incompatible")
+ private static RelDataType arrayInsertReturnType(SqlOperatorBinding
opBinding) {
+ final RelDataType arrayType = opBinding.collectOperandTypes().get(0);
+ final RelDataType componentType = arrayType.getComponentType();
+ final RelDataType elementType = opBinding.collectOperandTypes().get(2);
+ // we don't need to do leastRestrictive on componentType and elementType,
+ // because in operand checker we limit the elementType must equals array
component type.
+ // So we use componentType directly.
+ RelDataType type = componentType;
+ if (elementType.isNullable()) {
+ type = opBinding.getTypeFactory().createTypeWithNullability(type, true);
+ }
+ requireNonNull(type, "inferred array element type");
+ return SqlTypeUtil.createArrayType(opBinding.getTypeFactory(), type,
arrayType.isNullable());
+ }
+
+ /** The "ARRAY_INSERT(array, pos, val)" function (Spark). */
+ @LibraryOperator(libraries = {SPARK})
+ public static final SqlFunction ARRAY_INSERT =
+ SqlBasicFunction.create(SqlKind.ARRAY_INSERT,
+ SqlLibraryOperators::arrayInsertReturnType,
+ OperandTypes.ARRAY_INSERT);
+
/** The "ARRAY_INTERSECT(array1, array2)" function. */
@LibraryOperator(libraries = {SPARK})
public static final SqlFunction ARRAY_INTERSECT =
diff --git
a/core/src/main/java/org/apache/calcite/sql/type/ArrayInsertOperandTypeChecker.java
b/core/src/main/java/org/apache/calcite/sql/type/ArrayInsertOperandTypeChecker.java
new file mode 100644
index 0000000000..0310ad9eb1
--- /dev/null
+++
b/core/src/main/java/org/apache/calcite/sql/type/ArrayInsertOperandTypeChecker.java
@@ -0,0 +1,100 @@
+/*
+ * 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.calcite.sql.type;
+
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.sql.SqlCallBinding;
+import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.sql.SqlUtil;
+
+import com.google.common.collect.ImmutableList;
+
+import static
org.apache.calcite.sql.type.NonNullableAccessors.getComponentTypeOrThrow;
+import static org.apache.calcite.util.Static.RESOURCE;
+
+/**
+ * Parameter type checking strategy, where the operand size is 3, and the
first type is array,
+ * the second type is integer, and the third type is the component type of the
first array (
+ * strictly matched without leastRestrictive conversion).
+ */
+public class ArrayInsertOperandTypeChecker extends SameOperandTypeChecker {
+
+ public ArrayInsertOperandTypeChecker() {
+ super(3);
+ }
+
+ //~ Methods ----------------------------------------------------------------
+
+ @Override public boolean checkOperandTypes(
+ SqlCallBinding callBinding,
+ boolean throwOnFailure) {
+ // all operands can't be null (without cast)
+ for (SqlNode node : callBinding.operands()) {
+ if (SqlUtil.isNullLiteral(node, false)) {
+ if (throwOnFailure) {
+ throw callBinding.getValidator().newValidationError(node,
+
RESOURCE.argumentMustNotBeNull(callBinding.getOperator().getName()));
+ } else {
+ return false;
+ }
+ }
+ }
+
+ // op0 is an array type
+ final SqlNode op0 = callBinding.operand(0);
+ if (!OperandTypes.ARRAY.checkSingleOperandType(
+ callBinding,
+ op0,
+ 0,
+ throwOnFailure)) {
+ return false;
+ }
+
+ // op1 is an Integer type
+ final SqlNode op1 = callBinding.operand(1);
+ RelDataType op1Type = SqlTypeUtil.deriveType(callBinding, op1);
+ SqlTypeName op1TypeName = op1Type.getSqlTypeName();
+ if (op1TypeName != SqlTypeName.INTEGER) {
+ if (throwOnFailure) {
+ throw callBinding.newError(
+ RESOURCE.typeNotComparable(
+ op1TypeName.getName(), SqlTypeName.INTEGER.getName()));
+ }
+ return false;
+ }
+
+ // op2 must has same type with op0 component type strictly
+ final RelDataType op0ComponentType =
+ getComponentTypeOrThrow(SqlTypeUtil.deriveType(callBinding, op0));
+ SqlNode op2 = callBinding.operand(2);
+ RelDataType op2Type = SqlTypeUtil.deriveType(callBinding, op2);
+ RelDataType biggest =
+ callBinding.getTypeFactory().leastRestrictive(
+ ImmutableList.of(op0ComponentType, op2Type));
+ if (biggest == null) {
+ if (throwOnFailure) {
+ throw callBinding.newError(
+ RESOURCE.typeNotComparable(
+ op0ComponentType.getSqlTypeName().getName(),
+ op2Type.getSqlTypeName().getName()));
+ }
+ return false;
+ }
+
+ return true;
+ }
+}
diff --git a/core/src/main/java/org/apache/calcite/sql/type/OperandTypes.java
b/core/src/main/java/org/apache/calcite/sql/type/OperandTypes.java
index 7416a1549d..eac7e7450f 100644
--- a/core/src/main/java/org/apache/calcite/sql/type/OperandTypes.java
+++ b/core/src/main/java/org/apache/calcite/sql/type/OperandTypes.java
@@ -527,6 +527,9 @@ public abstract class OperandTypes {
public static final SqlOperandTypeChecker ARRAY_ELEMENT =
new ArrayElementOperandTypeChecker();
+ public static final SqlOperandTypeChecker ARRAY_INSERT =
+ new ArrayInsertOperandTypeChecker();
+
public static final SqlSingleOperandTypeChecker MAP_FROM_ENTRIES =
new MapFromEntriesOperandTypeChecker();
diff --git a/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java
b/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java
index e7bfbc0a1e..19aadc93d7 100644
--- a/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java
+++ b/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java
@@ -656,6 +656,7 @@ public enum BuiltInMethod {
ARRAY_REMOVE(SqlFunctions.class, "arrayRemove", List.class, Object.class),
ARRAY_REPEAT(SqlFunctions.class, "repeat", Object.class, Integer.class),
ARRAY_EXCEPT(SqlFunctions.class, "arrayExcept", List.class, List.class),
+ ARRAY_INSERT(SqlFunctions.class, "arrayInsert", List.class, Integer.class,
Object.class),
ARRAY_INTERSECT(SqlFunctions.class, "arrayIntersect", List.class,
List.class),
ARRAY_UNION(SqlFunctions.class, "arrayUnion", List.class, List.class),
ARRAY_REVERSE(SqlFunctions.class, "reverse", List.class),
diff --git a/site/_docs/reference.md b/site/_docs/reference.md
index ef80f5ea03..7daa15b7b2 100644
--- a/site/_docs/reference.md
+++ b/site/_docs/reference.md
@@ -2658,6 +2658,7 @@ BigQuery's type system uses confusingly different names
for types and functions:
| s | ARRAY_CONTAINS(array, element) | Returns true if the
*array* contains the *element*
| s | ARRAY_DISTINCT(array) | Removes duplicate
values from the *array* that keeps ordering of elements
| s | ARRAY_EXCEPT(array1, array2) | Returns an array of the
elements in *array1* but not in *array2*, without duplicates
+| s | ARRAY_INSERT(array, pos, element) | Places *element* into
index *pos* of *array*. Array index start at 1, or start from the end if index
is negative. Index above array size appends the array, or prepends the array if
index is negative, with `NULL` elements.
| s | ARRAY_INTERSECT(array1, array2) | Returns an array of the
elements in the intersection of *array1* and *array2*, without duplicates
| s | ARRAY_JOIN(array, delimiter [, nullText ]) | Synonym for
`ARRAY_TO_STRING`
| b | ARRAY_LENGTH(array) | Synonym for
`CARDINALITY`
diff --git a/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java
b/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java
index 82a2aa56be..1edd551984 100644
--- a/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java
+++ b/testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java
@@ -6090,6 +6090,81 @@ public class SqlOperatorTest {
f.checkNull("array_except(cast(null as integer array), cast(null as
integer array))");
}
+ /** Tests {@code ARRAY_INSERT} function from Spark. */
+ @Test void testArrayInsertFunc() {
+ final SqlOperatorFixture f0 = fixture();
+ f0.setFor(SqlLibraryOperators.ARRAY_INSERT);
+ f0.checkFails("^array_insert(null, 3, 4)^",
+ "No match found for function signature "
+ + "ARRAY_INSERT\\(<NULL>, <NUMERIC>, <NUMERIC>\\)", false);
+ f0.checkFails("^array_insert(array[1], null, 4)^",
+ "No match found for function signature "
+ + "ARRAY_INSERT\\(<INTEGER ARRAY>, <NULL>, <NUMERIC>\\)", false);
+ f0.checkFails("^array_insert(array[1], 3, null)^",
+ "No match found for function signature "
+ + "ARRAY_INSERT\\(<INTEGER ARRAY>, <NUMERIC>, <NULL>\\)", false);
+
+ final SqlOperatorFixture f1 = f0.withLibrary(SqlLibrary.SPARK);
+
+ // can't be NULL
+ f1.checkFails("array_insert(^null^, 3, 4)",
+ "Argument to function 'ARRAY_INSERT' must not be NULL", false);
+ f1.checkFails("array_insert(array[1], ^null^, 4)",
+ "Argument to function 'ARRAY_INSERT' must not be NULL", false);
+ f1.checkFails("array_insert(array[1], 3, ^null^)",
+ "Argument to function 'ARRAY_INSERT' must not be NULL", false);
+
+ // return null
+ f1.checkNull("array_insert(cast(null as integer array), 3, 4)");
+ f1.checkNull("array_insert(array[1], cast(null as integer), 4)");
+
+ // op1 must be Integer type
+ f1.checkFails("^array_insert(array[1, 2, 3], cast(3 as tinyint), 4)^",
+ "TINYINT is not comparable to INTEGER", false);
+ f1.checkFails("^array_insert(array[1, 2, 3], cast(3 as smallint), 4)^",
+ "SMALLINT is not comparable to INTEGER", false);
+ f1.checkFails("^array_insert(array[1, 2, 3], cast(3 as bigint), 4)^",
+ "BIGINT is not comparable to INTEGER", false);
+ f1.checkFails("^array_insert(array[1, 2, 3], 3.0, 4)^",
+ "DECIMAL is not comparable to INTEGER", false);
+ f1.checkFails("^array_insert(array[1, 2, 3], '3', 4)^",
+ "CHAR is not comparable to INTEGER", false);
+ // op1 can't be 0
+ f1.checkFails("array_insert(array[2, 3, 4], 0, 1)",
+ "The index 0 is invalid. "
+ + "An index shall be either < 0 or > 0 \\(the first element has
index 1\\) "
+ + "and not exceeds the allowed limit.", true);
+ // op1 overflow
+ f1.checkFails("array_insert(array[2, 3, 4], 2147483647, 1)",
+ "The index 0 is invalid. "
+ + "An index shall be either < 0 or > 0 \\(the first element has
index 1\\) "
+ + "and not exceeds the allowed limit.", true);
+ f1.checkFails("array_insert(array[2, 3, 4], -2147483648, 1)",
+ "The index 0 is invalid. "
+ + "An index shall be either < 0 or > 0 \\(the first element has
index 1\\) "
+ + "and not exceeds the allowed limit.", true);
+
+ f1.checkScalar("array_insert(array[1, 2, 3], 3, 4)",
+ "[1, 2, 4, 3]", "INTEGER NOT NULL ARRAY NOT NULL");
+ f1.checkScalar("array_insert(array[1, 2, 3], 3, cast(null as integer))",
+ "[1, 2, null, 3]", "INTEGER ARRAY NOT NULL");
+ f1.checkScalar("array_insert(array[2, 3, 4], 1, 1)",
+ "[1, 2, 3, 4]", "INTEGER NOT NULL ARRAY NOT NULL");
+ f1.checkScalar("array_insert(array[1, 3, 4], -2, 2)",
+ "[1, 2, 3, 4]", "INTEGER NOT NULL ARRAY NOT NULL");
+ f1.checkScalar("array_insert(array[2, 3, null, 4], -5, 1)",
+ "[1, null, 2, 3, null, 4]", "INTEGER ARRAY NOT NULL");
+ // check complex type
+ f1.checkScalar("array_insert(array[array[1,2]], 1, array[1])",
+ "[[1], [1, 2]]", "INTEGER NOT NULL ARRAY NOT NULL ARRAY NOT NULL");
+ f1.checkScalar("array_insert(array[array[1,2]], -1, array[1])",
+ "[[1], [1, 2]]", "INTEGER NOT NULL ARRAY NOT NULL ARRAY NOT NULL");
+ f1.checkScalar("array_insert(array[map[1, 'a']], 1, map[2, 'b'])",
"[{2=b}, {1=a}]",
+ "(INTEGER NOT NULL, CHAR(1) NOT NULL) MAP NOT NULL ARRAY NOT NULL");
+ f1.checkScalar("array_insert(array[map[1, 'a']], -1, map[2, 'b'])",
"[{2=b}, {1=a}]",
+ "(INTEGER NOT NULL, CHAR(1) NOT NULL) MAP NOT NULL ARRAY NOT NULL");
+ }
+
/** Tests {@code ARRAY_INTERSECT} function from Spark. */
@Test void testArrayIntersectFunc() {
final SqlOperatorFixture f0 = fixture();
