LuciferYang commented on code in PR #59984:
URL: https://github.com/apache/spark/pull/59984#discussion_r4481087414
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -2333,10 +2344,98 @@ class CodegenContext extends Logging {
// We need to make sure that we do not reuse stateful expressions. This is
needed for codegen
// as well because some expressions may implement `CodegenFallback`.
val cleanedExpressions =
expressions.map(_.freshCopyIfContainsStatefulExpression())
- if (doSubexpressionElimination)
subexpressionElimination(cleanedExpressions)
+ if (doSubexpressionElimination) {
+ subexprFunctions ++=
subexpressionElimination(cleanedExpressions).map(_._2)
+ }
cleanedExpressions.map(e => e.genCode(this))
}
+ /**
+ * Returns the index of the first of `expressions` that references the
common subexpression
+ * `subexpr`, or -1. When the expressions are evaluated in order and one of
them can stop the
+ * evaluation, such as a FALSE conjunct of a predicate, the subexpression is
evaluated right
+ * before that expression instead of before all of them (SPARK-56032,
SPARK-60332). An
+ * occurrence in a conditionally evaluated part of an expression, such as a
branch of CASE WHEN,
+ * counts as a reference, so the subexpression is still evaluated before
such an expression.
+ */
+ def firstReferencingExpression(expressions: Seq[Expression], subexpr:
ExpressionEquals): Int = {
+ expressions.indexWhere(_.exists(e => ExpressionEquals(e) == subexpr))
+ }
+
+ /**
+ * Generates the code of the top-level conjuncts of `predicate` with
subexpression elimination,
+ * for a predicate whose conjuncts are evaluated in order and whose
evaluation stops at a FALSE
+ * conjunct, as `And.eval` does.
+ *
+ * The common subexpressions are found in the whole predicate, as
`generateExpressions` finds
+ * them, so that the same subexpressions are eliminated: none inside a task
for a predicate that
+ * holds a plan expression, and only those of the first conjunct with
+ * `spark.sql.subexpressionElimination.skipForShortcutExpr`. The predicate
is then split into
+ * its conjuncts, keeping an AND of the top-level chain that is itself a
common subexpression as
+ * one conjunct, which reads the state of the subexpression. The common
subexpressions that a
+ * conjunct contains are evaluated by functions called just before the first
conjunct that
+ * references them, instead of before the whole predicate (SPARK-60332); any
other AND of the
+ * chain is not, since the conjuncts are generated one by one. Returns, per
conjunct, the code
+ * of the calls to put before it and the code of the conjunct, or None,
generating nothing, when
+ * no subexpression is eliminated: the caller then generates the predicate
as a whole.
+ *
+ * The function of a subexpression that contains an eliminated subexpression
reads the state
+ * of the latter, so the calls are ordered by their dependencies: a
contained subexpression is
+ * called before the conjunct that first references the containing one, even
when its own first
+ * reference is later. A conjunct can contain a copy of the containing
subexpression that lacks
+ * the contained one, e.g. `(a + c) + b` for `(a + b) + c` under the
canonicalization of a
+ * commutative operator.
+ *
+ * This must be the only subexpression elimination on this context.
+ */
+ def generateConjunctsWithSubexprCodes(predicate: Expression):
Option[Seq[(String, ExprCode)]] = {
+ val cleanedPredicate = predicate.freshCopyIfContainsStatefulExpression()
+ equivalentExpressions.addExprTree(cleanedPredicate)
+ val allCommonExprs = equivalentExpressions.getCommonSubexpressions
+ val commonEqs = allCommonExprs.map(ExpressionEquals(_)).toSet
+ def split(e: Expression): Seq[Expression] = e match {
+ case And(left, right) if !commonEqs.contains(ExpressionEquals(e)) =>
+ split(left) ++ split(right)
+ case _ => Seq(e)
+ }
+ val conjuncts = split(cleanedPredicate)
+ // The common subexpressions, a contained one before the ones that contain
it, that some
+ // conjunct contains.
+ val commonExprs = allCommonExprs.filter { expr =>
+ firstReferencingExpression(conjuncts, ExpressionEquals(expr)) >= 0
+ }
+ if (commonExprs.isEmpty) {
+ None
+ } else {
+ val calls = commonExprs.map(expr => ExpressionEquals(expr) ->
generateSubexprFunction(expr))
+ val codes = conjuncts.map(e => e.genCode(this))
+ val callOf = calls.toMap
+ val eliminated = calls.map(_._1)
+ // The eliminated subexpressions that `expr` contains, other than itself.
+ def dependencies(expr: Expression): Seq[ExpressionEquals] =
eliminated.filter { dep =>
+ expr.children.exists(_.exists(e => ExpressionEquals(e) == dep))
+ }
+ val emitted = mutable.HashSet.empty[ExpressionEquals]
+ Some(codes.zipWithIndex.map { case (code, k) =>
+ val callsBefore = mutable.ArrayBuffer.empty[String]
+ def emit(subexpr: ExpressionEquals): Unit = {
+ if (!emitted.contains(subexpr)) {
+ dependencies(subexpr.e).foreach(emit)
+ emitted += subexpr
+ callsBefore += callOf(subexpr)
+ }
+ }
+ eliminated.filter(e => firstReferencingExpression(conjuncts, e) ==
k).foreach(emit)
Review Comment:
`generateConjunctsWithSubexprCodes` calls `firstReferencingExpression` for
every eliminated subexpression inside the loop over the conjuncts, so predicate
generation is O(n²·c) for n conjuncts and c common subexpressions. A filter
whose input has more than `spark.sql.codegen.maxFields` columns after pruning
is not in whole-stage codegen and generates its predicate in every task. With
400 columns, each filtered by two conjuncts sharing `c_i + 1`, one
`Predicate.create` takes 6.7 s instead of 79 ms on master.
The first reference does not depend on `k`, so computing it once before the
loop is enough, as the CSE branch of `FilterExec` does with `groupBy`. With
`val byFirstReference = eliminated.groupBy(e =>
firstReferencingExpression(conjuncts, e))` before the loop and
`byFirstReference.getOrElse(k, Nil).foreach(emit)` inside it, generation takes
93 ms and the new tests pass.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
