[FLINK-38709][table][python] Fix ScalarFunctionSplitter to allow PythonFunction & AsyncFunction work when taking the recursive field of composite type as input

Author: dianfu

flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/RemoteCalcSplitRule.scala

@@ -434,6 +434,8 @@ class ScalarFunctionSplitter(
 
   private var fieldsRexCall: Map[Int, Int] = Map[Int, Int]()
 
+  private val extractedRexNodeRefs: mutable.HashSet[RexNode] = mutable.HashSet[RexNode]()
+
   override def visitCall(call: RexCall): RexNode = {
     if (needConvert(call)) {
       getExtractedRexNode(call)
@@ -454,7 +456,9 @@ class ScalarFunctionSplitter(
           new RexInputRef(field.getIndex, field.getType)
         case _ =>
           val newFieldAccess =
-            rexBuilder.makeFieldAccess(expr.accept(this), fieldAccess.getField.getIndex)
+            rexBuilder.makeFieldAccess(
+              convertInputRefToLocalRefIfNecessary(expr.accept(this)),
+              fieldAccess.getField.getIndex)
           getExtractedRexNode(newFieldAccess)
       }
     } else {
@@ -468,9 +472,18 @@ class ScalarFunctionSplitter(
 
   override def visitNode(rexNode: RexNode): RexNode = rexNode
 
+  private def convertInputRefToLocalRefIfNecessary(node: RexNode): RexNode = {
+    node match {
+      case inputRef: RexInputRef if extractedRexNodeRefs.contains(node) =>
+        new RexLocalRef(inputRef.getIndex, node.getType)
+      case _ => node
+    }
+  }
+
   private def getExtractedRexNode(node: RexNode): RexNode = {
     val newNode = new RexInputRef(extractedFunctionOffset + extractedRexNodes.length, node.getType)
     extractedRexNodes.append(node)
+    extractedRexNodeRefs.add(newNode)
     newNode
   }
 

flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/rules/logical/AsyncCalcSplitRuleTest.java

@@ -62,7 +62,8 @@ public void setup() {
                         + "  a int,\n"
                         + "  b bigint,\n"
                         + "  c string,\n"
-                        + "  d ARRAY<INT NOT NULL>\n"
+                        + "  d ARRAY<INT NOT NULL>,\n"
+                        + "  e ROW<f ROW<h int, i double>, g string>"
                         + ") WITH (\n"
                         + "  'connector' = 'test-simple-table-source'\n"
                         + ") ;");
@@ -89,6 +90,7 @@ public void setup() {
     @Test
     public void testSingleCall() {
         String sqlQuery = "SELECT func1(a) FROM MyTable";
+        util.getTableEnv().explainSql(sqlQuery);
         util.verifyRelPlan(sqlQuery);
     }
 
@@ -182,6 +184,12 @@ public void testFieldAccessAfter() {
         util.verifyRelPlan(sqlQuery);
     }
 
+    @Test
+    public void testCompositeFieldAsInput() {
+        String sqlQuery = "SELECT func1(e.f.h) from MyTable";
+        util.verifyRelPlan(sqlQuery);
+    }
+
     @Test
     public void testFieldOperand() {
         String sqlQuery = "SELECT func1(func5(a).f0) from MyTable";

flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/rules/logical/AsyncCorrelateSplitRuleTest.java

@@ -62,7 +62,8 @@ public void setup() {
                         + "  a int,\n"
                         + "  b bigint,\n"
                         + "  c string,\n"
-                        + "  d ARRAY<INT NOT NULL>\n"
+                        + "  d ARRAY<INT NOT NULL>,\n"
+                        + "  e ROW<f ROW<h int, i double>, g string>\n"
                         + ") WITH (\n"
                         + "  'connector' = 'test-simple-table-source'\n"
                         + ") ;");
@@ -110,6 +111,12 @@ public void testCorrelateWithCast() {
         util.verifyRelPlan(sqlQuery);
     }
 
+    @Test
+    public void testCorrelateWithCompositeFieldAsInput() {
+        String sqlQuery = "select * FROM MyTable, LATERAL TABLE(asyncTableFunc(e.f.h))";
+        util.verifyRelPlan(sqlQuery);
+    }
+
     /** Test function. */
     public static class AsyncFunc extends AsyncTableFunction<String> {