[FLINK-27891][Table] merge arrayelement strategies

apache · Mar 14, 2023 · 36c13d9 · 36c13d9
1 parent 0895252
commit 36c13d9
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 74 deletions.
diff --git a/...ble-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java b/...ble-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java
@@ -87,9 +87,9 @@
 import static org.apache.flink.table.types.inference.TypeStrategies.nullableIfAllArgs;
 import static org.apache.flink.table.types.inference.TypeStrategies.nullableIfArgs;
 import static org.apache.flink.table.types.inference.TypeStrategies.varyingString;
-import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.FIRST_ARRAY_ELEMENT_ARG;
+import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.ARRAY_ELEMENT_ARG;
+import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.ARRAY_ELEMENT_ARG1;
 import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.JSON_ARGUMENT;
-import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.LAST_ARRAY_ELEMENT_ARG;
 import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.TWO_EQUALS_COMPARABLE;
 import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.TWO_FULLY_COMPARABLE;
 
@@ -159,8 +159,7 @@ ANY, and(logical(LogicalTypeRoot.BOOLEAN), LITERAL)
                             sequence(
                                     Arrays.asList("haystack", "needle"),
                                     Arrays.asList(
-                                            logical(LogicalTypeRoot.ARRAY),
-                                            FIRST_ARRAY_ELEMENT_ARG)))
+                                            logical(LogicalTypeRoot.ARRAY), ARRAY_ELEMENT_ARG)))
                     .outputTypeStrategy(
                             nullableIfArgs(
                                     ConstantArgumentCount.of(0), explicit(DataTypes.BOOLEAN())))
@@ -189,8 +188,7 @@ ANY, and(logical(LogicalTypeRoot.BOOLEAN), LITERAL)
                             sequence(
                                     Arrays.asList("haystack", "element"),
                                     Arrays.asList(
-                                            logical(LogicalTypeRoot.ARRAY),
-                                            FIRST_ARRAY_ELEMENT_ARG)))
+                                            logical(LogicalTypeRoot.ARRAY), ARRAY_ELEMENT_ARG)))
                     .outputTypeStrategy(nullableIfArgs(argument(0)))
                     .runtimeClass(
                             "org.apache.flink.table.runtime.functions.scalar.ArrayAppendFunction")
@@ -204,8 +202,7 @@ ANY, and(logical(LogicalTypeRoot.BOOLEAN), LITERAL)
                             sequence(
                                     Arrays.asList("element", "haystack"),
                                     Arrays.asList(
-                                            LAST_ARRAY_ELEMENT_ARG,
-                                            logical(LogicalTypeRoot.ARRAY))))
+                                            ARRAY_ELEMENT_ARG1, logical(LogicalTypeRoot.ARRAY))))
                     .outputTypeStrategy(nullableIfArgs(argument(1)))
                     .runtimeClass(
                             "org.apache.flink.table.runtime.functions.scalar.ArrayPrependFunction")

diff --git a/...irstArrayElementArgumentTypeStrategy.java → ...ies/ArrayElementArgumentTypeStrategy.java b/...irstArrayElementArgumentTypeStrategy.java → ...ies/ArrayElementArgumentTypeStrategy.java
@@ -35,13 +35,19 @@
 
 /** Specific {@link ArgumentTypeStrategy} for {@link BuiltInFunctionDefinitions#ARRAY_CONTAINS}. */
 @Internal
-class FirstArrayElementArgumentTypeStrategy implements ArgumentTypeStrategy {
+class ArrayElementArgumentTypeStrategy implements ArgumentTypeStrategy {
+
+    private final int arrayPosition;
+
+    public ArrayElementArgumentTypeStrategy(int arrayPosition) {
+        this.arrayPosition = arrayPosition;
+    }
 
     @Override
     public Optional<DataType> inferArgumentType(
             CallContext callContext, int argumentPos, boolean throwOnFailure) {
         final ArrayType haystackType =
-                (ArrayType) callContext.getArgumentDataTypes().get(0).getLogicalType();
+                (ArrayType) callContext.getArgumentDataTypes().get(arrayPosition).getLogicalType();
         final LogicalType needleType =
                 callContext.getArgumentDataTypes().get(argumentPos).getLogicalType();
         LogicalType haystackElementType = haystackType.getElementType();

diff --git a/...g/apache/flink/table/types/inference/strategies/LastArrayElementArgumentTypeStrategy.java b/...g/apache/flink/table/types/inference/strategies/LastArrayElementArgumentTypeStrategy.java
diff --git a/...n/java/org/apache/flink/table/types/inference/strategies/SpecificInputTypeStrategies.java b/...n/java/org/apache/flink/table/types/inference/strategies/SpecificInputTypeStrategies.java
@@ -70,11 +70,11 @@ public final class SpecificInputTypeStrategies {
                     logical(LogicalTypeFamily.NUMERIC));
 
     /** Argument type derived from the array element type. */
-    public static final ArgumentTypeStrategy FIRST_ARRAY_ELEMENT_ARG =
-            new FirstArrayElementArgumentTypeStrategy();
+    public static final ArgumentTypeStrategy ARRAY_ELEMENT_ARG =
+            new ArrayElementArgumentTypeStrategy(0);
 
-    public static final ArgumentTypeStrategy LAST_ARRAY_ELEMENT_ARG =
-            new LastArrayElementArgumentTypeStrategy();
+    public static final ArgumentTypeStrategy ARRAY_ELEMENT_ARG1 =
+            new ArrayElementArgumentTypeStrategy(1);
 
     /**
      * Input strategy for {@link BuiltInFunctionDefinitions#JSON_OBJECT}.