diff --git a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java index a69cbcbf0c6128..2e84bb97bb28ef 100644 --- a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java +++ b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java @@ -401,6 +401,16 @@ ANY, and(logical(LogicalTypeRoot.BOOLEAN), LITERAL) "org.apache.flink.table.runtime.functions.scalar.ArrayConcatFunction") .build(); + public static final BuiltInFunctionDefinition ARRAY_FLATTEN = + BuiltInFunctionDefinition.newBuilder() + .name("ARRAY_FLATTEN") + .kind(SCALAR) + .inputTypeStrategy(sequence(logical(LogicalTypeRoot.ARRAY))) + .outputTypeStrategy(nullableIfArgs(SpecificTypeStrategies.ARRAY_FLATTEN)) + .runtimeClass( + "org.apache.flink.table.runtime.functions.scalar.ArrayFlattenFunction") + .build(); + public static final BuiltInFunctionDefinition ARRAY_MAX = BuiltInFunctionDefinition.newBuilder() .name("ARRAY_MAX") diff --git a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/SpecificTypeStrategies.java b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/SpecificTypeStrategies.java index 7771211444257f..1610e7d427dcf3 100644 --- a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/SpecificTypeStrategies.java +++ b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/SpecificTypeStrategies.java @@ -27,6 +27,7 @@ import org.apache.flink.table.types.inference.TypeStrategies; import org.apache.flink.table.types.inference.TypeStrategy; import org.apache.flink.table.types.logical.LogicalTypeRoot; +import org.apache.flink.table.types.utils.TypeConversions; import java.util.List; import java.util.Optional; @@ -67,6 +68,37 @@ public final class SpecificTypeStrategies { /** Type strategy specific for array element. */ public static final TypeStrategy ARRAY_ELEMENT = new ArrayElementTypeStrategy(); + /** Type strategy specific for {@link BuiltInFunctionDefinitions#ARRAY_FLATTEN}. */ + public static final TypeStrategy ARRAY_FLATTEN = + callContext -> { + // Input type is ARRAY> + DataType inputType = callContext.getArgumentDataTypes().get(0); + + if (!(inputType.getLogicalType() + instanceof org.apache.flink.table.types.logical.ArrayType)) { + return Optional.empty(); + } + + org.apache.flink.table.types.logical.ArrayType outerArrayType = + (org.apache.flink.table.types.logical.ArrayType) inputType.getLogicalType(); + + if (!(outerArrayType.getElementType() + instanceof org.apache.flink.table.types.logical.ArrayType)) { + return Optional.empty(); + } + + org.apache.flink.table.types.logical.ArrayType innerArrayType = + (org.apache.flink.table.types.logical.ArrayType) + outerArrayType.getElementType(); + + // Output type is ARRAY where T is the element type of the inner array + return Optional.of( + DataTypes.ARRAY( + TypeConversions.fromLogicalToDataType( + innerArrayType.getElementType())) + .nullable()); + }; + public static final TypeStrategy ITEM_AT = new ItemAtTypeStrategy(); /** See {@link ArrayAppendPrependTypeStrategy}. */ diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/CollectionFunctionsITCase.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/CollectionFunctionsITCase.java index f67d422c994a6a..1c00b59b303b1f 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/CollectionFunctionsITCase.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/CollectionFunctionsITCase.java @@ -51,6 +51,7 @@ Stream getTestSetSpecs() { arrayReverseTestCases(), arrayUnionTestCases(), arrayConcatTestCases(), + arrayFlattenTestCases(), arrayMaxTestCases(), arrayJoinTestCases(), arraySliceTestCases(), @@ -1900,4 +1901,59 @@ private Stream arrayElementTestCases() { .testResult($("f2").element(), "ELEMENT(f2)", 4.0F, DataTypes.FLOAT()) .testResult($("f3").element(), "ELEMENT(f3)", null, DataTypes.INT())); } + + private Stream arrayFlattenTestCases() { + return Stream.of( + TestSetSpec.forFunction(BuiltInFunctionDefinitions.ARRAY_FLATTEN) + .onFieldsWithData( + new Integer[][] {new Integer[] {1, 2}, new Integer[] {3, 4}}, + new String[][] {new String[] {"a", "b"}, new String[] {"c"}}, + null, + new Integer[][] {new Integer[] {1, 2}, null, new Integer[] {3}}, + new Integer[][] {new Integer[] {1, null, 2}, new Integer[] {3}}, + new Integer[][] {new Integer[] {1}}) + .andDataTypes( + DataTypes.ARRAY(DataTypes.ARRAY(DataTypes.INT())), + DataTypes.ARRAY(DataTypes.ARRAY(DataTypes.STRING())), + DataTypes.ARRAY(DataTypes.ARRAY(DataTypes.INT())), + DataTypes.ARRAY(DataTypes.ARRAY(DataTypes.INT())), + DataTypes.ARRAY(DataTypes.ARRAY(DataTypes.INT())), + DataTypes.ARRAY(DataTypes.ARRAY(DataTypes.INT()))) + // Basic flattening + .testResult( + call("ARRAY_FLATTEN", $("f0")), + "ARRAY_FLATTEN(f0)", + new Integer[] {1, 2, 3, 4}, + DataTypes.ARRAY(DataTypes.INT())) + // String arrays + .testResult( + call("ARRAY_FLATTEN", $("f1")), + "ARRAY_FLATTEN(f1)", + new String[] {"a", "b", "c"}, + DataTypes.ARRAY(DataTypes.STRING())) + // NULL input + .testResult( + call("ARRAY_FLATTEN", $("f2")), + "ARRAY_FLATTEN(f2)", + null, + DataTypes.ARRAY(DataTypes.INT()).nullable()) + // NULL inner arrays - should be skipped + .testResult( + call("ARRAY_FLATTEN", $("f3")), + "ARRAY_FLATTEN(f3)", + new Integer[] {1, 2, 3}, + DataTypes.ARRAY(DataTypes.INT())) + // NULL elements - should be preserved + .testResult( + call("ARRAY_FLATTEN", $("f4")), + "ARRAY_FLATTEN(f4)", + new Integer[] {1, null, 2, 3}, + DataTypes.ARRAY(DataTypes.INT())) + // Single element + .testResult( + call("ARRAY_FLATTEN", $("f5")), + "ARRAY_FLATTEN(f5)", + new Integer[] {1}, + DataTypes.ARRAY(DataTypes.INT()))); + } } diff --git a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/functions/scalar/ArrayFlattenFunction.java b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/functions/scalar/ArrayFlattenFunction.java new file mode 100644 index 00000000000000..87f4770f95a193 --- /dev/null +++ b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/functions/scalar/ArrayFlattenFunction.java @@ -0,0 +1,107 @@ +/* + * 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.flink.table.runtime.functions.scalar; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.table.data.ArrayData; +import org.apache.flink.table.data.GenericArrayData; +import org.apache.flink.table.functions.BuiltInFunctionDefinitions; +import org.apache.flink.table.functions.SpecializedFunction; +import org.apache.flink.table.types.CollectionDataType; +import org.apache.flink.table.types.DataType; +import org.apache.flink.util.FlinkRuntimeException; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.List; + +/** + * Implementation of {@link BuiltInFunctionDefinitions#ARRAY_FLATTEN}. + * + *

Flattens a nested array by one level. + * + *

NULL handling: + * + *

    + *
  • If the input array is NULL, returns NULL + *
  • NULL inner arrays are skipped + *
  • NULL elements within arrays are preserved + *
+ */ +@Internal +public class ArrayFlattenFunction extends BuiltInScalarFunction { + private final ArrayData.ElementGetter outerElementGetter; + private final ArrayData.ElementGetter innerElementGetter; + + public ArrayFlattenFunction(SpecializedFunction.SpecializedContext context) { + super(BuiltInFunctionDefinitions.ARRAY_FLATTEN, context); + + // Get the input data type (ARRAY>) + final DataType inputDataType = context.getCallContext().getArgumentDataTypes().get(0); + // Get the inner array type (ARRAY) + final DataType innerArrayDataType = + ((CollectionDataType) inputDataType).getElementDataType(); + // Get the element type (T) + final DataType elementDataType = + ((CollectionDataType) innerArrayDataType).getElementDataType(); + + // Create element getters + // Outer getter retrieves inner arrays from the outer array + outerElementGetter = ArrayData.createElementGetter(innerArrayDataType.getLogicalType()); + // Inner getter retrieves elements from inner arrays + innerElementGetter = ArrayData.createElementGetter(elementDataType.getLogicalType()); + } + + /** + * Flattens a nested array by one level. + * + * @param array the input array of arrays + * @return the flattened array, or NULL if input is NULL + */ + public @Nullable ArrayData eval(ArrayData array) { + if (array == null) { + return null; + } + + try { + List result = new ArrayList<>(); + + // Iterate through outer array + for (int i = 0; i < array.size(); i++) { + ArrayData innerArray = (ArrayData) outerElementGetter.getElementOrNull(array, i); + + if (innerArray == null) { + // Skip NULL inner arrays + continue; + } + + // Iterate through inner array and add all elements (including NULL) + for (int j = 0; j < innerArray.size(); j++) { + Object element = innerElementGetter.getElementOrNull(innerArray, j); + result.add(element); // Preserve NULL elements + } + } + + return new GenericArrayData(result.toArray()); + } catch (Throwable t) { + throw new FlinkRuntimeException(t); + } + } +}