[FLINK-39253][table] Preserve field names in ROW function from AS aliases

Author: autophagy

flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/expressions/resolver/rules/ResolveCallByArgumentsRule.java

@@ -681,6 +681,23 @@ public List<DataType> getArgumentDataTypes() {
                     .collect(Collectors.toList());
         }
 
+        @Override
+        public Optional<String> getArgumentName(int pos) {
+            final ResolvedExpression arg = getArgument(pos);
+
+            if (arg instanceof CallExpression) {
+                final CallExpression call = (CallExpression) arg;
+                if (call.getFunctionDefinition() == BuiltInFunctionDefinitions.AS) {
+                    final List<ResolvedExpression> children = call.getResolvedChildren();
+                    if (children.size() >= 2 && children.get(1) instanceof ValueLiteralExpression) {
+                        return ((ValueLiteralExpression) children.get(1)).getValueAs(String.class);
+                    }
+                }
+            }
+
+            return Optional.empty();
+        }
+
         @Override
         public Optional<DataType> getOutputDataType() {
             return Optional.empty();

flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/CallContext.java

@@ -113,6 +113,11 @@ default Optional<ChangelogMode> getOutputChangelogMode() {
      */
     List<DataType> getArgumentDataTypes();
 
+    /** Returns the name/alias of the argument at the given position if one is available. */
+    default Optional<String> getArgumentName(int pos) {
+        return Optional.empty();
+    }
+
     /**
      * Returns the inferred output data type of the function call.
      *

flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/RowTypeStrategy.java

@@ -37,7 +37,12 @@ public Optional<DataType> inferType(CallContext callContext) {
         List<DataType> argumentDataTypes = callContext.getArgumentDataTypes();
         DataTypes.Field[] fields =
                 IntStream.range(0, argumentDataTypes.size())
-                        .mapToObj(idx -> DataTypes.FIELD("f" + idx, argumentDataTypes.get(idx)))
+                        .mapToObj(
+                                idx -> {
+                                    String fieldName =
+                                            callContext.getArgumentName(idx).orElse("f" + idx);
+                                    return DataTypes.FIELD(fieldName, argumentDataTypes.get(idx));
+                                })
                         .toArray(DataTypes.Field[]::new);
 
         return Optional.of(DataTypes.ROW(fields).notNull());

flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/utils/CastCallContext.java

@@ -115,6 +115,12 @@ public List<DataType> getArgumentDataTypes() {
         return expectedArguments;
     }
 
+    @Override
+    public Optional<String> getArgumentName(int pos) {
+        // argument names remain regardless of casting
+        return originalContext.getArgumentName(pos);
+    }
+
     @Override
     public Optional<DataType> getOutputDataType() {
         return Optional.ofNullable(outputDataType);

flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/RowFunctionITCase.java

@@ -115,7 +115,23 @@ Stream<TestSetSpec> getTestSetSpecs() {
                                                 DataTypes.FIELD("b", DataTypes.TINYINT()),
                                                 DataTypes.FIELD("c", DataTypes.BIGINT()),
                                                 DataTypes.FIELD("d", DataTypes.BOOLEAN()))
-                                        .notNull()));
+                                        .notNull()),
+                TestSetSpec.forFunction(
+                                BuiltInFunctionDefinitions.ROW, "with aliased fields using .as()")
+                        .onFieldsWithData(100, "abc", 75.50)
+                        .andDataTypes(DataTypes.INT(), DataTypes.STRING(), DataTypes.DOUBLE())
+                        .testTableApiResult(
+                                row($("f0").as("a"), $("f1").as("b"), $("f2").as("c")),
+                                Row.of(100, "abc", 75.50),
+                                DataTypes.ROW(
+                                                DataTypes.FIELD("a", DataTypes.INT()),
+                                                DataTypes.FIELD("b", DataTypes.STRING()),
+                                                DataTypes.FIELD("c", DataTypes.DOUBLE()))
+                                        .notNull())
+                        .testTableApiResult(
+                                row($("f0").as("a"), $("f1").as("b"), $("f2").as("c")).get("a"),
+                                100,
+                                DataTypes.INT()));
     }
 
     // --------------------------------------------------------------------------------------------