HyukjinKwon commented on code in PR #58403:
URL: https://github.com/apache/spark/pull/58403#discussion_r3974084968
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -2484,6 +2498,125 @@ object PushPredicateThroughNonJoin extends
Rule[LogicalPlan] with PredicateHelpe
}
}
+ /**
+ * Splits `project` into a stack of [[Project]]s with the expensive
conditions in
+ * `expensiveConds` between them, so an expensive expression only runs on
the rows the filters
+ * below it kept (SPARK-55014). For example, with both regexes considered
expensive and a child
+ * producing `a` and `e`:
+ *
+ * {{{
+ * Filter f AND g
+ * Project a, rlike(e, 'magic') AS f, rlike(e, 'other') AS g
+ * child
+ * }}}
+ * becomes
+ * {{{
+ * Filter g
+ * Project a, f, rlike(e, 'other') AS g
+ * Filter f
+ * Project a, e, rlike(e, 'magic') AS f
+ * child
+ * }}}
+ * so `rlike(e, 'other')` only runs on the rows where `rlike(e, 'magic')`
was true.
+ *
+ * Conditions are grouped by the aliases they reference, then the group
needing the fewest
+ * not-yet-computed aliases is split off first -- least demanding first
keeps the expensive
+ * expressions as high in the stack, and so over as few rows, as we can
manage. On a tie, the
+ * group making the most conditions evaluable wins, falling back to the
condition written first
+ * (the only selectivity signal we have). Whole-stage codegen already defers
a projected
+ * expression past a filter that does not need it
(`CodegenSupport.evaluateRequiredVariables`),
+ * so this rule only buys the same saving on the paths codegen does not
cover -- interpreted
+ * projections, operators it bails out of, Python UDFs.
+ *
+ * Two restrictions keep the split from costing more than it saves:
+ * - Only split a [[Project]] off while it leaves an expensive alias for a
later layer; an extra
+ * operator has to buy a real deferral.
+ * - Aliases sharing an expensive sub-expression move together.
Subexpression elimination
+ * works within one projection and cannot reach across a [[Filter]], so
splitting them apart
+ * would re-evaluate the shared part on every row below the filter and
again on every
+ * survivor (a struct-returning UDF read field by field is the common
shape).
+ *
+ * The most demanding group is left un-placed for the caller to put above
the projection --
+ * where expensive conditions go when there is nothing worth splitting.
+ *
+ * Returns the new plan, the alias attributes it has already computed, and
the un-placed
+ * conditions (in their original, alias-referencing form).
+ */
+ private def splitProjectForExpensiveConditions(
+ project: Project,
+ aliasMap: AttributeMap[Alias],
+ baseChild: LogicalPlan,
+ expensiveConds: Seq[(Expression, AttributeMap[Alias])])
+ : (LogicalPlan, AttributeSet, Seq[Expression]) = {
+ val expensiveAliases = AttributeSet(
+ aliasMap.collect { case (attr, alias) if alias.child.expensive => attr })
+ // Expensive sub-expressions behind each expensive alias, to spot when two
aliases are the
+ // same piece of work. Cheap aliases have none, so we skip the (possibly
wide) cheap rest.
+ val expensivePartsOf = AttributeMap(expensiveAliases.toSeq.map { attr =>
+ attr -> aliasMap(attr).child.collect { case e if e.expensive =>
e.canonicalized }.toSet
+ })
+ def expensiveParts(attr: Attribute): Set[Expression] =
+ expensivePartsOf.getOrElse(attr, Set.empty)
+ val aliasesSharing = expensivePartsOf.toSeq
+ .flatMap { case (attr, parts) => parts.map(part => part -> attr) }
+ .groupBy(_._1)
+ .map { case (part, pairs) => part -> AttributeSet(pairs.map(_._2)) }
+ // Grow a set of aliases into the whole unit of shared expensive work it
belongs to,
+ // transitively (see the docstring restriction on shared sub-expressions).
+ def unitOf(aliases: AttributeSet): AttributeSet = {
+ var unit = aliases
+ var grew = true
+ while (grew) {
+ val next = unit.flatMap(expensiveParts).foldLeft(unit) {
+ case (acc, part) => acc ++ aliasesSharing(part)
+ }
+ grew = next.size > unit.size
+ unit = next
+ }
+ unit
+ }
+ val indexed = expensiveConds.zipWithIndex.map {
+ case ((cond, used), idx) => (unitOf(AttributeSet(used.keys)), idx, cond)
+ }
+ // Group conditions over the same aliases; group on sorted exprIds (not
AttributeSet) for
+ // stable hashing, and order groups by first condition to keep plans
stable.
+ var pending: Seq[(AttributeSet, Seq[(Int, Expression)])] = indexed
+ .groupBy { case (used, _, _) => used.toSeq.map(_.exprId.id).sorted }
+ .toSeq
+ .map { case (_, group) => (group.head._1, group.map { case (_, idx,
cond) => (idx, cond) }) }
+ .sortBy { case (_, conds) => conds.head._1 }
+ def conditionsOf(groups: Seq[(AttributeSet, Seq[(Int, Expression)])]):
Seq[Expression] =
+ groups.flatMap(_._2).sortBy(_._1).map(_._2)
+
+ // Projection order is stable, unlike the alias map's iteration order.
+ val orderedAliases = project.projectList.collect { case a: Alias => a }
+
+ var plan = baseChild
+ var computed = AttributeSet.empty
+ var searching = true
+ while (searching) {
+ // Only split while it leaves an expensive alias for a later layer.
+ val stillToCompute = expensiveAliases -- computed
+ val candidates = pending.filter { case (used, _) =>
!stillToCompute.subsetOf(used) }
+ if (candidates.isEmpty) {
+ searching = false
+ } else {
+ val (bestUsed, _) = candidates.minBy {
+ case (used, conds) => ((used -- computed).size, -conds.size)
+ }
+ val newAliases = orderedAliases.filter(a => bestUsed.contains(a) &&
!computed.contains(a))
+ computed ++= AttributeSet(newAliases.map(_.toAttribute))
+ // Keep the plan's existing outputs available for later projections;
column pruning drops
+ // the extras.
+ val newProject = project.copy(projectList = plan.output ++ newAliases,
child = plan)
+ val (evaluable, rest) = pending.partition { case (used, _) =>
used.subsetOf(computed) }
+ plan = Filter(conditionsOf(evaluable).reduce(And), newProject)
Review Comment:
This reorders the expensive conjuncts across layers by the cost heuristic,
but only the projection fields are guarded for determinism
(`fields.forall(_.deterministic)` on the match), not the conditions. A
non-deterministic conjunct that references an expensive alias (say
`nondetUDF(f)`) would then be evaluated at whichever layer the heuristic picks
-- i.e. on a different row set than the single top `Filter` it landed in before
this PR. Is a non-deterministic predicate over an expensive projected alias out
of scope here? If not, keeping non-deterministic conjuncts in the top `Filter`
(excluding them from the split) would avoid reordering them.
--
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]
