Fix

Author: maropu

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.catalyst.expressions
 
 import java.util.Comparator
-import java.util.concurrent.atomic.AtomicReference
+import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
 
 import scala.collection.mutable
 
@@ -54,6 +54,16 @@ case class UnresolvedNamedLambdaVariable(nameParts: Seq[String])
   override def sql: String = name
 }
 
+object UnresolvedNamedLambdaVariable {
+
+  // Counter to ensure lambda variable names are unique
+  private val nextVarNameId = new AtomicInteger(0)
+
+  def freshVarName(name: String): String = {
+    s"${name}_${nextVarNameId.getAndIncrement()}"
+  }
+}
+
 /**
  * A named lambda variable.
  */

sql/core/src/main/scala/org/apache/spark/sql/functions.scala

@@ -17,8 +17,6 @@
 
 package org.apache.spark.sql
 
-import java.util.concurrent.atomic.AtomicInteger
-
 import scala.collection.JavaConverters._
 import scala.reflect.runtime.universe.TypeTag
 import scala.util.Try
@@ -3801,26 +3799,23 @@ object functions {
     ArrayExcept(col1.expr, col2.expr)
   }
 
-  // counter to ensure lambda variable names are unique
-  private val lambdaVarNameCounter = new AtomicInteger(0)
-
   private def createLambda(f: Column => Column) = {
-    val x = UnresolvedNamedLambdaVariable(Seq("x_" + lambdaVarNameCounter.incrementAndGet()))
+    val x = UnresolvedNamedLambdaVariable(Seq(UnresolvedNamedLambdaVariable.freshVarName("x")))
     val function = f(Column(x)).expr
     LambdaFunction(function, Seq(x))
   }
 
   private def createLambda(f: (Column, Column) => Column) = {
-    val x = UnresolvedNamedLambdaVariable(Seq("x_" + lambdaVarNameCounter.incrementAndGet()))
-    val y = UnresolvedNamedLambdaVariable(Seq("y_" + lambdaVarNameCounter.incrementAndGet()))
+    val x = UnresolvedNamedLambdaVariable(Seq(UnresolvedNamedLambdaVariable.freshVarName("x")))
+    val y = UnresolvedNamedLambdaVariable(Seq(UnresolvedNamedLambdaVariable.freshVarName("y")))
     val function = f(Column(x), Column(y)).expr
     LambdaFunction(function, Seq(x, y))
   }
 
   private def createLambda(f: (Column, Column, Column) => Column) = {
-    val x = UnresolvedNamedLambdaVariable(Seq("x_" + lambdaVarNameCounter.incrementAndGet()))
-    val y = UnresolvedNamedLambdaVariable(Seq("y_" + lambdaVarNameCounter.incrementAndGet()))
-    val z = UnresolvedNamedLambdaVariable(Seq("z_" + lambdaVarNameCounter.incrementAndGet()))
+    val x = UnresolvedNamedLambdaVariable(Seq(UnresolvedNamedLambdaVariable.freshVarName("x")))
+    val y = UnresolvedNamedLambdaVariable(Seq(UnresolvedNamedLambdaVariable.freshVarName("y")))
+    val z = UnresolvedNamedLambdaVariable(Seq(UnresolvedNamedLambdaVariable.freshVarName("z")))
     val function = f(Column(x), Column(y), Column(z)).expr
     LambdaFunction(function, Seq(x, y, z))
   }

sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala

@@ -2261,32 +2261,6 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSparkSession {
     assert(ex3.getMessage.contains("cannot resolve 'a'"))
   }
 
-  test("nested transform (DSL)") {
-    val df = Seq(
-        (Seq(1, 2, 3), Seq("a", "b", "c"))
-    ).toDF("numbers", "letters")
-
-    checkAnswer(
-        df.select(
-          flatten(
-            transform(
-              $"numbers",
-              (number: Column) => transform(
-                $"letters",
-                (letter: Column) => struct(
-                  number.as("number"),
-                  letter.as("letter")
-                )
-              )
-            )
-          ).as("zipped")
-        ),
-        Seq(Row(Seq(Row(1, "a"), Row(1, "b"), Row(1, "c"), Row(2, "a"), Row(2, "b"),
-            Row(2, "c"), Row(3, "a"), Row(3, "b"), Row(3, "c")
-        )))
-    )
-  }
-
   test("map_filter") {
     val dfInts = Seq(
       Map(1 -> 10, 2 -> 20, 3 -> 30),
@@ -3655,6 +3629,29 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSparkSession {
       df.select(map(map_entries($"m"), lit(1))),
       Row(Map(Seq(Row(1, "a")) -> 1)))
   }
+
+  test("SPARK-34794: lambda variable name issues in nested functions") {
+    val df1 = Seq((Seq(1, 2), Seq("a", "b"))).toDF("numbers", "letters")
+
+    checkAnswer(df1.select(flatten(transform($"numbers", (number: Column) =>
+        transform($"letters", (letter: Column) =>
+          struct(number, letter))))),
+        Seq(Row(Seq(Row(1, "a"), Row(1, "b"), Row(2, "a"), Row(2, "b"))))
+    )
+    checkAnswer(df1.select(flatten(transform($"numbers", (number: Column, i: Column) =>
+        transform($"letters", (letter: Column, j: Column) =>
+          struct(number + j, concat(letter, i)))))),
+        Seq(Row(Seq(Row(1, "a0"), Row(2, "b0"), Row(2, "a1"), Row(3, "b1"))))
+    )
+
+    val df2 = Seq((Map("a" -> 1, "b" -> 2), Map("a" -> 2, "b" -> 3))).toDF("m1", "m2")
+
+    checkAnswer(df2.select(map_zip_with($"m1", $"m2", (k1: Column, ov1: Column, ov2: Column) =>
+        map_zip_with($"m1", $"m2", (k2: Column, iv1: Column, iv2: Column) =>
+          ov1 + iv1 + ov2 + iv2))),
+        Seq(Row(Map("a" -> Map("a" -> 6, "b" -> 8), "b" -> Map("a" -> 8, "b" -> 10))))
+    )
+  }
 }
 
 object DataFrameFunctionsSuite {