HyukjinKwon commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4235965395
##########
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:
**Blocking (P1):** This only shares the input when the generated read's
child is the *same object* as `structExpr`. On the common Column API path the
struct is unresolved when `UpdateFields.apply` builds the tree, for example
`col("s").withField("a.b", lit(1)).withField("a.c", lit(2))`.
`ColumnResolutionHelper.innerResolve` then resolves each occurrence of that
shared instance separately (`mapChildren`, no memoization), so after analysis
the two positions hold distinct but equal copies and `eq` fails. Because the
new guards in `OptimizeUpdateFields` and `SimplifyExtractValueOps` also skip
these chains, nothing collapses the copies either. At each level of the chain
the previous update appears twice, so `n` chained nested-path calls run about
`2^(n-1)` lowered copies of the first update per row, and CSE can't merge them
because it doesn't descend into `With`. The merge target reduced the same chain
to a single `UpdateFields(s, a -> UpdateFields(GetStructField(s, a), ...))`.
For non
deterministic inputs, the evaluate-once guarantee doesn't hold on this path
either.
See **Shared repair plan 1** in the review body.
##########
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:
**Non-blocking (P2):** The previous lowering read each copied field through
`GetStructField` codegen, i.e. `CodeGenerator.getValue` with typed getters
(`getLong`, `getBinaryView`, `getInterval`, ...). The generic `get(ordinal,
DataType)` here isn't supported by every `InternalRow`: `ColumnarRow.get`
(ColumnarRow.java:191-235) has no branch for `TimeType`, `GeometryType`,
`GeographyType`, or `CalendarIntervalType` and throws
`_LEGACY_ERROR_TEMP_3155`. Structs from the vectorized Parquet reader arrive as
`ColumnarRow` (`vector.getStruct(rowId)`), so with geospatial enabled by
default, `spark.read.parquet(p).select($"s".withField("b", lit(1)))` on `s
STRUCT<g: GEOMETRY(4326), a: INT>` now fails at runtime, where it used to
succeed. The same goes for TIME fields when the time type is enabled. This only
affects the code-generated path; interpreted `GetStructField` already used the
generic `get`. Emitting `CodeGenerator.getValue(sourceRow, fieldType, ordinal)`
per copied ordinal would k
eep the previous compatibility.
**Verification:**
- **Regression:** A vectorized Parquet struct with an unchanged GEOMETRY or
TIME field passes through withField under default codegen without
SparkUnsupportedOperationException.
##########
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:
**Non-blocking (P2):** `newExpr`'s only caller is `SimplifyExtractValueOps`,
so this tag ends up on ordinary single-level extractions, not just on
nested-path reads. `$"t".withField("y", lit(2)).getField("a")` becomes a tagged
`GetStructField(t, a)`, and the tag stays because `copyTagsFrom` only fills
untagged nodes. If `t` is later inlined as `s.withField("x", lit(1) / $"s.b")`
(e.g. by CollapseProject), the new `!isGeneratedStructRead(field)` guard
refuses to simplify `GetStructField(UpdateFields(s, x), a)` to `s.a`. The full
update, including `x`, is then evaluated, so with ANSI on and `s.b = 0` this
now fails with DIVIDE_BY_ZERO where it used to return rows. A simplified read
like this inside an ops list also trips the flatten guard at
UpdateFields.scala:73, even though no nested-path read would be separated. This
conflicts with the comment that says generated reads are "field accesses
introduced while expanding nested paths".
See **Shared repair plan 1** in the review body.
##########
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.
Review Comment:
**Nit (P3):** This reads as the opposite of what the next lines do.
`replaceStruct` is applied to every `Updated(expr)`, and for a nested path that
value is the generated `UpdateFields(GetStructField(structExpr, a), ...)` whose
tagged read *is* rewritten to `structRef`. The actual rule is narrower: only
reads generated by this update's own nested-path expansion are rebound, and
other reads of `structExpr` inside withField values are left alone. Maybe
something like: "Leave other reads in withField values unchanged; they must be
evaluated independently."
**Verification:**
- **Inspection:** The comment states the owner-tagged nested-read rule and
no longer claims that withField values are never rewritten.
##########
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)
Review Comment:
**Nit (P3):** `outputOrdinals` and `updatedValues` are Scala 2.13 `List`s
here whenever a `WithField` is present (`ArrayBuffer.toSeq` returns a `List`,
and so does its `.map`). This loop calls `outputOrdinals.length`,
`outputOrdinals(i)`, and `updatedValues(updatedIndex)` for every field of every
row, so interpreted evaluation is O(width^2) per row. That's roughly 125k+ list
hops per row for a 500-field struct. `needsSource` is also recomputed with an
`exists` scan on every row. Codegen isn't affected because it copies into an
`int[]`, but interpreted mode is what you get when codegen falls back on very
large expressions. Building indexed collections (`toArray`/`IndexedSeq`, or
`@transient lazy val` arrays) and hoisting `needsSource` into a lazy val would
keep this linear.
**Verification:**
- **Inspection:** The per-row eval loop indexes only array-backed or
IndexedSeq state, and needsSource is computed once per expression instance.
##########
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:
**Blocking (P1):** This guard also runs in the analyzer:
`OptimizeUpdateFields` is in the Substitution batch (Analyzer.scala:641-645)
specifically because "when manipulating deeply nested schema, UpdateFields
expression tree could be very complex and make analysis impossible". Each
nested-path step `u_k = UpdateFields(u_{k-1}, a ->
UpdateFields(ExtractValue(u_{k-1}, a), op_k))` references `u_{k-1}` twice.
Before, flatten plus dedupe collapsed every step before resolution, so the tree
stayed linear. Now `(fieldOps1 ++
fieldOps2).exists(_.exists(isGeneratedStructRead))` is true for every
nested-path step, so nothing collapses, and resolution and later traversals
walk about `2^N` occurrences for N chained calls. Examples are
`col("s").withField("a.x1", ..)...withField("a.xN", ..)` or `dropFields("a.x1",
..., "a.xN")`, which `Column.dropFields` folds into such a chain. Programmatic
updates of many nested fields can stall analysis before execution even starts.
See **Shared repair plan 1** in the review body.
##########
sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:
##########
@@ -445,6 +446,141 @@ class ColumnExpressionSuite extends SharedSparkSession {
SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY")(f)
}
+ test("withField should evaluate a nondeterministic struct expression once") {
+ onEachEvalPath {
+ val counter = new java.util.concurrent.atomic.AtomicLong()
+ val nextStruct = udf(() => {
+ val value = counter.getAndIncrement()
+ (value, value)
+ }).asNondeterministic()
+ val structExpr = nextStruct()
+ val df = spark.range(0, 10, 1, 1)
+ df.select(structExpr.withField("copy", structExpr)).collect().foreach {
row =>
+ val result = row.getStruct(0)
+ val copy = result.getStruct(2)
+ assert(result.getLong(0) == result.getLong(1))
+ assert(copy.getLong(0) == copy.getLong(1))
+ assert(result.getLong(0) != copy.getLong(0))
+ }
+ }
+ }
+
+ test("withField should evaluate a reused nondeterministic value separately")
{
+ onEachEvalPath {
+ val counter = new java.util.concurrent.atomic.AtomicLong()
+ val nextStruct = udf(() => {
+ val value = counter.getAndIncrement()
+ (value, value)
+ }).asNondeterministic()
+ val structExpr = nextStruct().expr
+ val updated = Column(UpdateFields(structExpr,
+ Seq(WithField("copy", GetStructField(structExpr, 0)))))
+ val df = spark.range(0, 10, 1, 1)
+ df.select(updated).collect().foreach { row =>
+ val result = row.getStruct(0)
+ assert(result.getLong(0) == result.getLong(1))
+ assert(result.getLong(0) != result.getLong(2))
+ }
+ }
+ }
+
+ test("withField should evaluate a reused nested update separately") {
+ onEachEvalPath {
+ val counter = new java.util.concurrent.atomic.AtomicLong()
+ val nextStruct = udf(() => {
+ val value = counter.getAndIncrement()
+ (value, value)
+ }).asNondeterministic()
+ val structExpr = nextStruct().expr
+ val nestedUpdate = UpdateFields(structExpr, Seq(WithField("extra",
Literal(1))))
+ val updated = Column(UpdateFields(structExpr,
+ Seq(WithField("copy", GetStructField(nestedUpdate, 0)))))
+ val df = spark.range(0, 10, 1, 1)
+ df.select(updated).collect().foreach { row =>
+ val result = row.getStruct(0)
+ assert(result.getLong(0) == result.getLong(1))
+ assert(result.getLong(0) != result.getLong(2))
+ }
+ }
+ }
+
+ test("withField should reuse an updated struct for a nested path") {
Review Comment:
**Non-blocking (P2):** This test (and `preserve a nested path through a
later update`) builds `UpdateFields` directly over `nextStruct().expr`, a
zero-argument UDF that is already resolved. Identity therefore survives
analysis and the `eq`-based sharing works. Users reach this code through
`Column.withField` on input-dependent expressions such as
`udf(...)($"id").withField("_1.copy", lit(1))` or `col("s")...`, where analysis
rebuilds the struct. `ReplaceUpdateFieldsExpressionSuite` also runs only
`ReplaceUpdateFieldsExpression` on pre-resolved `AttributeReference` inputs,
skipping both the analyzer and the operator-optimization batch. So none of the
new tests would fail for the exponential chains, the evaluated dropped values,
or the double evaluation on the normal DataFrame/Connect path. Could you add
full-pipeline DataFrame tests for chained nested-path `withField`/`dropFields`
on unresolved columns, including a drop or overwrite after a nested update?
See **Shared repair plan 1** in the review body.
##########
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:
**Nit (P3):** `fieldSources` re-implements the `WithField`/`DropField` name
matching instead of going through `StructFieldsOperation.apply`, and nothing in
production calls `apply` anymore. That leaves two copies of the same semantics
that can drift apart, a trait scaladoc ("...used as the fields argument for
StructType and as the children argument for CreateNamedStruct inside
UpdateFields") that is no longer true, and an internal error for any other
operation. Could `apply` be made generic over the value type (e.g.
`apply[T](values: Seq[(StructField, T)], updated: Expression => T)`) and called
from here? Or remove it and update the doc so there's a single source of truth.
**Verification:**
- **Compatibility:** Existing withField/dropFields name-matching behavior,
including case sensitivity and duplicate names, is unchanged.
--
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]
