HyukjinKwon commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4241669003
##########
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:
Thanks, the per-row loop now uses the arrays and is linear. The same List
indexing is still in `dataType` (outputFields(index),
updatedValues(updatedIndex)) and in `doGenCode`
(outputOrdinals.indices.filter(outputOrdinals(_) < 0)), which makes planning
quadratic in width. Reusing indexedOrdinals/indexedUpdatedValues or zipped
iterators there would fix it.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4235965417","thread_id":"inline:4235965417","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.
Review Comment:
Thanks, verified. The comment went away with the rewriting machinery it
described.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4235965418","thread_id":"inline:4235965418","verdict_sha256":"ab2b0fd9bcfc655d854e51bb3e2b7b7c0e60761ac6cbf11e41a5d24bd9a888fd"}
-->
##########
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:
Thanks, verified. The new DataFrame tests go through analysis and
optimization for unresolved nested chains, later drops and overwrites, and an
input-dependent nondeterministic struct, which is what I asked for.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4235965422","thread_id":"inline:4235965422","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]
