uros-b commented on code in PR #58119:
URL: https://github.com/apache/spark/pull/58119#discussion_r3872880106
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -66,6 +66,105 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
}
}
+ /**
+ * Rewrites the `With` expressions in a single expression tree by inlining
their common
+ * expressions. Uses `transformUp` to handle nested `With`.
+ *
+ * Inlining duplicates a definition at every reference, so a definition
referenced more than once
+ * must be safe to duplicate (see `isSafeToDuplicate`); anything else
throws, as there is no plan
+ * here to pre-evaluate it into. Run [[ComputeCurrentTime]],
[[ReplaceCurrentLike]] and
+ * [[SpecialDatetimeValues]] first to fold the current date/time and catalog
expressions to
+ * literals.
+ *
+ * Does not descend into subquery plans (e.g. `ScalarSubquery`). A caller
whose expression
+ * may contain a subquery must rewrite those plans separately.
+ */
+ private[sql] def applyForExpression(expression: Expression): Expression =
+ inlineWith(expression, checkDuplication = true)
+
+ // The plan-level rewrite shares this to inline `With` in conditional
branches, which may not be
+ // evaluated and so can't be pulled into a Project; it passes false to
inline unconditionally.
+ private def inlineWith(expression: Expression, checkDuplication: Boolean):
Expression = {
+ // Which definitions are safe to duplicate can only be decided
outer-first, so it is collected
+ // in a separate pass before inlining bottom-up.
+ val safeIds = if (checkDuplication) {
+ safeToDuplicateIds(expression)
+ } else {
+ Set.empty[CommonExpressionId]
+ }
+ expression.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case With(child, defs) =>
+ if (checkDuplication) {
+ // Nested `With` is already rewritten, so a ref left in `child`
belongs to `defs` or to
+ // an enclosing `With`. Only the ids defined here are checked.
+ val refCounts = child.collect { case ref: CommonExpressionRef =>
ref.id }
+ .groupBy(identity)
+ .transform((_, refs) => refs.size)
+ defs.foreach { commonExprDef =>
+ if (refCounts.getOrElse(commonExprDef.id, 0) > 1 &&
+ !safeIds.contains(commonExprDef.id)) {
+ throw SparkException.internalError(
+ "Cannot inline a common expression definition that is
referenced more than " +
+ s"once and is not safe to duplicate:
${commonExprDef.child.sql}")
+ }
+ }
+ }
+ val refToExpr = defs.map(commonExprDef => commonExprDef.id ->
commonExprDef.child).toMap
+ child.transformWithPruning(_.containsPattern(COMMON_EXPR_REF)) {
+ case ref: CommonExpressionRef if refToExpr.contains(ref.id) =>
refToExpr(ref.id)
+ }
+ }
+ }
+
+ /**
+ * Ids of the definitions that are safe to duplicate. Walks outer-first, so
a definition that
+ * references an enclosing one is decided after it, and visits a
definition's own nested `With`
+ * before the definition itself. Ids are globally unique, so one flat set is
enough.
+ */
+ private def safeToDuplicateIds(expression: Expression):
Set[CommonExpressionId] = {
Review Comment:
safeIds assumes globally unique IDs, but sibling canonicalized With
expressions both start numbering at 1. A safe sibling can therefore authorize
duplication of an unsafe definition. Reject canonicalized IDs as the plan-level
path does, or scope safety per With.
--
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]
