Should we modify the logical plan to insert pre/post-project?

[liujiayi771]

In #4245, we hope to insert pre/post-project through spark Rule, which could simplify the handling of pre/post-project in operators such as agg, sort, join, window, etc.

Recently, I've been attempting to add the Rule for logical plan, and there are numerous advantages to this approach:

• The logical plan is more concise compared to the physical plan. For example, in agg, there's no need to consider whether it's partial or final.
• The necessary projects for gluten have been added right from the source. In most cases, gluten doesn't need to be aware of the insertion of the project, so the code will become more concise.
• Validation will not require additional modifications.
• It supports individual fallback for pre/post-project without the need to fallback on time-consuming computations like aggregate or sort. For example, in aggExprs, if there is a func that gluten does not support, we may only need a pre-project fallback. The agg can then use the result computed by the pre-project, eliminating the need for the agg to fallback.

During the process of fixing UT, I find that the modifications to the logical plan can indeed cause many side effects. While these side effects usually do not pose correctness issues, they can lead to failures when gluten executes Spark UTs.

For example,

• There are many UTs that count the number of plan nodes, such as the number of ReusedSubqueryExec. Before the modification, some ReusedSubqueryExec appear inside an agg. After being converted to a physical plan, both the partial and final agg will have ReusedSubqueryExec. However, if we move the subqueries from agg to a pre-project, the count will not double. As a result, the number of ReusedSubqueryExec will not match the number asserted in the Spark UTs.
• Some UTs verify the consistency of the logical plan and will re-execute some optimizer rules on the logical plan. For example, in the InjectRuntimeFilterSuite, before comparison, the ColumnPruning rule is executed. Our newly added InsertPreProject rule's output plan would be re-optimized by the ColumnPruning rule, which changes the plan and leads to a validation failure.
  https://github.com/apache/spark/blob/018808236708bea7a78618abf750bea39be3c9f8/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala#L278-L290

Currently, in the PR, we have only dealt with agg and sort. I believe that if we were to add support for join and window, many more issues would come to light.

I've been considering whether gluten, as a columnar plan transformer, should be modifying the logical plan. Should we perhaps limit ourselves to operating at the columnar plan level? Alternatively, should we spend time investigating the causes of these unit test failures? If the comparison is against the vanilla plan, then we should consider excluding or overriding the new tests.

Alternatively, we could consider using columnar plan Rule to modify the plan, but this might result in losing some of the advantages mentioned above.

• Fallback will be the same as the current situation; if there are unsupported expressions in pre/post-project, they will fallback along with agg and sort. However, I believe that optimizations for individual fallbacks can also be further supported; it's just that a proper solution has not yet been conceived.

But it can also bring some benefits.

• The code for the columnar transformer is entirely under our control, so to address certain issues, we can add the information we need within it.
• It won't affect Spark's original plan; some Spark unit tests can now pass normally.
• The processing of the bloom filter will become simpler, as it requires complex matching in the logical plan to handle it.

[liujiayi771]

cc @ulysses-you @rui-mo @zhouyuan

[liujiayi771]

Another bad case in UT, however, modifying the physical plan should probably lead to the same issues as well. Spark UT checks whether the plan's outputOrdering meets the expectations.

The sortOrder in a sort operation might contain expressions, which would require adding a pre-project to execute them beforehand. Consider the following plan fragment:

Sort [func(data#9) DESC NULLS FIRST, func(id#8) DESC NULLS FIRST], false
 +- Project [key#3 AS id#8, value#4 AS data#9]

When calculating the outputOrdering for this plan, we directly take the sort order from the sort operation, get [func(data#9) DESC NULLS FIRST, func(id#8) DESC NULLS FIRST]. Now, to support native execution of this plan, we need to insert a pre-project so that the expressions in sort.order can be executed beforehand. The resulting plan will look like this:

Sort [pre_1#17 DESC NULLS FIRST, pre_2#18 DESC NULLS FIRST], false
 +- Project [func(data#9) AS pre_1#17, func(id#8) AS pre_2#18]
   +- Project [key#3 AS id#8, value#4 AS data#9]

However, the output of the sort in the modified plan will not match the original plan's output. The original plan's sort output might be identical to the child's output, for example, [id#8, data#9], but now it becomes something like [pre_1#17, pre_2#18], which represents the computed values of functions in the order.

To ensure that the output also matches the expected, you would need to include the original input columns [id#8, data#9] in the pre-project and add a post-project to extract these columns, discarding the additional computed function results. The final plan could look something like this:

Project [id#8, data#9]
+- Sort [pre_1#17 DESC NULLS FIRST, pre_2#18 DESC NULLS FIRST], false
  +- Project [id#8, data#9, func(data#9) AS pre_1#17, func(id#8) AS pre_2#18]
    +- Project [key#3 AS id#8, value#4 AS data#9]

This seems to be fine, but we encountered a consistency issue when determining the outputOrdering of the current plan segment. Let's first take a look at how Spark determines the outputOrdering of a project. The specific logic is in AliasAwareQueryOutputOrdering.outputOrdering.

https://github.com/apache/spark/blob/v3.4.1/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/AliasAwareOutputExpression.scala#L121-L129

val outputSet = AttributeSet(outputExpressions.map(_.toAttribute))
orderingExpressions.iterator.map { order =>
  val validChildren = order.children.filter(_.references.subsetOf(outputSet))
  if (validChildren.nonEmpty) {
    Some(order.copy(child = validChildren.head, sameOrderExpressions = validChildren.tail))
  } else {
    None
  }
}

The outputSet of the most downstream project is [id#8, data#9], and the orderingExpressions are child.orderingExpressions = [pre_1#17 DESC NULLS FIRST, pre_2#18 DESC NULLS FIRST]. Each order in orderingExpressions has references pre_1 and pre_2, which do not match id#8 and data#9. This means that the calculation result of order.children.filter(_.references.subsetOf(outputSet)) in the above code snippet will be None, failing to find validChildren, ultimately causing the project's outputOrdering to be None, inconsistent with the original plan.

Let's explain this using id#8 as an example. The fundamental issue here is that after order.child was modified to Attribute, its reference was changed to pre_2#18. Although pre_2#18 is computed from id#8, in Spark, an attribute is the origin of a reference, and once it gets to pre_2#18, it can no longer be traced back to its source. Originally, order.child was func(id#8), and tracing the reference of func(id#8) would correctly lead back to id#8. This can be simply illustrated with the following diagram.

[rui-mo]

@liujiayi771 Thanks for taking efforts to prototype on this optimization. I recognize this work can be divided into two parts. One is about code refactor, and the other is about performance optimization. As you desribed above, 1) to modify logical plan provides several advantages on the code simplicity, but also introduces possible disturbances and risks for Spark plan. 2) to modify logical plan brings extra performance improvement as it reduces fallback.
For 1), I'm wondering if it's necessary to bring suck risks just for advantages on the code implementation. For 2), After modifying the physical plan, seems the fallback of pre/post project could also be separated from the operator contains it by improving the validation logic. So seems 2) is not a must reason to tackle pre and post project on logical plan level.
Would it make sense to conduct this work through two steps? Firstly refactor current process of creating pre/post project on physical plan level. After that, conduct the optimization to make it possible to fallback pre/post project separately.
@liujiayi771 @ulysses-you Glad to hear your opinions, thanks!

[ulysses-you]

I think pull out pre-project at logical side sounds better to me. Besides the fallback, it's more easy and efficient.

Let's say a plan with aggregate and distinct. We can reduce two expression evaluation and project if we pull out project at logical side.

SELECT c1, count(distinct (c2 + 1)) FROM t GROUP BY c1

-- If we pull out `c2 + 1` at logical side, then the final physical plan is:
Aggregate c1, count(__gluten_c1)
  Exchange c1
    Aggregate c1, partial_count(__gluten_c1)
      Aggregate c1, __gluten_c1
        Exchange c1, __gluten_c1
          Aggregate c1, __gluten_c1
            Project c1, (c2 + 1) as __gluten_c1


-- if we pull out `c2 + 1` at columnar side, then the final physical plan is:
Aggregate c1, count(__gluten_c1)
  Exchange c1
    Aggregate c1, partial_count(__gluten_c1)
      Project c1, (c2 + 1) as __gluten_c3
        Aggregate c1, __gluten_c1
          Exchange c1, __gluten_c2
            Project c1, __gluten_c1, (c2 + 1) as __gluten_c2
              Aggregate c1, __gluten_c1
                Project c1, (c2 + 1) as __gluten_c1

Let's say a plan with shuffled hash join. We can also reduce one expression evaluation and project if we pull out project at logical side. BTW, we can reduce two expression evaluation if the join is SMJ, one extra project for sort.

SELECT * FROM t1 JOIN t2 ON c1 + 1 = c2

-- If we pull out `c1 + 1` at logical side, then the final physical plan is:
ShuffledHashJoin c1, __gluten_c1
  Exchange c1
    t1
  Exchange __gluten_c1
    Project (c2 + 1) as __gluten_c1
      t2

-- if we pull out `c1 + 1` at columnar side, then the final physical plan is:
ShuffledHashJoin c1, __gluten_c1
  Project c1, (c2 + 1) as __gluten_c1
    Exchange c1
      t1
    Exchange __gluten_c1
      Project c2, (c2 + 1) as __gluten_c1
        t2

I also did not see vanilla Spark has significant plan changes at physical side. In general, the physcial rules are used to eliminate the operators. One more reference is that, there already exists an logical rule in vanilla Spark to pull out the complex expression for aggregate grouping keys PullOutGroupingExpressions. I think it's good to follow it.

[liujiayi771]

@ulysses-you I would also prefer to address it at the logical plan level, but during the process of fixing unit tests, I encountered many bad cases. The current blocking issue I'm encountering is the aforementioned outputOrdering. Do you have any ideas for a solution?

[liujiayi771]

The issue related to sort outputOrdering discussed above needs to be addressed in the physical plan (SortExecTransformer).

We will handle the parts that are suitable for logical plan processing in the logical plan rule, and those that are suitable for physical plan processing in the physical plan rule.

[liujiayi771]

Some Strategy in Spark will add function to Agg's groupingExpressions. For example, StatefulAggregationStrategy will apply NormalizeFloatingNumbers.normalize to groupingExpressions, which will add KnownFloatingPointNormalized expr to float column.

[liujiayi771]

We need to trim KnownFloatingPointNormalized in gluten.