bhollis-dbx commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4276221590
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:
##########
@@ -806,35 +989,151 @@ 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)
+ private[sql] def sqlString(
+ render: Expression => String,
+ quoteField: String => String): String = {
+ def updateSql(inputSql: String, operations: Seq[StructFieldsOperation]):
String = {
+ val operationsSql = operations.map {
+ case nested: UpdateNestedField =>
+ val fieldSql = if (structExpr.resolved) {
+ s"$inputSql.${quoteField(nested.name)}"
+ } else {
+ s"$inputSql[${render(Literal(nested.name))}]"
+ }
+ s"WithField(${updateSql(fieldSql, nested.fieldOps)})"
+ case operation => render(operation)
+ }
+ s"$prettyName(${(inputSql +: operationsSql).mkString(", ")})"
}
- 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))
+ updateSql(render(structExpr), fieldOps)
}
- private lazy val newFields: Seq[StructField] = newFieldExprs.map(_._1)
+ override def sql: String = sqlString(_.sql, quoteIdentifier)
- lazy val newExprs: Seq[Expression] = newFieldExprs.map(_._2)
+ private lazy val newFieldSources: Seq[(StructField,
UpdateFields.FieldSource)] =
+ UpdateFields.fieldSources(structExpr, fieldOps)
- lazy val evalExpr: Expression = {
- val createNamedStructExpr = CreateNamedStruct(newFieldExprs.flatMap {
- case (field, expr) => Seq(Literal(field.name), expr)
- })
+ private lazy val newFields: Seq[StructField] = newFieldSources.map(_._1)
- if (structExpr.nullable) {
- If(IsNull(structExpr), Literal(null, dataType), createNamedStructExpr)
- } else {
- createNamedStructExpr
+ /** 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
+ }
+
+ private def getFieldExpr(ordinal: Int): Expression = structExpr match {
+ case c: CreateNamedStruct => c.valExprs(ordinal)
+ case _ => GetStructField(structExpr, ordinal)
+ }
+
+ lazy val evalExpr: Expression = {
+ def buildExpression(structRef: Expression): Expression = {
+ val fieldSources = UpdateFields.fieldSources(structRef, fieldOps)
+ val updatedValues = ArrayBuffer.empty[Expression]
+ val outputOrdinals = fieldSources.map {
+ case (_, UpdateFields.Original(ordinal)) => ordinal
+ case (_, UpdateFields.Updated(expr)) =>
+ updatedValues += expr
+ -1
+ }
+ UpdateFieldsExpression(structRef, updatedValues.toSeq, outputOrdinals,
newFields)
+ }
+ structExpr match {
+ case namedStruct: CreateNamedStruct =>
+ val values = namedStruct.valExprs.toArray
+ CreateNamedStruct(newFieldSources.flatMap { case (field, source) =>
+ val value = source match {
+ case UpdateFields.Original(ordinal) => values(ordinal)
+ case UpdateFields.Updated(expr) => expr
+ }
+ Seq(Literal(field.name), value)
+ })
+ case _ if fieldOps.exists(_.isInstanceOf[UpdateNestedField]) =>
+ With(structExpr) { case Seq(structRef) => buildExpression(structRef) }
Review Comment:
Fixed in 1c5c87f6b. CSE now traverses always-evaluated With definitions
while respecting guards and nested scopes. Added withField/dropFields tests
confirming a shared struct input evaluates once per row on all three evaluation
paths.
##########
sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:
##########
@@ -445,6 +446,428 @@ 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("field updates should preserve shared value subexpression elimination")
{
+ withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") {
+ for {
+ nested <- Seq(false, true)
+ external <- Seq(false, true)
+ opaque <- Seq(false, true)
+ } {
+ val calls = spark.sparkContext.longAccumulator
+ val sourceCalls = spark.sparkContext.longAccumulator
+ val value = udf((id: Long) => {
+ calls.add(1)
+ id + 1
+ }).apply($"id")
+ val source = if (opaque) {
+ udf((id: Long) => {
+ sourceCalls.add(1)
+ ((id, id), id)
+ }).asNonNullable().asNondeterministic()($"id")
+ } else {
+ struct(struct($"id", $"id").as("_1"), $"id".as("_2"))
+ }
+ val prefix = if (nested) "_1." else ""
+ val updated = source.withField(s"${prefix}first", value)
+ .withField(s"${prefix}second", value)
+ val columns = if (external) Seq(updated, value) else Seq(updated)
+ val rows = spark.range(0, 10, 1, 1).select(columns: _*).collect()
+ rows.zipWithIndex.foreach { case (row, index) =>
+ val result = if (nested) row.getStruct(0).getStruct(0) else
row.getStruct(0)
+ assert(result.getLong(2) == index + 1L)
+ assert(result.getLong(3) == index + 1L)
+ if (external) assert(row.getLong(1) == index + 1L)
+ }
+ val expectedCalls = if (nested && opaque) 10 * (if (external) 3 else
2) else 10
Review Comment:
Added a test with a non-nullable nested parent, sharing one value across two
nested updates and another column. It verifies results and once-per-row
evaluation on all three paths. Disabling With-body traversal makes it fail with
30 UDF calls instead of 10.
--
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]
