address comments

snowflakedb · Jan 3, 2024 · 69c1cfa · 69c1cfa
1 parent c5e7c58
commit 69c1cfa
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/src/main/scala/com/snowflake/snowpark/internal/analyzer/MultiChildrenNode.scala b/src/main/scala/com/snowflake/snowpark/internal/analyzer/MultiChildrenNode.scala
@@ -14,11 +14,10 @@ private[snowpark] trait MultiChildrenNode extends LogicalPlan {
   protected def updateChildren(newChildren: Seq[LogicalPlan]): MultiChildrenNode
 
 
-  override lazy val dfAliasMap: Map[String, Seq[Attribute]] = {
-    var result: Map[String, Seq[Attribute]] = Map.empty
-    children.foreach(child => result = Utils.addToDataframeAliasMap(result, child))
-    result
-  }
+  override lazy val dfAliasMap: Map[String, Seq[Attribute]] =
+    children.foldLeft(Map.empty[String, Seq[Attribute]]) {
+      case (map, child) => Utils.addToDataframeAliasMap(map, child)
+    }
 
   override protected def analyze: LogicalPlan =
     createFromAnalyzedChildren(children.map(_.analyzed))

diff --git a/src/test/scala/com/snowflake/snowpark_test/DataFrameAliasSuite.scala b/src/test/scala/com/snowflake/snowpark_test/DataFrameAliasSuite.scala
@@ -38,6 +38,15 @@ class DataFrameAliasSuite extends TestData with BeforeAndAfterEach with EagerSes
     checkAnswer(df1.select($"B.num"), Seq(Row(1), Row(2), Row(3)))
   }
 
+  test("Test for alias with dot in column name") {
+    createTable(tableName1, "\"num.col\" int")
+    runQuery(s"insert into $tableName1 values(1),(2),(3)", session)
+    val df = session.table(tableName1).alias("A")
+    checkAnswer(df.select(df.col("A.num.col")), Seq(Row(1), Row(2), Row(3)))
+    checkAnswer(df.select(col("A.num.col")), Seq(Row(1), Row(2), Row(3)))
+    checkAnswer(df.select($"A.num.col"), Seq(Row(1), Row(2), Row(3)))
+  }
+
   test("Test for alias with join") {
     createTable(tableName1, "id1 int, num1 int")
     createTable(tableName2, "id2 int, num2 int")