[FLINK-39531][table-planner] Fix duplicated Python UDF extraction in calc split rule

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

@@ -432,9 +432,7 @@ class ScalarFunctionSplitter(
     callFinder: RemoteCallFinder)
   extends RexDefaultVisitor[RexNode] {
 
-  private var fieldsRexCall: Map[Int, Int] = Map[Int, Int]()
-
-  private val extractedRexNodeRefs: mutable.HashSet[RexNode] = mutable.HashSet[RexNode]()
+  private val extractedRexNodeToIndex = mutable.HashMap.empty[RexNode, Int]
 
   override def visitCall(call: RexCall): RexNode = {
     if (needConvert(call)) {
@@ -474,30 +472,33 @@ class ScalarFunctionSplitter(
 
   private def convertInputRefToLocalRefIfNecessary(node: RexNode): RexNode = {
     node match {
-      case inputRef: RexInputRef if extractedRexNodeRefs.contains(node) =>
+      case inputRef: RexInputRef if inputRef.getIndex >= extractedFunctionOffset =>
         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
+    new RexInputRef(extractedFunctionOffset + getExtractedRexNodeIndex(node), node.getType)
   }
 
   private def getExtractedRexFieldAccess(node: RexFieldAccess, rexCallIndex: Int): RexNode = {
     val remoteCall: RexCall =
       program.expandLocalRef(node.getReferenceExpr.asInstanceOf[RexLocalRef]).asInstanceOf[RexCall]
-    if (!fieldsRexCall.contains(rexCallIndex)) {
-      extractedRexNodes.append(remoteCall)
-      fieldsRexCall += rexCallIndex -> (extractedFunctionOffset + extractedRexNodes.length - 1)
-    }
     rexBuilder.makeFieldAccess(
-      new RexInputRef(fieldsRexCall(rexCallIndex), remoteCall.getType),
+      new RexInputRef(
+        extractedFunctionOffset + getExtractedRexNodeIndex(remoteCall),
+        remoteCall.getType),
       node.getField.getIndex)
   }
+
+  private def getExtractedRexNodeIndex(node: RexNode): Int = {
+    extractedRexNodeToIndex.getOrElseUpdate(
+      node, {
+        extractedRexNodes.append(node)
+        extractedRexNodes.length - 1
+      })
+  }
 }
 
 /**

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

@@ -716,12 +716,10 @@ LogicalProject(EXPR$0=[func1(func1($0))], EXPR$1=[func1(func1(func1($0)))], EXPR
     </Resource>
     <Resource name="optimized rel plan">
       <![CDATA[
-AsyncCalc(select=[f0 AS EXPR$0, func1(f1) AS EXPR$1, f2 AS EXPR$2])
-+- AsyncCalc(select=[f2, f0, func1(f1) AS f1])
-   +- AsyncCalc(select=[f2, f1, func1(f0) AS f0])
-      +- Calc(select=[f0, f0 AS f1, f0 AS f2])
-         +- AsyncCalc(select=[func1(a) AS f0])
-            +- TableSourceScan(table=[[default_catalog, default_database, MyTable]], fields=[a, b, c, d, e])
+AsyncCalc(select=[f00 AS EXPR$0, func1(f00) AS EXPR$1, f0 AS EXPR$2])
++- AsyncCalc(select=[f0, func1(f0) AS f00])
+   +- AsyncCalc(select=[func1(a) AS f0])
+      +- TableSourceScan(table=[[default_catalog, default_database, MyTable]], fields=[a, b, c, d, e])
 ]]>
     </Resource>
   </TestCase>

flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRuleTest.xml

@@ -520,6 +520,24 @@ LogicalProject(a=[$0], EXPR$1=[pyFunc1($0, $2)], b=[$1])
 FlinkLogicalCalc(select=[a, f0 AS EXPR$1, b])
 +- FlinkLogicalCalc(select=[a, b, pyFunc1(a, c) AS f0])
    +- FlinkLogicalTableSourceScan(table=[[default_catalog, default_database, MyTable]], fields=[a, b, c, d])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testSamePythonFunctionUsedInMultipleProjectionExpressions">
+    <Resource name="sql">
+      <![CDATA[SELECT pyFunc1(a, c) + 1, pyFunc1(a, c) + 2 FROM MyTable]]>
+    </Resource>
+    <Resource name="ast">
+      <![CDATA[
+LogicalProject(EXPR$0=[+(pyFunc1($0, $2), 1)], EXPR$1=[+(pyFunc1($0, $2), 2)])
++- LogicalTableScan(table=[[default_catalog, default_database, MyTable]])
+]]>
+    </Resource>
+    <Resource name="optimized rel plan">
+      <![CDATA[
+FlinkLogicalCalc(select=[+(f0, 1) AS EXPR$0, +(f0, 2) AS EXPR$1])
++- FlinkLogicalCalc(select=[pyFunc1(a, c) AS f0])
+   +- FlinkLogicalTableSourceScan(table=[[default_catalog, default_database, MyTable]], fields=[a, b, c, d])
 ]]>
     </Resource>
   </TestCase>

flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRuleTest.scala

@@ -120,6 +120,12 @@ class PythonCalcSplitRuleTest extends TableTestBase {
     util.verifyRelPlan(sqlQuery)
   }
 
+  @Test
+  def testSamePythonFunctionUsedInMultipleProjectionExpressions(): Unit = {
+    val sqlQuery = "SELECT pyFunc1(a, c) + 1, pyFunc1(a, c) + 2 FROM MyTable"
+    util.verifyRelPlan(sqlQuery)
+  }
+
   @Test
   def testReorderPythonCalc(): Unit = {
     val sqlQuery = "SELECT a, pyFunc1(a, c), b FROM MyTable"