diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala index 6695e24ad618..8f43fac0e1c1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala @@ -19,8 +19,8 @@ package org.apache.spark.sql.execution.datasources import scala.collection.mutable -import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeSet, Cast, Expression, - NamedExpression, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, AttributeSet, Cast, + Expression, NamedExpression, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression import org.apache.spark.sql.catalyst.expressions.variant.VariantGet import org.apache.spark.sql.catalyst.plans.LeftExistence @@ -105,20 +105,10 @@ import org.apache.spark.sql.types.VariantType * [[SQLConf.PUSH_VARIANT_INTO_SCAN]] is also enabled) and only fires when a hoistable extraction is * present, so non-variant plans are untouched. * - * Cast-error surface: relocating a strict extraction (`variant_get(..., failOnError = true)` or a - * strict `Cast`) below a `Join` means it is evaluated at the scan on rows the join later eliminates - * -- so a cast failure can surface for a row the un-hoisted plan would never have cast (here the - * eliminating rows come from the *other* joined table). This is the same pre-existing trade-off as - * [[PushVariantIntoScan]] pushing casts below a `Filter`, not a new error class: in both, the - * strict cast runs before the operator that would have discarded the failing row. This rule only - * relocates the extraction into a `Project`; [[PushVariantIntoScan]] still does the scan-level - * materialization and, when [[SQLConf.PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR]] is set (default - * false), wraps the cast with a per-row cast-error companion slot so the error is only raised when - * the original expression consumes the failing row. That deferral is provenance-agnostic -- it acts - * on the relocated extraction regardless of how it reached the `Project` -- so enabling the flag - * suppresses the join-eliminated-row error exactly as it does the filter-eliminated-row case. With - * the flag off (the default), the strict cast raises immediately on any failing scanned row, as - * documented for below-`Filter` pushdown. + * Cast-error surface: relocating a throwable extraction below a `Join` means it is evaluated at the + * scan on rows the join later eliminates. Such extractions cross a join only when + * [[SQLConf.PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR]] is enabled. Expressions not classified as + * throwable retain the existing behavior, including strict `VariantGet` and `Cast` expressions. */ object PullOutVariantExtractions extends Rule[LogicalPlan] { @@ -177,11 +167,16 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { case _ => false } + private def isJoinHoistable(e: Expression): Boolean = { + isHoistable(e) && (!e.throwable || + SQLConf.get.getConf(SQLConf.PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR)) + } + /** * Collects hoisted extractions as aliases, de-duplicated by canonical form so a repeated * extraction maps to a single output slot. */ - private class ExtractionHoister { + private class ExtractionHoister(hoistable: Expression => Boolean = isHoistable) { private val extracted = mutable.LinkedHashMap.empty[Expression, Alias] def aliases: Seq[NamedExpression] = extracted.values.toSeq @@ -194,10 +189,14 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { /** Replaces every hoistable extraction in `e` with a reference to its (new) alias. */ def hoist(e: Expression): Expression = e.transformDown { - case ex if isHoistable(ex) => aliasFor(ex) + case ex if hoistable(ex) => aliasFor(ex) } } + private def extractionHoister(crossesJoin: Boolean): ExtractionHoister = { + if (crossesJoin) new ExtractionHoister(isJoinHoistable) else new ExtractionHoister + } + // Recursively pushes hoisted extraction aliases down through a join tree until each lands in a // `Project` directly above a non-join child (the scan side, or the `Filter`/`Project` chain above // it that `PhysicalOperation` collapses). Hoisting an extraction into a `Project` above a `Join` @@ -299,7 +298,7 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { rightHoister: ExtractionHoister): Seq[NamedExpression] = { projectList.map { e => e.transformDown { - case ex if isHoistable(ex) => + case ex if isJoinHoistable(ex) => if (ex.references.subsetOf(leftOutput)) { leftHoister.aliasFor(ex) } else if (ex.references.subsetOf(rightOutput)) { @@ -312,7 +311,7 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { } private def rewriteAggregate(agg: Aggregate): LogicalPlan = { - val hoister = new ExtractionHoister + val hoister = extractionHoister(agg.child.containsPattern(JOIN)) // Only hoist extractions that sit inside an aggregate function's arguments (or filter). A // top-level extraction in `aggregateExpressions` is a grouping-key reference (grouping keys are // already pulled out by `PullOutGroupingExpressions`); hoisting it would leave the `Aggregate` @@ -388,7 +387,7 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { private def rewriteSortUnderProject( project: Project, projectList: Seq[NamedExpression], sort: Sort): LogicalPlan = { - val hoister = new ExtractionHoister + val hoister = extractionHoister(sort.child.containsPattern(JOIN)) val newOrder = sort.order.map(hoister.hoist(_).asInstanceOf[SortOrder]) if (hoister.isEmpty) { project @@ -419,7 +418,7 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { // yielding a multi-slot shredded struct with a full-variant slot -- the same shape as a // Sort-under-Project whose `v` is also selected. See the class doc. private def rewriteBareSort(sort: Sort): LogicalPlan = { - val hoister = new ExtractionHoister + val hoister = extractionHoister(sort.child.containsPattern(JOIN)) val newOrder = sort.order.map(hoister.hoist(_).asInstanceOf[SortOrder]) if (hoister.isEmpty) { sort @@ -442,8 +441,8 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { project: Project, projectList: Seq[NamedExpression], join: Join): LogicalPlan = { val leftOutput = join.left.outputSet val rightOutput = join.right.outputSet - val leftHoister = new ExtractionHoister - val rightHoister = new ExtractionHoister + val leftHoister = extractionHoister(crossesJoin = true) + val rightHoister = extractionHoister(crossesJoin = true) // A single variant extraction references exactly one attribute, hence one join side; route it // to that side's hoister. Anything not cleanly on one side is left in place. We hoist from both @@ -451,7 +450,7 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { // join (its extractions -- e.g. aggregate arguments hoisted here by `rewriteAggregate`, or a // user's `SELECT variant_get(...)`), so the pushdown sees them below the join. val newCondition = join.condition.map(_.transformDown { - case ex if isHoistable(ex) => + case ex if isJoinHoistable(ex) => if (ex.references.subsetOf(leftOutput)) { leftHoister.aliasFor(ex) } else if (ex.references.subsetOf(rightOutput)) { @@ -481,7 +480,13 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { left = pushSideAliases(join.left, leftHoister.aliases, needed), right = pushSideAliases(join.right, rightHoister.aliases, needed), condition = newCondition) - project.copy(projectList = newProjectList, child = newJoin) + // Outer joins widen attributes on their nullable side. Remap references in the parent + // Project to the copied Join's actual output so their nullability matches the child. + val newJoinOutput = AttributeMap(newJoin.output.map(attr => attr -> attr)) + val remappedProjectList = newProjectList.map(_.transformDown { + case attr: Attribute => newJoinOutput.getOrElse(attr, attr) + }.asInstanceOf[NamedExpression]) + project.copy(projectList = remappedProjectList, child = newJoin) } } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala index d6ed274327a6..f796c9c0050f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala @@ -21,6 +21,7 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.QueryTest import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.variant._ +import org.apache.spark.sql.catalyst.plans.LeftOuter import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation import org.apache.spark.sql.internal.SQLConf @@ -776,6 +777,34 @@ trait PushVariantIntoScanSuiteBase extends SharedSparkSession { } } + test("left outer join: projected right-side attributes use widened nullability") { + val leftKey = AttributeReference("leftKey", IntegerType, nullable = false)() + val rightKey = AttributeReference("rightKey", IntegerType, nullable = false)() + val rightVariant = AttributeReference("rightVariant", VariantType, nullable = false)() + val join = Join( + LocalRelation(leftKey), + LocalRelation(rightKey, rightVariant), + LeftOuter, + Some(EqualTo(leftKey, rightKey)), + JoinHint.NONE) + val extraction = VariantGet( + rightVariant, + path = Literal("$.label"), + targetType = StringType, + failOnError = false, + timeZoneId = Some(localTimeZone)) + val plan = Project(Seq(rightKey, Alias(extraction, "label")()), join) + + PullOutVariantExtractions(plan) match { + case Project(projectList, rewrittenJoin: Join) => + val projectedRightKey = projectList.head.asInstanceOf[Attribute] + val childRightKey = rewrittenJoin.output.find(_.exprId == rightKey.exprId).get + assert(projectedRightKey.nullable) + assert(projectedRightKey.nullable == childRightKey.nullable) + case other => fail(s"Expected a Project over Join, but got:\n$other") + } + } + test("full outer join: extractions on both (nullable) sides are pushed and value-preserving") { // In a FULL OUTER join both sides are nullable: unmatched rows are null-padded on the opposite // side. Extractions on each side push onto that side, computed before null-padding; the join