HyukjinKwon commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4241668312
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:
##########
@@ -806,35 +987,131 @@ case class UpdateFields(structExpr: Expression,
fieldOps: Seq[StructFieldsOperat
override def prettyName: String = "update_fields"
- private lazy val newFieldExprs: Seq[(StructField, Expression)] = {
- def getFieldExpr(i: Int): Expression = structExpr match {
- case c: CreateNamedStruct => c.valExprs(i)
- case _ => GetStructField(structExpr, i)
- }
- val fieldsWithIndex =
structExpr.dataType.asInstanceOf[StructType].fields.zipWithIndex
- val existingFieldExprs: Seq[(StructField, Expression)] =
- fieldsWithIndex.map { case (field, i) => (field, getFieldExpr(i))
}.toImmutableArraySeq
- fieldOps.foldLeft(existingFieldExprs)((exprs, op) => op(exprs))
- }
+ private lazy val newFieldSources: Seq[(StructField,
UpdateFields.FieldSource)] =
+ UpdateFields.fieldSources(structExpr, fieldOps)
- private lazy val newFields: Seq[StructField] = newFieldExprs.map(_._1)
+ private lazy val newFields: Seq[StructField] = newFieldSources.map(_._1)
- lazy val newExprs: Seq[Expression] = newFieldExprs.map(_._2)
+ /** Returns the expression for one output field without constructing the
complete struct. */
+ def newExpr(ordinal: Int): Expression = newFieldSources(ordinal)._2 match {
+ case UpdateFields.Original(sourceOrdinal) => getFieldExpr(sourceOrdinal)
+ case UpdateFields.Updated(expr) => expr
+ }
- lazy val evalExpr: Expression = {
- val createNamedStructExpr = CreateNamedStruct(newFieldExprs.flatMap {
- case (field, expr) => Seq(Literal(field.name), expr)
- })
+ private def getFieldExpr(ordinal: Int): Expression = structExpr match {
+ case c: CreateNamedStruct => c.valExprs(ordinal)
+ case _ => GetStructField(structExpr, ordinal)
+ }
- if (structExpr.nullable) {
- If(IsNull(structExpr), Literal(null, dataType), createNamedStructExpr)
+ lazy val evalExpr: Expression = {
+ def buildExpression(structRef: Expression): Expression = {
+ val source = structExpr match {
+ case namedStruct: CreateNamedStruct => namedStruct
+ case _ => structRef
+ }
+ val fieldSources = UpdateFields.fieldSources(source, fieldOps)
+ val updatedValues = ArrayBuffer.empty[Expression]
+ val outputOrdinals = fieldSources.map { case (_, source) =>
+ source match {
+ case UpdateFields.Original(ordinal) => structExpr match {
+ case namedStruct: CreateNamedStruct =>
+ updatedValues += namedStruct.valExprs(ordinal)
+ -1
+ case _ => ordinal
+ }
+ case UpdateFields.Updated(expr) =>
+ updatedValues += expr
+ -1
+ }
+ }
+ UpdateFieldsExpression(structRef, updatedValues.toSeq, outputOrdinals,
newFields)
Review Comment:
**Non-blocking (P2):** For a `CreateNamedStruct` input, every output value
already comes from `namedStruct.valExprs` (all `outputOrdinals` are -1), but
the original struct is still passed here as `structExpr`, and it is also kept
as the `With` definition on line 1030 for nested paths. `needsSource` is false,
so it is never evaluated per row, but it is still a child, together with the
values that `dropFields` or a later `withField` removed.
Rules that walk the whole expression tree after
`ReplaceUpdateFieldsExpression` still see those values. For example, with ANSI
on:
```scala
Seq((1L, 0L)).toDF("x", "z")
.agg(struct(sum($"x").as("a"), sum($"x" / $"z").as("b")).dropFields("b"))
```
`PhysicalAggregation` collects `AggregateExpression`s from the whole result
expression (`planning/patterns.scala:301-307`), so it still computes `sum(x /
z)` and fails with DIVIDE_BY_ZERO. Before this PR the same query returned
`Row(Row(1))`. Python UDFs in a dropped field are extracted and run the same
way, because `ExtractPythonUDFs` runs after the optimizer batches, and a scalar
subquery in a dropped field is still planned. Inputs reach this path through
direct `struct(...).dropFields/withField` calls and also when `CollapseProject`
inlines a projected `struct(...)` into a later update.
Row-level evaluation in a plain projection is not affected, which is why
`dropFields should not evaluate omitted CreateNamedStruct fields` passes. For
nested paths over `struct(...)` inputs, the retained `With` definition also
keeps a full copy of the input struct at every nesting level, so the lowered
tree grows quadratically with depth there.
**Verification:**
- **Regression:** An aggregate query that drops a struct(...) field
computing a failing ANSI expression returns the surviving fields instead of
throwing.
- **Regression:** A dropped struct(...) field holding a side-effecting or
failing value is not evaluated when the update is consumed by an aggregate or a
projection.
- **Inspection:** The lowered expression for a nested-path update over a
struct(...) input contains no With and holds each surviving value once.
- **Compatibility:** Existing withField/dropFields results and nullability
over struct(...) inputs are unchanged.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala:
##########
@@ -32,9 +32,9 @@ object SimplifyExtractValueOps extends Rule[LogicalPlan] {
// Remove redundant field extraction.
case GetStructField(createNamedStruct: CreateNamedStruct, ordinal, _) =>
createNamedStruct.valExprs(ordinal)
- case GetStructField(u: UpdateFields, ordinal, _)if
!u.structExpr.isInstanceOf[UpdateFields] =>
+ case GetStructField(u: UpdateFields, ordinal, _) =>
Review Comment:
**Non-blocking (P2):** Before this PR, the guard `if
!u.structExpr.isInstanceOf[UpdateFields]` left an extraction from an update
chain alone until `OptimizeUpdateFields` had merged the chain.
`OptimizeUpdateFields` runs before `SimplifyExtractValueOps` in the batch, so
without the guard this arm can see a chain that is not merged yet and rewrite
it to `If(IsNull(innerUpdate), null, value)` whenever the inner update is
nullable. `NullPropagation` only folds `IsNull` over non-nullable children, so
after lowering `IsNull` evaluates the whole inner update, including values the
extracted field does not depend on.
For example, with a non-nullable struct `s` whose field `a` is a nullable
struct, ANSI on, and `z = 0`:
```scala
col("s").withField("a.x", lit(1) / col("z")).getField("a")
.withField("y", lit(2)).getField("y")
```
In one pass the inner extraction becomes `UpdateFields(s.a, [x = 1 / z])`
and the outer one becomes `If(IsNull(UpdateFields(s.a, [x = 1 / z])), null,
2)`, which throws DIVIDE_BY_ZERO. With the guard, the next iteration merged the
chain into `UpdateFields(s.a, [x, y])` and simplified the read to
`If(IsNull(s.a), null, 2)`, which returns 2. The same shape appears when
`CollapseProject` inlines an earlier `withColumn` that updated `s.a`.
Nullable top-level inputs already behaved this way before the PR;
non-nullable ones are newly affected. Keeping the guard costs one more
optimizer iteration, which the merge target already paid.
**Verification:**
- **Regression:** An extraction from a re-updated, extracted nested struct
does not evaluate the discarded inner update values, for non-nullable and
nullable inputs.
- **Compatibility:** Existing SimplifyExtractValueOps and UpdateFields
optimizer results for single-level updates are unchanged.
##########
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:
**Nit (P3):** `outputFields(index)` and `updatedValues(updatedIndex)` are
positional lookups, and the sequences `evalExpr` passes in are `List`s on Scala
2.13 (`updatedValues.toSeq` from an `ArrayBuffer`, and the `fieldSources`
result after any `WithField`). Each lookup is O(index), so building this type
is O(F^2) for an F-field struct, and because it is a per-instance lazy val it
is recomputed for every copy a tree transform makes. `doGenCode` has the same
pattern in `outputOrdinals.indices.filter(outputOrdinals(_) < 0)` (line 891).
The per-row `eval` loop already uses `indexedOrdinals` and
`indexedUpdatedValues` after the earlier review. Using those arrays here too,
or walking the sequences with iterators, would keep planning and codegen linear
in width for the wide structs this PR targets.
**Verification:**
- **Inspection:** dataType and doGenCode access outputs through arrays or a
single iterator pass.
- **Behavior:** Wide-struct updates (hundreds of fields) keep identical
output types and results under interpreted and codegen evaluation.
##########
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:
**Nit (P3):** This lead comment now describes the bound-reference rule, but
two later comments in the same test still describe the removed one: around line
629 the guard is said to drop the `With` "along with its subtree", and around
line 637 the comment says that below a `With` "anything below it is not" a
legal candidate.
`childrenToRecurse` now returns `Seq(withExpr.child)` unless
`refUnderCodegenFallback` is set, and `supportedExpression` only rejects
candidates with unbound references. The new `With exposes independent
subexpressions but not unbound references` test asserts that a reference-free
subexpression below a `With` is a candidate, which contradicts both comments.
Updating them to the bound-reference rule would keep this test's explanation
accurate.
**Verification:**
- **Inspection:** The test comments describe the bound-reference rule
enforced by supportedExpression and asserted by the new With tests.
##########
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:
**Nit (P3):** `ReplaceUpdateFieldsExpression` is
`transformAllExpressionsWithPruning(...) { case u: UpdateFields => u.evalExpr
}` (`UpdateFields.scala:99-103`), a top-down transform that also recurses into
the children of each replacement. Every `UpdateFields` that `evalExpr` creates
(the nested updates built by `fieldSources`, or an input kept as a `With`
definition) is a child of the returned `UpdateFieldsExpression` (`structExpr +:
updatedValues`) or `With` (`child +: defs`), so the same pass lowers it. The
second fixed-point iteration only confirms that nothing changed, so the comment
describes a requirement the rule does not have.
Moving off `Once` also takes this batch out of the test-mode idempotence
check that `RuleExecutor` runs on `Once` batches, and adds one more traversal
of the plan.
**Verification:**
- **Regression:** Nested-path and chained withField/dropFields queries
optimize with no remaining UpdateFields, and the test-mode Once idempotence
check passes.
##########
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:
**Nit (P3):** `UpdateNestedField` overrides neither `prettyName` nor `sql`,
so `Expression.sql` renders it as `updatenestedfield(<child ops>)`: the
lowercased class name, without the parent field name. Unaliased nested-path
`withField`/`dropFields` columns get their name from this rendering.
`df.select(col("s").withField("a.b", lit(1)))` used to produce a column named
`update_fields(s, WithField(update_fields(s.a, WithField(1))))` and now
produces `update_fields(s, updatenestedfield(WithField(1)))`;
`dropFields("a.b")` changes the same way to `update_fields(s,
updatenestedfield(dropfield()))`.
This exposes an internal class name in the output schema, in classic and
Connect, and updates on different parents now collide: `withField("a.x",
lit(1))` and `withField("c.x", lit(1))` get the same name. Analyzed-plan
EXPLAIN output and `sqlExpr` error parameters that render the outer update
change the same way. Query results are not affected, and no test pins the
generated names.
**Verification:**
- **Compatibility:** Generated column names of unaliased nested-path
withField and dropFields match the merge-target names.
- **Regression:** Two nested updates on different parent fields produce
distinct generated names.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:
##########
@@ -806,35 +969,129 @@ case class UpdateFields(structExpr: Expression,
fieldOps: Seq[StructFieldsOperat
override def prettyName: String = "update_fields"
- private lazy val newFieldExprs: Seq[(StructField, Expression)] = {
- def getFieldExpr(i: Int): Expression = structExpr match {
- case c: CreateNamedStruct => c.valExprs(i)
- case _ => GetStructField(structExpr, i)
- }
- val fieldsWithIndex =
structExpr.dataType.asInstanceOf[StructType].fields.zipWithIndex
- val existingFieldExprs: Seq[(StructField, Expression)] =
- fieldsWithIndex.map { case (field, i) => (field, getFieldExpr(i))
}.toImmutableArraySeq
- fieldOps.foldLeft(existingFieldExprs)((exprs, op) => op(exprs))
- }
+ private lazy val newFieldSources: Seq[(StructField,
UpdateFields.FieldSource)] =
+ UpdateFields.fieldSources(structExpr.dataType.asInstanceOf[StructType],
fieldOps)
- private lazy val newFields: Seq[StructField] = newFieldExprs.map(_._1)
+ private lazy val newFields: Seq[StructField] = newFieldSources.map(_._1)
- lazy val newExprs: Seq[Expression] = newFieldExprs.map(_._2)
+ /** Returns the expression for one output field without constructing the
complete struct. */
+ def newExpr(ordinal: Int): Expression = newFieldSources(ordinal)._2 match {
+ case UpdateFields.Original(sourceOrdinal) => getFieldExpr(sourceOrdinal)
+ case UpdateFields.Updated(expr) => expr
+ }
- lazy val evalExpr: Expression = {
- val createNamedStructExpr = CreateNamedStruct(newFieldExprs.flatMap {
- case (field, expr) => Seq(Literal(field.name), expr)
- })
+ /**
+ * Returns an input field expression and marks generated reads so they can
share `structExpr`
+ * without changing equivalent expressions supplied to `withField`.
+ */
+ private def getFieldExpr(ordinal: Int): Expression = structExpr match {
+ case c: CreateNamedStruct => c.valExprs(ordinal)
+ case _ => UpdateFields.markGeneratedStructRead(
+ GetStructField(structExpr, ordinal), generatedStructReadOwner)
+ }
- if (structExpr.nullable) {
- If(IsNull(structExpr), Literal(null, dataType), createNamedStructExpr)
- } else {
- createNamedStructExpr
+ lazy val evalExpr: Expression = With(structExpr) { case Seq(structRef) =>
+ // Read fields used to rebuild the input struct from structRef so
structExpr is evaluated once.
+ // Do not rewrite expressions passed to withField; they must be evaluated
independently.
+ def replaceStruct(expr: Expression): Expression = expr.transformDown {
+ case field @ GetStructField(child, _, _)
+ if UpdateFields.isGeneratedStructRead(field,
generatedStructReadOwner) &&
+ child.eq(structExpr) =>
Review Comment:
Thanks, verified on the current head. UpdateNestedField applies nested
operations to the current field value, so the parent input is no longer
duplicated, and the With in evalExpr evaluates a non-literal input once. The
new unresolved-column and nondeterministic-input tests cover the path I raised.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4235965395","thread_id":"inline:4235965395","verdict_sha256":"ab2b0fd9bcfc655d854e51bb3e2b7b7c0e60761ac6cbf11e41a5d24bd9a888fd"}
-->
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:
##########
@@ -806,35 +969,129 @@ case class UpdateFields(structExpr: Expression,
fieldOps: Seq[StructFieldsOperat
override def prettyName: String = "update_fields"
- private lazy val newFieldExprs: Seq[(StructField, Expression)] = {
- def getFieldExpr(i: Int): Expression = structExpr match {
- case c: CreateNamedStruct => c.valExprs(i)
- case _ => GetStructField(structExpr, i)
- }
- val fieldsWithIndex =
structExpr.dataType.asInstanceOf[StructType].fields.zipWithIndex
- val existingFieldExprs: Seq[(StructField, Expression)] =
- fieldsWithIndex.map { case (field, i) => (field, getFieldExpr(i))
}.toImmutableArraySeq
- fieldOps.foldLeft(existingFieldExprs)((exprs, op) => op(exprs))
- }
+ private lazy val newFieldSources: Seq[(StructField,
UpdateFields.FieldSource)] =
+ UpdateFields.fieldSources(structExpr.dataType.asInstanceOf[StructType],
fieldOps)
- private lazy val newFields: Seq[StructField] = newFieldExprs.map(_._1)
+ private lazy val newFields: Seq[StructField] = newFieldSources.map(_._1)
- lazy val newExprs: Seq[Expression] = newFieldExprs.map(_._2)
+ /** Returns the expression for one output field without constructing the
complete struct. */
+ def newExpr(ordinal: Int): Expression = newFieldSources(ordinal)._2 match {
+ case UpdateFields.Original(sourceOrdinal) => getFieldExpr(sourceOrdinal)
+ case UpdateFields.Updated(expr) => expr
+ }
- lazy val evalExpr: Expression = {
- val createNamedStructExpr = CreateNamedStruct(newFieldExprs.flatMap {
- case (field, expr) => Seq(Literal(field.name), expr)
- })
+ /**
+ * Returns an input field expression and marks generated reads so they can
share `structExpr`
+ * without changing equivalent expressions supplied to `withField`.
+ */
+ private def getFieldExpr(ordinal: Int): Expression = structExpr match {
+ case c: CreateNamedStruct => c.valExprs(ordinal)
+ case _ => UpdateFields.markGeneratedStructRead(
+ GetStructField(structExpr, ordinal), generatedStructReadOwner)
+ }
- if (structExpr.nullable) {
- If(IsNull(structExpr), Literal(null, dataType), createNamedStructExpr)
- } else {
- createNamedStructExpr
+ lazy val evalExpr: Expression = With(structExpr) { case Seq(structRef) =>
+ // Read fields used to rebuild the input struct from structRef so
structExpr is evaluated once.
+ // Do not rewrite expressions passed to withField; they must be evaluated
independently.
+ def replaceStruct(expr: Expression): Expression = expr.transformDown {
+ case field @ GetStructField(child, _, _)
+ if UpdateFields.isGeneratedStructRead(field,
generatedStructReadOwner) &&
+ child.eq(structExpr) =>
+ field.copy(child = structRef)
}
+
+ val updatedValues = ArrayBuffer.empty[Expression]
+ val outputOrdinals = newFieldSources.map { case (_, source) =>
+ source match {
+ case UpdateFields.Original(ordinal) => structExpr match {
+ case namedStruct: CreateNamedStruct =>
+ updatedValues += namedStruct.valExprs(ordinal)
+ -1
+ case _ => ordinal
+ }
+ case UpdateFields.Updated(expr) =>
+ updatedValues += replaceStruct(expr)
+ -1
+ }
+ }
+ UpdateFieldsExpression(structRef, updatedValues.toSeq, outputOrdinals,
newFields)
}
}
object UpdateFields {
+ /** Describes where an output field gets its value without creating an
expression for a copy. */
+ private sealed trait FieldSource
+
+ /** Copies the field at `ordinal` from the input struct. */
+ private case class Original(ordinal: Int) extends FieldSource
+
+ /** Computes the field from an expression supplied to `withField`. */
+ private case class Updated(expr: Expression) extends FieldSource
+
+ /**
+ * Applies the requested field operations while retaining whether each
output field is copied or
+ * computed. This provenance lets evaluation copy unchanged fields by
ordinal and lets extraction
+ * return one field without constructing the complete struct.
+ */
+ private def fieldSources(
+ structType: StructType,
+ fieldOps: Seq[StructFieldsOperation]): Seq[(StructField, FieldSource)] =
{
+ val fields: Seq[(StructField, FieldSource)] =
structType.fields.zipWithIndex.map {
+ case (field, ordinal) =>
+ field -> Original(ordinal)
+ }.toImmutableArraySeq
+ fieldOps.foldLeft(fields) {
+ case (current, op: WithField) =>
+ val updated = StructField(op.name, op.valExpr.dataType,
op.valExpr.nullable) ->
+ Updated(op.valExpr)
+ val result = ArrayBuffer.empty[(StructField, FieldSource)]
+ var hasMatch = false
+ current.foreach { fieldSource =>
+ if (op.resolver(fieldSource._1.name, op.name)) {
+ hasMatch = true
+ result += updated
+ } else {
+ result += fieldSource
+ }
+ }
+ if (!hasMatch) result += updated
+ result.toSeq
+ case (current, op: DropField) =>
+ current.filterNot { case (field, _) => op.resolver(field.name,
op.name) }
+ case (_, op) =>
Review Comment:
Thanks, verified. The apply implementations and their scaladoc are gone, and
name matching now lives only in fieldSources.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4235965400","thread_id":"inline:4235965400","verdict_sha256":"ab2b0fd9bcfc655d854e51bb3e2b7b7c0e60761ac6cbf11e41a5d24bd9a888fd"}
-->
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala:
##########
@@ -68,7 +68,9 @@ object OptimizeUpdateFields extends Rule[LogicalPlan] {
val newWithFields = newNames.map(n => WithField(nameMap(n),
newValues(n)))
UpdateFields(structExpr, newWithFields.toSeq)
- case UpdateFields(UpdateFields(struct, fieldOps1), fieldOps2) =>
+ // Flattening would separate a nested field read from the struct it reads.
+ case UpdateFields(UpdateFields(struct, fieldOps1), fieldOps2)
+ if !(fieldOps1 ++
fieldOps2).exists(_.exists(UpdateFields.isGeneratedStructRead)) =>
Review Comment:
Thanks, verified. The flatten case is back without the guard, nested
operations merge in optimizeFieldOps, and the new test bounds analyzed and
optimized plan size for chained nested updates on unresolved columns.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4235965404","thread_id":"inline:4235965404","verdict_sha256":"ab2b0fd9bcfc655d854e51bb3e2b7b7c0e60761ac6cbf11e41a5d24bd9a888fd"}
-->
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:
##########
@@ -764,12 +764,175 @@ case class DropField(name: String) extends
StructFieldsOperation with LeafLike[E
values.filterNot { case (field, _) => resolver(field.name, name) }
}
+/**
+ * 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 Expression {
+
+ override def children: Seq[Expression] = structExpr +: updatedValues
+
+ override def nullable: Boolean = structExpr.nullable
+
+ override lazy val dataType: StructType = {
+ val sourceFields = structExpr.dataType.asInstanceOf[StructType]
+ var updatedIndex = 0
+ StructType(outputOrdinals.zipWithIndex.map { case (sourceOrdinal, index) =>
+ if (sourceOrdinal >= 0) {
+ sourceFields(sourceOrdinal)
+ } else {
+ val updatedValue = updatedValues(updatedIndex)
+ updatedIndex += 1
+ outputFields(index).copy(
+ dataType = updatedValue.dataType,
+ nullable = updatedValue.nullable)
+ }
+ })
+ }
+
+ override def eval(input: InternalRow): Any = {
+ val needsSource = nullable || outputOrdinals.exists(_ >= 0)
+ val sourceRow = if (needsSource) {
+ structExpr.eval(input).asInstanceOf[InternalRow]
+ } else {
+ null
+ }
+ if (sourceRow == null && nullable) {
+ null
+ } else {
+ val values = new Array[Any](outputOrdinals.length)
+ var updatedIndex = 0
+ var i = 0
+ while (i < outputOrdinals.length) {
+ val sourceOrdinal = outputOrdinals(i)
+ if (sourceOrdinal >= 0) {
+ values(i) = if (sourceRow == null) {
+ null
+ } else {
+ sourceRow.get(sourceOrdinal, dataType(i).dataType)
+ }
+ } else {
+ values(i) = updatedValues(updatedIndex).eval(input)
+ updatedIndex += 1
+ }
+ i += 1
+ }
+ new GenericInternalRow(values)
+ }
+ }
+
+ override protected def withNewChildrenInternal(
+ newChildren: IndexedSeq[Expression]): UpdateFieldsExpression =
+ copy(structExpr = newChildren.head, updatedValues = newChildren.tail)
+
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
+ val rowClass = classOf[GenericInternalRow].getName
+ val values = ctx.freshName("values")
+ val needsSource = nullable || outputOrdinals.exists(_ >= 0)
+ val sourceCode = if (needsSource) Some(structExpr.genCode(ctx)) else None
+ val sourceRow = ctx.freshName("sourceRow")
+
+ val copyValuesCode = if (outputOrdinals.exists(_ >= 0)) {
+ val index = ctx.freshName("i")
+ val sourceOrdinal = ctx.freshName("sourceOrdinal")
+ val sourceOrdinals = ctx.addReferenceObj("sourceOrdinals",
outputOrdinals.toArray, "int[]")
+ val fieldTypes = ctx.addReferenceObj(
+ "fieldTypes",
+ dataType.fields.map(_.dataType).toArray,
+ s"${classOf[DataType].getName}[]")
+ s"""
+ |for (int $index = 0; $index < ${outputOrdinals.length}; $index++) {
+ | int $sourceOrdinal = $sourceOrdinals[$index];
+ | if ($sourceOrdinal >= 0) {
+ | $values[$index] = $sourceRow == null ? null :
+ | $sourceRow.get($sourceOrdinal, $fieldTypes[$index]);
Review Comment:
Thanks, verified. Copied fields now go through CodeGenerator.getValue per
type group, and the new vectorized Parquet tests cover TIME, GEOMETRY, and
GEOGRAPHY.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4235965407","thread_id":"inline:4235965407","verdict_sha256":"ab2b0fd9bcfc655d854e51bb3e2b7b7c0e60761ac6cbf11e41a5d24bd9a888fd"}
-->
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:
##########
@@ -806,35 +969,129 @@ case class UpdateFields(structExpr: Expression,
fieldOps: Seq[StructFieldsOperat
override def prettyName: String = "update_fields"
- private lazy val newFieldExprs: Seq[(StructField, Expression)] = {
- def getFieldExpr(i: Int): Expression = structExpr match {
- case c: CreateNamedStruct => c.valExprs(i)
- case _ => GetStructField(structExpr, i)
- }
- val fieldsWithIndex =
structExpr.dataType.asInstanceOf[StructType].fields.zipWithIndex
- val existingFieldExprs: Seq[(StructField, Expression)] =
- fieldsWithIndex.map { case (field, i) => (field, getFieldExpr(i))
}.toImmutableArraySeq
- fieldOps.foldLeft(existingFieldExprs)((exprs, op) => op(exprs))
- }
+ private lazy val newFieldSources: Seq[(StructField,
UpdateFields.FieldSource)] =
+ UpdateFields.fieldSources(structExpr.dataType.asInstanceOf[StructType],
fieldOps)
- private lazy val newFields: Seq[StructField] = newFieldExprs.map(_._1)
+ private lazy val newFields: Seq[StructField] = newFieldSources.map(_._1)
- lazy val newExprs: Seq[Expression] = newFieldExprs.map(_._2)
+ /** Returns the expression for one output field without constructing the
complete struct. */
+ def newExpr(ordinal: Int): Expression = newFieldSources(ordinal)._2 match {
+ case UpdateFields.Original(sourceOrdinal) => getFieldExpr(sourceOrdinal)
+ case UpdateFields.Updated(expr) => expr
+ }
- lazy val evalExpr: Expression = {
- val createNamedStructExpr = CreateNamedStruct(newFieldExprs.flatMap {
- case (field, expr) => Seq(Literal(field.name), expr)
- })
+ /**
+ * Returns an input field expression and marks generated reads so they can
share `structExpr`
+ * without changing equivalent expressions supplied to `withField`.
+ */
+ private def getFieldExpr(ordinal: Int): Expression = structExpr match {
+ case c: CreateNamedStruct => c.valExprs(ordinal)
+ case _ => UpdateFields.markGeneratedStructRead(
Review Comment:
Thanks, the tag-based block is gone and ordinary extraction simplifies
again. One related regression: the revision also removed the merge-target
`!u.structExpr.isInstanceOf[UpdateFields]` guard in SimplifyExtractValueOps. An
extraction from a two-level chain is now rewritten to `If(IsNull(innerUpdate),
...)` before OptimizeUpdateFields can merge the chain, so the inner update's
values run. I left the details in a new comment on ComplexTypes.scala.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4235965413","thread_id":"inline:4235965413","verdict_sha256":"ab2b0fd9bcfc655d854e51bb3e2b7b7c0e60761ac6cbf11e41a5d24bd9a888fd"}
-->
--
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]
