From 69c1cfaeaa360da98c2762506ba74c4c6a629a03 Mon Sep 17 00:00:00 2001 From: Zihan Li Date: Wed, 3 Jan 2024 15:12:47 -0800 Subject: [PATCH] address comments --- .../snowpark/internal/analyzer/MultiChildrenNode.scala | 9 ++++----- .../snowflake/snowpark_test/DataFrameAliasSuite.scala | 9 +++++++++ 2 files changed, 13 insertions(+), 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 index c3ad8f0e..b5594184 100644 --- 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 index 397a1000..5deca2c9 100644 --- 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")