HyukjinKwon commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4254326119
##########
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])
Review Comment:
Thanks, single nested updates now render with the legacy names, and the test
checks that different parents get distinct names. Chained nested calls still
get different names than before. For example, withField("a.x",
...).withField("a.y", ...) or dropFields("a.b", "a.c") used to embed the first
update in the name, and now render as one merged update on s.a. I raised this
as a question in this review: keep the old names or document the new ones.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4241668318","thread_id":"inline:4241668318","verdict_sha256":"376a843e0801b33f093c16ee386ce4fd6a12672aa472f4cfb03d6ca9bc61f79e"}
-->
##########
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:
Thanks, verified. The batch is Once again and the comment is gone. The new
suite checks that nested updates lower fully in one pass and that a second
application changes nothing.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4241668320","thread_id":"inline:4241668320","verdict_sha256":"376a843e0801b33f093c16ee386ce4fd6a12672aa472f4cfb03d6ca9bc61f79e"}
-->
##########
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:
Thanks, verified. Both comments now describe the bound-reference rule that
the new With tests assert.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4241668321","thread_id":"inline:4241668321","verdict_sha256":"376a843e0801b33f093c16ee386ce4fd6a12672aa472f4cfb03d6ca9bc61f79e"}
-->
##########
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:
Thanks, verified. dataType now walks the indexed ordinals and updated values
with iterators, and doGenCode uses indexedOrdinals. The 100- and 500-field
tests cover types, nullability, and interpreted and codegen results.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4241668324","thread_id":"inline:4241668324","verdict_sha256":"376a843e0801b33f093c16ee386ce4fd6a12672aa472f4cfb03d6ca9bc61f79e"}
-->
--
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]
