Clean expression stack before resolving project list (#1059)

LantaoJin · web-flow · commit 09aba20a7f31 · 2025-02-19T10:31:22.000+08:00
Signed-off-by: Lantao Jin &lt;ltjin@amazon.com&gt;
diff --git a/integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLBasicITSuite.scala b/integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLBasicITSuite.scala
@@ -8,7 +8,7 @@ package org.opensearch.flint.spark.ppl
 import org.apache.spark.sql.{AnalysisException, QueryTest, Row}
 import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedFunction, UnresolvedRelation, UnresolvedStar}
-import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Descending, EqualTo, IsNotNull, Literal, Not, SortOrder}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Descending, EqualTo, GreaterThanOrEqual, IsNotNull, Literal, Not, SortOrder}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.ExplainMode
 import org.apache.spark.sql.execution.command.{DescribeTableCommand, ExplainCommand}
@@ -671,4 +671,22 @@ class FlintSparkPPLBasicITSuite
 
     comparePlans(expectedPlan, logicalPlan, checkAnalysis = false)
   }
+
+  test("test fields with alias") {
+    val frame = sql(s"""
+         | source = $testTable as l | where l.age >= 30 | fields l.name, l.age
+         | """.stripMargin)
+
+    val logicalPlan: LogicalPlan = frame.queryExecution.logical
+    val expectedPlan = Project(
+      Seq(UnresolvedAttribute("l.name"), UnresolvedAttribute("l.age")),
+      Filter(
+        GreaterThanOrEqual(UnresolvedAttribute("l.age"), Literal(30)),
+        SubqueryAlias(
+          "l",
+          UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test")))))
+    comparePlans(expectedPlan, logicalPlan, checkAnalysis = false)
+    assertSameRows(Seq(Row("Jake", 70), Row("Hello", 30)), frame)
+
+  }
 }
diff --git a/integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLInSubqueryITSuite.scala b/integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLInSubqueryITSuite.scala
@@ -475,4 +475,45 @@ class FlintSparkPPLInSubqueryITSuite
     assert(ex.getMessage.contains(
       "The number of columns in the left hand side of an IN subquery does not match the number of columns in the output of subquery"))
   }
+
+  test("test in subquery with table alias") {
+    val frame = sql(s"""
+         | source = $outerTable as o
+         | | where id in [
+         |   source = $innerTable as i
+         | | where i.department = 'DATA'
+         | | fields i.uid
+         | ]
+         | | sort  - o.salary
+         | | fields o.id, o.name, o.salary
+         | """.stripMargin)
+    val expectedResults: Array[Row] = Array(Row(1002, "John", 120000), Row(1005, "Jane", 90000))
+    assertSameRows(expectedResults, frame)
+
+    val logicalPlan: LogicalPlan = frame.queryExecution.logical
+    val outer =
+      SubqueryAlias("o", UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test1")))
+    val inner =
+      SubqueryAlias("i", UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test2")))
+    val inSubquery =
+      Filter(
+        InSubquery(
+          Seq(UnresolvedAttribute("id")),
+          ListQuery(
+            Project(
+              Seq(UnresolvedAttribute("i.uid")),
+              Filter(EqualTo(UnresolvedAttribute("i.department"), Literal("DATA")), inner)))),
+        outer)
+    val sortedPlan: LogicalPlan =
+      Sort(Seq(SortOrder(UnresolvedAttribute("o.salary"), Descending)), global = true, inSubquery)
+    val expectedPlan =
+      Project(
+        Seq(
+          UnresolvedAttribute("o.id"),
+          UnresolvedAttribute("o.name"),
+          UnresolvedAttribute("o.salary")),
+        sortedPlan)
+
+    comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
+  }
 }
diff --git a/integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLScalarSubqueryITSuite.scala b/integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLScalarSubqueryITSuite.scala
@@ -481,4 +481,103 @@ class FlintSparkPPLScalarSubqueryITSuite
     implicit val rowOrdering: Ordering[Row] = Ordering.by[Row, Integer](_.getAs[Integer](0))
     assert(results.sorted.sameElements(expectedResults.sorted))
   }
+
+  test("test nested scalar subquery with table alias") {
+    val frame = sql(s"""
+                       | source = $outerTable as o
+                       | | where id = [
+                       |     source = $innerTable as i
+                       |     | where uid = [
+                       |         source = $nestedInnerTable as n
+                       |         | stats min(n.salary)
+                       |       ] + 1000
+                       |     | sort i.department
+                       |     | stats max(i.uid)
+                       |   ]
+                       | | fields o.id, o.name
+                       | """.stripMargin)
+    val expectedResults: Array[Row] = Array(Row(1000, "Jake"))
+    assertSameRows(expectedResults, frame)
+  }
+
+  test("test correlated scalar subquery with table alias") {
+    val frame = sql(s"""
+                       | source = $outerTable as o
+                       | | where id = [
+                       |     source = $innerTable as i | where o.id = i.uid | stats max(i.uid)
+                       |   ]
+                       | | fields o.id, o.name
+                       | """.stripMargin)
+    val expectedResults: Array[Row] = Array(
+      Row(1000, "Jake"),
+      Row(1002, "John"),
+      Row(1003, "David"),
+      Row(1005, "Jane"),
+      Row(1006, "Tommy"))
+    assertSameRows(expectedResults, frame)
+
+    val logicalPlan: LogicalPlan = frame.queryExecution.logical
+
+    val outer =
+      SubqueryAlias("o", UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test1")))
+    val inner =
+      SubqueryAlias("i", UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test2")))
+    val aggregateExpressions = Seq(
+      Alias(
+        UnresolvedFunction(Seq("MAX"), Seq(UnresolvedAttribute("i.uid")), isDistinct = false),
+        "max(i.uid)")())
+    val innerFilter =
+      Filter(EqualTo(UnresolvedAttribute("o.id"), UnresolvedAttribute("i.uid")), inner)
+    val aggregatePlan = Aggregate(Seq(), aggregateExpressions, innerFilter)
+    val scalarSubqueryExpr = ScalarSubquery(aggregatePlan)
+    val outerFilter = Filter(EqualTo(UnresolvedAttribute("id"), scalarSubqueryExpr), outer)
+    val expectedPlan =
+      Project(Seq(UnresolvedAttribute("o.id"), UnresolvedAttribute("o.name")), outerFilter)
+
+    comparePlans(expectedPlan, logicalPlan, checkAnalysis = false)
+  }
+
+  test("test uncorrelated scalar subquery with table alias") {
+    val frame = sql(s"""
+                       | source = $outerTable as o
+                       | | eval max_uid = [
+                       |     source = $innerTable as i | where i.department = 'DATA' | stats max(i.uid)
+                       |   ]
+                       | | fields o.id, o.name, max_uid
+                       | """.stripMargin)
+    val expectedResults: Array[Row] = Array(
+      Row(1000, "Jake", 1005),
+      Row(1001, "Hello", 1005),
+      Row(1002, "John", 1005),
+      Row(1003, "David", 1005),
+      Row(1004, "David", 1005),
+      Row(1005, "Jane", 1005),
+      Row(1006, "Tommy", 1005))
+    assertSameRows(expectedResults, frame)
+
+    val logicalPlan: LogicalPlan = frame.queryExecution.logical
+
+    val outer =
+      SubqueryAlias("o", UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test1")))
+    val inner =
+      SubqueryAlias("i", UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test2")))
+    val aggregateExpressions = Seq(
+      Alias(
+        UnresolvedFunction(Seq("MAX"), Seq(UnresolvedAttribute("i.uid")), isDistinct = false),
+        "max(i.uid)")())
+    val innerFilter =
+      Filter(EqualTo(UnresolvedAttribute("i.department"), Literal("DATA")), inner)
+    val aggregatePlan = Aggregate(Seq(), aggregateExpressions, innerFilter)
+    val scalarSubqueryExpr = Alias(ScalarSubquery(aggregatePlan), "max_uid")()
+    val outerFilter = Project(Seq(UnresolvedStar(None), scalarSubqueryExpr), outer)
+    val expectedPlan =
+      Project(
+        Seq(
+          UnresolvedAttribute("o.id"),
+          UnresolvedAttribute("o.name"),
+          UnresolvedAttribute("max_uid")),
+        outerFilter)
+
+    comparePlans(expectedPlan, logicalPlan, checkAnalysis = false)
+  }
 }
diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystExpressionVisitor.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystExpressionVisitor.java
@@ -349,7 +349,7 @@ public Expression visitCase(Case node, CatalystPlanContext context) {
                         )
                 );
             }
-            context.retainAllNamedParseExpressions(e -> e);
+            context.resetNamedParseExpressions();
         }
         context.setNamedParseExpressions(initialNameExpressions);
         return context.getNamedParseExpressions().push(new CaseWhen(seq(whens), Option.apply(elseValue)));
@@ -421,7 +421,7 @@ public Expression visitBetween(Between node, CatalystPlanContext context) {
         Expression value = analyze(node.getValue(), context);
         Expression lower = analyze(node.getLowerBound(), context);
         Expression upper = analyze(node.getUpperBound(), context);
-        context.retainAllNamedParseExpressions(p -> p);
+        context.resetNamedParseExpressions();
         return context.getNamedParseExpressions().push(new org.apache.spark.sql.catalyst.expressions.And(new GreaterThanOrEqual(value, lower), new LessThanOrEqual(value, upper)));
     }
 
diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystPlanContext.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystPlanContext.java
@@ -8,15 +8,11 @@
 import lombok.Getter;
 import org.apache.spark.sql.SparkSession;
 import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation;
-import org.apache.spark.sql.catalyst.expressions.AttributeReference;
 import org.apache.spark.sql.catalyst.expressions.Expression;
-import org.apache.spark.sql.catalyst.expressions.NamedExpression;
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan;
 import org.apache.spark.sql.catalyst.plans.logical.SubqueryAlias;
 import org.apache.spark.sql.catalyst.plans.logical.Union;
-import org.apache.spark.sql.types.Metadata;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
-import org.opensearch.sql.data.type.ExprType;
 import scala.collection.Iterator;
 import scala.collection.Seq;
 
@@ -27,7 +23,6 @@
 import java.util.Stack;
 import java.util.function.BiFunction;
 import java.util.function.Function;
-import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyList;
@@ -218,6 +213,14 @@ public <T> Seq<T> retainAllPlans(Function<LogicalPlan, T> transformFunction) {
         return plans;
     }
 
+    /**
+     * Reset all expressions in stack,
+     * generally use it after calling visitFirstChild() in visit methods.
+     */
+    public void resetNamedParseExpressions() {
+        getNamedParseExpressions().retainAll(emptyList());
+    }
+
     /**
      * retain all expressions and clear expression stack
      *
@@ -226,7 +229,7 @@ public <T> Seq<T> retainAllPlans(Function<LogicalPlan, T> transformFunction) {
     public <T> Seq<T> retainAllNamedParseExpressions(Function<Expression, T> transformFunction) {
         Seq<T> aggregateExpressions = seq(getNamedParseExpressions().stream()
                 .map(transformFunction).collect(Collectors.toList()));
-        getNamedParseExpressions().retainAll(emptyList());
+        resetNamedParseExpressions();
         return aggregateExpressions;
     }
 
diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystQueryPlanVisitor.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystQueryPlanVisitor.java
@@ -89,12 +89,14 @@
 import scala.collection.Seq;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static java.util.Collections.emptyList;
 import static java.util.List.of;
@@ -196,7 +198,7 @@ public LogicalPlan visitFilter(Filter node, CatalystPlanContext context) {
     public LogicalPlan visitLookup(Lookup node, CatalystPlanContext context) {
         visitFirstChild(node, context);
         return context.apply( searchSide -> {
-            context.retainAllNamedParseExpressions(p -> p);
+            context.resetNamedParseExpressions();
             context.retainAllPlans(p -> p);
             LogicalPlan target;
             LogicalPlan lookupTable = node.getLookupRelation().accept(this, context);
@@ -257,7 +259,7 @@ public LogicalPlan visitLookup(Lookup node, CatalystPlanContext context) {
 
             LogicalPlan outputWithDropped = DataFrameDropColumns$.MODULE$.apply(seq(toDrop), target);
 
-            context.retainAllNamedParseExpressions(p -> p);
+            context.resetNamedParseExpressions();
             context.retainAllPlans(p -> p);
             return outputWithDropped;
         });
@@ -304,7 +306,7 @@ public LogicalPlan visitAppendCol(AppendCol node, CatalystPlanContext context) {
             var subSearchWithRowNumber = getRowNumStarProjection(context, subSearch, TABLE_RHS);
 
             context.withSubqueryAlias(subSearchWithRowNumber);
-            context.retainAllNamedParseExpressions(p -> p);
+            context.resetNamedParseExpressions();
             context.retainAllPlans(p -> p);
 
             // Join both Main and Sub search with _ROW_NUMBER_ column
@@ -351,7 +353,7 @@ public LogicalPlan visitJoin(Join node, CatalystPlanContext context) {
             LogicalPlan right = node.getRight().accept(this, context);
             Optional<Expression> joinCondition = node.getJoinCondition()
                 .map(c -> expressionAnalyzer.analyzeJoinCondition(c, context));
-            context.retainAllNamedParseExpressions(p -> p);
+            context.resetNamedParseExpressions();
             context.retainAllPlans(p -> p);
             return join(left, right, node.getJoinType(), joinCondition, node.getJoinHint());
         });
@@ -371,6 +373,8 @@ public LogicalPlan visitSubqueryAlias(SubqueryAlias node, CatalystPlanContext co
     @Override
     public LogicalPlan visitAggregation(Aggregation node, CatalystPlanContext context) {
         visitFirstChild(node, context);
+        // clean before to go
+        context.resetNamedParseExpressions();
         List<Expression> aggsExpList = visitExpressionList(node.getAggExprList(), context);
         List<Expression> groupExpList = visitExpressionList(node.getGroupExprList(), context);
         if (!groupExpList.isEmpty()) {
@@ -460,6 +464,9 @@ public LogicalPlan visitProject(Project node, CatalystPlanContext context) {
             context.withProjectedFields(node.getProjectList());
         }
         LogicalPlan child = visitFirstChild(node, context);
+
+        // reset expression stack before resolving project
+        context.resetNamedParseExpressions();
         visitExpressionList(node.getProjectList(), context);
 
         // Create a projection list from the existing expressions
@@ -481,6 +488,7 @@ public LogicalPlan visitProject(Project node, CatalystPlanContext context) {
     @Override
     public LogicalPlan visitSort(Sort node, CatalystPlanContext context) {
         visitFirstChild(node, context);
+        context.resetNamedParseExpressions();
         visitFieldList(node.getSortList(), context);
         Seq<SortOrder> sortElements = context.retainAllNamedParseExpressions(exp -> SortUtils.getSortDirection(node, (NamedExpression) exp));
         return context.apply(p -> (LogicalPlan) new org.apache.spark.sql.catalyst.plans.logical.Sort(sortElements, true, p));
diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/utils/GeoIpCatalystLogicalPlanTranslator.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/utils/GeoIpCatalystLogicalPlanTranslator.java
@@ -116,7 +116,7 @@ static private LogicalPlan applyJoin(Expression ipAddress, CatalystPlanContext c
                             UnresolvedAttribute$.MODULE$.apply(seq(GEOIP_TABLE_ALIAS,GEOIP_IPV4_COLUMN_NAME))
                     )
             ));
-            context.retainAllNamedParseExpressions(p -> p);
+            context.resetNamedParseExpressions();
             context.retainAllPlans(p -> p);
             return join(leftAlias,
                     rightAlias,
diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/utils/LookupTransformer.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/utils/LookupTransformer.java
@@ -75,7 +75,7 @@ static Expression buildLookupMappingCondition(
             Expression equalTo = EqualTo$.MODULE$.apply(lookupNamedExpression, sourceNamedExpression);
             equiConditions.add(equalTo);
         }
-        context.retainAllNamedParseExpressions(e -> e);
+        context.resetNamedParseExpressions();
         return equiConditions.stream().reduce(org.apache.spark.sql.catalyst.expressions.And::new).orElse(null);
     }
 
@@ -114,7 +114,7 @@ static List<NamedExpression> buildOutputProjectList(
                 seq(new java.util.ArrayList<String>()));
             outputProjectList.add(output);
         }
-        context.retainAllNamedParseExpressions(p -> p);
+        context.resetNamedParseExpressions();
         return outputProjectList;
     }
 

Original file line number	Diff line number	Diff line change
`@@ -349,7 +349,7 @@ public Expression visitCase(Case node, CatalystPlanContext context) {`
`349`	`349`	`)`
`350`	`350`	`);`
`351`	`351`	`}`
`352`		`- context.retainAllNamedParseExpressions(e -> e);`
	`352`	`+ context.resetNamedParseExpressions();`
`353`	`353`	`}`
`354`	`354`	`context.setNamedParseExpressions(initialNameExpressions);`
`355`	`355`	`return context.getNamedParseExpressions().push(new CaseWhen(seq(whens), Option.apply(elseValue)));`
`@@ -421,7 +421,7 @@ public Expression visitBetween(Between node, CatalystPlanContext context) {`
`421`	`421`	`Expression value = analyze(node.getValue(), context);`
`422`	`422`	`Expression lower = analyze(node.getLowerBound(), context);`
`423`	`423`	`Expression upper = analyze(node.getUpperBound(), context);`
`424`		`- context.retainAllNamedParseExpressions(p -> p);`
	`424`	`+ context.resetNamedParseExpressions();`
`425`	`425`	`return context.getNamedParseExpressions().push(new org.apache.spark.sql.catalyst.expressions.And(new GreaterThanOrEqual(value, lower), new LessThanOrEqual(value, upper)));`
`426`	`426`	`}`
`427`	`427`
Original file line number	Diff line number	Diff line change
`@@ -116,7 +116,7 @@ static private LogicalPlan applyJoin(Expression ipAddress, CatalystPlanContext c`
`116`	`116`	`UnresolvedAttribute$.MODULE$.apply(seq(GEOIP_TABLE_ALIAS,GEOIP_IPV4_COLUMN_NAME))`
`117`	`117`	`)`
`118`	`118`	`));`
`119`		`- context.retainAllNamedParseExpressions(p -> p);`
	`119`	`+ context.resetNamedParseExpressions();`
`120`	`120`	`context.retainAllPlans(p -> p);`
`121`	`121`	`return join(leftAlias,`
`122`	`122`	`rightAlias,`
Original file line number	Diff line number	Diff line change
`@@ -75,7 +75,7 @@ static Expression buildLookupMappingCondition(`
`75`	`75`	`Expression equalTo = EqualTo$.MODULE$.apply(lookupNamedExpression, sourceNamedExpression);`
`76`	`76`	`equiConditions.add(equalTo);`
`77`	`77`	`}`
`78`		`- context.retainAllNamedParseExpressions(e -> e);`
	`78`	`+ context.resetNamedParseExpressions();`
`79`	`79`	`return equiConditions.stream().reduce(org.apache.spark.sql.catalyst.expressions.And::new).orElse(null);`
`80`	`80`	`}`
`81`	`81`
`@@ -114,7 +114,7 @@ static List<NamedExpression> buildOutputProjectList(`
`114`	`114`	`seq(new java.util.ArrayList<String>()));`
`115`	`115`	`outputProjectList.add(output);`
`116`	`116`	`}`
`117`		`- context.retainAllNamedParseExpressions(p -> p);`
	`117`	`+ context.resetNamedParseExpressions();`
`118`	`118`	`return outputProjectList;`
`119`	`119`	`}`
`120`	`120`