fix fallback due to project

apache · May 28, 2024 · 07c1af3 · 07c1af3
1 parent b2e3de8
commit 07c1af3
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/...-core/src/main/scala/org/apache/gluten/extension/columnar/rewrite/PullOutPreProject.scala b/...-core/src/main/scala/org/apache/gluten/extension/columnar/rewrite/PullOutPreProject.scala
@@ -158,8 +158,12 @@ object PullOutPreProject extends RewriteSingleNode with PullOutProjectHelper {
       val newAgg = copyBaseAggregateExec(agg)(
         newGroupingExpressions = newGroupingExpressions,
         newAggregateExpressions = newAggregateExpressions)
+      // ISSUE-5852: literals with same names are lost in expressionMap, need addback to Project
+      val missingLiterals =
+          getMissingLiterals(agg.groupingExpressions)
+      val expressions = expressionMap.values.toSeq ++ missingLiterals
       val preProject = ProjectExec(
-        eliminateProjectList(agg.child.outputSet, expressionMap.values.toSeq),
+        eliminateProjectList(agg.child.outputSet, expressions),
         agg.child)
       newAgg.withNewChildren(Seq(preProject))
 

diff --git a/gluten-core/src/main/scala/org/apache/gluten/utils/PullOutProjectHelper.scala b/gluten-core/src/main/scala/org/apache/gluten/utils/PullOutProjectHelper.scala
@@ -63,6 +63,7 @@ trait PullOutProjectHelper {
       case alias: Alias =>
         alias.child match {
           case _: Literal =>
+            projectExprsMap.getOrElseUpdate(alias.child.canonicalized, alias)
             alias.toAttribute
           case _ =>
             projectExprsMap.getOrElseUpdate(alias.child.canonicalized, alias).toAttribute
@@ -75,6 +76,19 @@ trait PullOutProjectHelper {
           .toAttribute
     }
 
+  protected def getMissingLiterals(exprs: Seq[NamedExpression]): Seq[NamedExpression] = {
+    val literalMap = new mutable.HashMap[Expression, NamedExpression]()
+    exprs.filter({
+      case alias: Alias =>
+        alias.child match {
+          case _: Literal =>
+            literalMap.put(alias.child.canonicalized, alias).nonEmpty
+          case _ => false
+        }
+      case _ => false
+    })
+  }
+
   /**
    * Append the pulled-out NamedExpressions after the child output and eliminate the duplicated
    * parts in the append.