bhollis-dbx commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4250541052
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -301,7 +301,9 @@ abstract class Optimizer(catalogManager: CatalogManager)
RemoveNoopOperators),
// This batch must be executed after the `RewriteSubquery` batch, which
creates joins.
Batch("NormalizeFloatingNumbers", Once, NormalizeFloatingNumbers),
- Batch("ReplaceUpdateFieldsExpression", Once,
ReplaceUpdateFieldsExpression)))
+ // Repeat because an update can contain another update. Keep the resulting
With expressions so
Review Comment:
Restored `Once` and removed the incorrect comment in 049f52ef8. Tests verify
that nested updates fully lower in one pass and that applying the rule again
changes nothing.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala:
##########
@@ -546,12 +546,65 @@ class SubexpressionEliminationSuite extends SparkFunSuite
with ExpressionEvalHel
assert(equivalence1.getCommonSubexpressions.size == 1)
}
- test("SPARK-58902: no candidate below a With, with or without the
short-circuit peel") {
+ test("With exposes independent subexpressions but not unbound references") {
+ val input = AttributeReference("input", IntegerType)()
+ val shared = Add(input, Literal(1))
+ val expression = With(input) { case Seq(ref) =>
+ Add(ref, Multiply(shared, shared))
+ }
+ val equivalence = new EquivalentExpressions
+ equivalence.addExprTree(expression)
+ assert(equivalence.getCommonSubexpressions == Seq(shared))
+ assert(equivalence.getAllExprStates().forall { state =>
+ state.expr == expression ||
!state.expr.exists(_.isInstanceOf[CommonExpressionRef])
+ })
+ }
+
+ test("With candidates must bind references in their lexical scope") {
+ val input = AttributeReference("input", IntegerType)()
+ val shared = Add(input, Literal(1))
+ var inner: Expression = null
+ val outer = With(input) { case Seq(outerRef) =>
+ inner = With(input) { case Seq(innerRef) =>
+ Add(Add(outerRef, innerRef), shared)
+ }
+ Add(inner, shared)
+ }
+ val equivalence = new EquivalentExpressions
+ equivalence.addExprTree(outer)
+ assert(equivalence.getExprState(outer).nonEmpty)
+ assert(equivalence.getExprState(inner).isEmpty)
+ assert(equivalence.getCommonSubexpressions == Seq(shared))
+ assert(!equivalence.addExpr(inner))
+ val closed = With(input) { case Seq(ref) => Add(ref, Literal(1)) }
+ val sibling = new CommonExpressionRef(closed.defs.head)
+ assert(!equivalence.addExpr(Add(closed, sibling)))
+ }
+
+ test("With does not expose lazy definitions or interpreted children") {
+ val input = AttributeReference("input", IntegerType)()
+ val shared = Add(input, Literal(1))
+ val lazyDefinition = With(Multiply(shared, shared)) { case Seq(ref) =>
+ If(Literal(false), ref, input)
+ }
+ val fallback = With(input) { case Seq(ref) =>
+ Add(CodegenFallbackExpression(ref), Multiply(shared, shared))
+ }
+ val guarded = With(input) { case Seq(ref) =>
+ If(GreaterThan(ref, Literal(0)), Multiply(shared, shared), Literal(0))
+ }
+ Seq(lazyDefinition, fallback, guarded).foreach { expression =>
+ val equivalence = new EquivalentExpressions
+ equivalence.addExprTree(expression)
+ assert(equivalence.getCommonSubexpressions.isEmpty)
+ }
+ }
+
+ test("SPARK-58902: candidates below a With cannot hold unbound references") {
// A `CommonExpressionRef` can only be evaluated inside the `With` that
binds it: the codegen
// slots exist only while that `With` is being generated, and
`getCommonExpr` throws otherwise.
// Subexpression elimination generates its candidates outside every `With`
scope, so a candidate
- // holding a reference would fail the query. `childrenToRecurse` therefore
stops at a `With`,
- // including after `skipForShortcut` has peeled an `And`/`Or` down to one.
+ // holding an unbound reference would fail the query.
Review Comment:
Updated in 049f52ef8. The comments now describe the bound-reference rule:
independent expressions below `With` remain eligible; candidates with unbound
references are rejected.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:
##########
@@ -759,9 +737,212 @@ case class WithField(name: String, valExpr: Expression)
/**
* Drop a field by name.
*/
-case class DropField(name: String) extends StructFieldsOperation with
LeafLike[Expression] {
- override def apply(values: Seq[(StructField, Expression)]):
Seq[(StructField, Expression)] =
- values.filterNot { case (field, _) => resolver(field.name, name) }
+case class DropField(name: String) extends StructFieldsOperation with
LeafLike[Expression]
+
+/**
+ * Updates the current value of a field without duplicating the parent struct
expression.
+ */
+case class UpdateNestedField(name: String, fieldOps:
Seq[StructFieldsOperation])
+ extends StructFieldsOperation {
+
+ override def children: Seq[Expression] = fieldOps
+
+ override protected def withNewChildrenInternal(
+ newChildren: IndexedSeq[Expression]): Expression =
+ copy(fieldOps = newChildren.map(_.asInstanceOf[StructFieldsOperation]))
+}
+
+/**
+ * Builds a struct from fields copied from `structExpr` and values computed by
`updatedValues`.
+ *
+ * `outputOrdinals` describes each output position. A nonnegative entry copies
that field from
+ * `structExpr`; a negative entry evaluates the next expression in
`updatedValues`. Storing copied
+ * fields as ordinals avoids adding a `GetStructField` child for every copied
field. Row construction
+ * still visits every output field.
+ *
+ * For example, to transform `(a, b, c)` into `(a, newB, c, newD)`,
`outputOrdinals` is
+ * `Seq(0, -1, 2, -1)` and `updatedValues` is `Seq(newB, newD)`. The first and
third output fields
+ * copy source fields 0 and 2. The negative entries use `newB` and `newD`, in
that order.
+ *
+ * If a nullable `structExpr` evaluates to null, the result is null and
`updatedValues` are not
+ * evaluated.
+ */
+private case class UpdateFieldsExpression(
+ structExpr: Expression,
+ updatedValues: Seq[Expression],
+ outputOrdinals: Seq[Int],
+ outputFields: Seq[StructField]) extends ConditionalExpression {
+
+ override def children: Seq[Expression] = structExpr +: updatedValues
+
+ override def nullable: Boolean = structExpr.nullable
+
+ private lazy val indexedOrdinals = outputOrdinals.toArray
+ private lazy val indexedUpdatedValues = updatedValues.toArray
+ private lazy val needsSource = nullable || indexedOrdinals.exists(_ >= 0)
+
+ override def alwaysEvaluatedInputs: Seq[Expression] =
+ (if (needsSource) Seq(structExpr) else Nil) ++ (if (nullable) Nil else
updatedValues)
+
+ override def branchGroups: Seq[Seq[Expression]] = Nil
+
+ override def withNewAlwaysEvaluatedInputs(
+ inputs: Seq[Expression]): ConditionalExpression =
+ copy(
+ structExpr = if (needsSource) inputs.head else structExpr,
+ updatedValues = if (nullable) updatedValues else if (needsSource)
inputs.tail else inputs)
+
+ override lazy val dataType: StructType = {
Review Comment:
Fixed in 049f52ef8. Type construction uses indexed ordinals and iterators;
codegen uses indexed ordinals. Added 100/500-field checks for output types,
nullability, null inputs, and interpreted/codegen results.
--
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]
