ohad-definity commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4416818939
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruning.scala:
##########
@@ -133,6 +134,63 @@ object SchemaPruning extends SQLConfHelper {
} ++ rootFields
}
+ /**
+ * Gets the root field that the array `input` of an `explode` or
`posexplode` needs, when the
+ * generator's element output `element` is only read by `exprs` through
nested field accesses.
+ * For example, for `explode(items) AS item` read as `item.a` and `item.b`,
with
+ * `items: array<struct<a, b, c>>`, it returns `items: array<struct<a, b>>`.
+ *
+ * `elementFields` are fields of the element read by other means, e.g. as
the input of a
+ * generator above this one, named like the element.
+ *
+ * If nothing of the element is read, only the array sizes are needed, so a
single smallest
+ * field is kept, as in
[[org.apache.spark.sql.catalyst.optimizer.GenerateOptimization]].
+ *
+ * If the element is used whole, `input` is needed whole, e.g. `items.parts`
for
+ * `explode(items.parts)`, which still narrows `items` to `parts`.
+ *
+ * `input` must be expressed over the scan's output, or over the element of
a generator below
+ * it. Returns None if `input` is not an access path [[SelectedField]] can
narrow.
+ */
+ def getGeneratorInputRootField(
+ input: Expression,
+ element: Attribute,
+ exprs: Seq[Expression],
+ elementFields: Seq[StructField] = Nil): Option[RootField] =
input.dataType match {
+ case ArrayType(elementType, containsNull) =>
+ // Rename the element, and every other attribute too, so that the
element's fields can be
+ // told apart by name from other columns and generator outputs.
+ val name = s"${element.name}#${element.exprId.id}"
+ val renamed = exprs.map(_.transform {
+ case a: Attribute if a.exprId == element.exprId => a.withName(name)
+ case a: Attribute => a.withName(s"${a.name}#${a.exprId.id}")
+ })
+ val fields = identifyRootFields(Nil, renamed,
+ elementFields.map(f => RootField(f.copy(name = name), derivedFromAtt =
false)))
+ .filter(_.field.name == name)
+ val prunedElementType = (fields, elementType) match {
+ // not read: only the array sizes are needed
+ case (Seq(), StructType(structFields)) if structFields.length > 1 =>
+ StructType(Array(structFields.minBy(_.dataType.defaultSize)))
+ case (Seq(), _) => elementType
+ // used whole: the input is needed whole, which still narrows what it
is extracted from
+ case _ if fields.exists(_.derivedFromAtt) => elementType
+ case _ =>
+ pruneSchema(StructType(Array(StructField(name, elementType))),
fields).head.dataType
+ }
+ SelectedField.withDataType(input, ArrayType(prunedElementType,
containsNull))
Review Comment:
Confirmed and fixed in d9120db1: a generator input whose path goes through
`map_keys` is no longer narrowed, so the keys are read whole as on master
(where ordinary nested pruning does not narrow map keys either). Your repro is
a new test, "an explode of map keys keeps the keys whole", with keys equal
without their `id`; it checks the scan schema and the rows (positions 0 and 1)
under all reader modes, ORC row reader included.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala:
##########
@@ -40,6 +49,19 @@ object SchemaPruning extends Rule[LogicalPlan] {
override def apply(plan: LogicalPlan): LogicalPlan =
plan transformDown {
+ case op @ GeneratorsOverScan(_, _,
+ l @ LogicalRelationWithTable(hadoopFsRelation: HadoopFsRelation, _))
=>
+ // Fields read from the generators' elements above them narrow the
scan below them.
+ pruneThroughGenerators(op, l) { (scan, otherProjects, inputRootFields)
=>
+ val allFilters = scan.filtersPushDown.reduceOption(And).toSeq ++
scan.filtersStayUp
+ prunePhysicalColumns(l, scan.projects, allFilters, hadoopFsRelation,
+ (prunedDataSchema, prunedMetadataSchema) => {
+ val prunedHadoopRelation = hadoopFsRelation.copy(
+ dataSchema = prunedDataSchema)(hadoopFsRelation.sparkSession)
Review Comment:
Thanks for the precise repro. This one is not specific to generators: on
current master (nightly 2026-10-29), ordinary nested schema pruning hits the
same check without any explode. With your file (`items =
array(named_struct('a', TIMESTAMP ..., 'b', TIMESTAMP_NTZ ..., 'c', 22))`) on
ORC V1:
```sql
SELECT items.b, items.c FROM t
```
| `nestedSchemaPruning.enabled` | vectorized reader | row reader |
|---|---|---|
| false | `[2025-01-02T01:02:03], [22]` | same |
| true | `UNSUPPORTED_FEATURE.ORC_TYPE_CAST` "TIMESTAMP" to "TIMESTAMP_NTZ"
| same |
So `OrcUtils.requestedColumnIds` already compares the pruned nested schema
with the physical one by ordinal (`checkTimestampCompatibility`) for any pruned
array of structs; this PR makes the same path reachable through an explode. The
fix belongs there (match nested fields by name before the compatibility check,
as the top-level case does), so I'd suggest a separate JIRA and PR for it,
which this PR would then pick up. ORC V2 is not affected on master. Happy to
send that fix first if you prefer.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -1015,8 +1016,56 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
Project(rewrittenProjectList, withFilter)
}
- def pruneColumns(plan: LogicalPlan): LogicalPlan = plan.transform {
+ def pruneColumns(plan: LogicalPlan, isSubquery: Boolean): LogicalPlan =
plan.transform {
+ case op @ SchemaPruning.GeneratorsOverScan(_, generators, sHolder:
ScanBuilderHolder)
+ if !hasBlockingPushdown(sHolder) =>
+ // Fields read from the generators' elements above them narrow the scan
below them. This
+ // prunes and builds the holder's builder before the plan above is
re-resolved. If that is
+ // then rejected, the builder must not be pruned again with a wider
schema (a builder may
+ // only narrow, e.g. Iceberg's), so the plain scan is built from a fresh
builder instead, as
+ // scan merging does, and the lower generators of the chain do not try
again.
+ var builderUsed = false
+ SchemaPruning.pruneThroughGenerators(op, sHolder) { (scan,
otherProjects, inputRootFields) =>
+ builderUsed = true
+ Some(pruneScan(scan.projects, scan.filtersStayUp,
scan.filtersPushDown, sHolder,
Review Comment:
Same here: this reproduces on current master (nightly 2026-10-29) with
ordinary nested schema pruning, without generators. Your file read through ORC
V2 with `orc.force.positional.evolution=true`:
```sql
SELECT items.b, items.c FROM t
```
| `nestedSchemaPruning.enabled` | ORC V2, vectorized | ORC V2, row reader |
|---|---|---|
| false | `[22], [33]` | `[22], [33]` |
| true | `[11], [22]` (wrong) | `[11], [22]` (wrong) |
ORC V1 returns the right values in both cases. So nested pruning on ORC V2
under positional evolution is already wrong on master; this PR only adds
another way to reach it. I'd suggest a separate JIRA and PR that skips nested
pruning (in both paths) when positional evolution is enabled for the scan;
happy to send that first, or to add the same guard to the generator path here
if you'd rather not wait.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -1015,8 +1016,56 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
Project(rewrittenProjectList, withFilter)
}
- def pruneColumns(plan: LogicalPlan): LogicalPlan = plan.transform {
+ def pruneColumns(plan: LogicalPlan, isSubquery: Boolean): LogicalPlan =
plan.transform {
+ case op @ SchemaPruning.GeneratorsOverScan(_, generators, sHolder:
ScanBuilderHolder)
+ if !hasBlockingPushdown(sHolder) =>
+ // Fields read from the generators' elements above them narrow the scan
below them. This
+ // prunes and builds the holder's builder before the plan above is
re-resolved. If that is
+ // then rejected, the builder must not be pruned again with a wider
schema (a builder may
+ // only narrow, e.g. Iceberg's), so the plain scan is built from a fresh
builder instead, as
+ // scan merging does, and the lower generators of the chain do not try
again.
+ var builderUsed = false
+ SchemaPruning.pruneThroughGenerators(op, sHolder) { (scan,
otherProjects, inputRootFields) =>
+ builderUsed = true
+ Some(pruneScan(scan.projects, scan.filtersStayUp,
scan.filtersPushDown, sHolder,
+ Some(otherProjects), inputRootFields))
+ }.getOrElse {
+ if (builderUsed) {
+ val bottom = generators.last.child
+ val rebuilt = rebuildWithoutGenerators(bottom, sHolder, isSubquery)
+ op.transformDown { case p if p eq bottom => rebuilt }
+ } else {
+ op
+ }
+ }
+
case ScanOperation(project, filtersStayUp, filtersPushDown, sHolder:
ScanBuilderHolder) =>
+ pruneScan(project, filtersStayUp, filtersPushDown, sHolder, None, Nil)
+ }
+
+ /**
+ * Plans `bottom`, the projects and filters over `sHolder` below a chain of
generators, from a
+ * fresh `ScanBuilder` of the same relation, re-applying the filters pushed
to `sHolder`, like
+ * [[rebuildScan]]. Only valid when `sHolder` has no pushdown such a rebuild
cannot reproduce.
+ */
+ private def rebuildWithoutGenerators(
+ bottom: LogicalPlan,
+ sHolder: ScanBuilderHolder,
+ isSubquery: Boolean): LogicalPlan = {
+ val relation = sHolder.relation
+ val withPushedFilters = sHolder.pushedFilterExpressions.reduceOption(And)
+ .map(Filter(_, relation)).getOrElse(relation)
Review Comment:
Fixed in d9120db1: the fresh builder now gets the filters pushed to the
holder and the scan segment's pushable filters together, in one `Filter`
directly over the relation (like `rebuildScan`), with the non-deterministic
ones above it. New test "a rejected rewrite offers the data filters to the
fresh builder" forces rejection on V2 Parquet and asserts the `n = 7` filter is
in the `ParquetScan`'s pushed filters as well as the rows; the
narrow-only-builder test now also asserts that the rejection used a second
builder.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala:
##########
@@ -372,6 +408,525 @@ abstract class SchemaPruningSuite
}
}
+ testSchemaPruning("SPARK-34956: select multiple fields of an exploded array
of struct") {
+ val exploded =
spark.table("contacts").select(explode(col("friends")).as("friend"))
+
+ val query1 = exploded.select(col("friend.first"), col("friend.last"))
+ checkScan(query1,
"struct<friends:array<struct<first:string,last:string>>>")
+ checkAnswer(query1, Row("Susan", "Smith") :: Nil)
+
+ // a filter between the generator and the projection reads its fields too
+ // (`first` is also read by the `size(friends) > 0` filter inferred for
the explode)
+ val query2 =
exploded.where(col("friend.middle").isNotNull).select(col("friend.first"))
+ checkScan(query2,
"struct<friends:array<struct<first:string,middle:string>>>")
+ checkAnswer(query2, Row("Susan") :: Nil)
+
+ val query3 = sql(
+ "select id, pos, friend.first, friend.last from contacts " +
+ "lateral view outer posexplode(friends) t as pos, friend")
+ checkScan(query3,
"struct<id:int,friends:array<struct<first:string,last:string>>>")
+ // contacts 2 and 3 have no friends column, so the outer generator gives
them a null row too
+ checkAnswer(query3,
+ Row(0, 0, "Susan", "Smith") :: Row(1, null, null, null) ::
+ Row(2, null, null, null) :: Row(3, null, null, null) :: Nil)
+ }
+
+ testSchemaPruning("SPARK-34956: an exploded element used whole is not
pruned") {
+ val query = spark.table("contacts")
+ .select(explode(col("friends")).as("friend"))
+ .select(col("friend"), col("friend.first"))
+ checkScan(query,
"struct<friends:array<struct<first:string,middle:string,last:string>>>")
+ checkAnswer(query, Row(Row("Susan", "Z.", "Smith"), "Susan") :: Nil)
+ }
+
+ testSchemaPruning("SPARK-34956: an unused exploded element reads a single
field") {
+ // only the array sizes are needed; the `size(friends) > 0` filter
inferred for the inner
+ // explode already reads the same (smallest) field
+ val query1 = sql("select id from contacts lateral view explode(friends) t
as friend")
+ checkScan(query1, "struct<id:int,friends:array<struct<first:string>>>")
+ checkAnswer(query1, Row(0) :: Nil)
+
+ val query2 = sql("select id from contacts lateral view outer
explode(friends) t as friend")
+ checkScan(query2, "struct<id:int,friends:array<struct<first:string>>>")
+ checkAnswer(query2, Row(0) :: Row(1) :: Row(2) :: Row(3) :: Nil)
+ }
+
+ testSchemaPruning("SPARK-34956: aggregate over fields of an exploded array
of struct") {
+ val query = sql(
+ "select id, count(friend.first), max(friend.last) " +
+ "from contacts lateral view explode(friends) t as friend group by id")
+ checkScan(query,
"struct<id:int,friends:array<struct<first:string,last:string>>>")
+ checkAnswer(query, Row(0, 1L, "Smith") :: Nil)
+ }
+
+ testSchemaPruning("SPARK-34956: fields of an exploded array of struct -
case-insensitivity") {
+ withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") {
+ val query = sql(
+ "select friend.FIRST, friend.Last from contacts lateral view
explode(friends) t as friend")
+ checkScan(query,
"struct<friends:array<struct<first:string,last:string>>>")
+ checkAnswer(query, Row("Susan", "Smith") :: Nil)
+ }
+ }
+
+ testSchemaPruning("SPARK-34956: pruning through a generator can be
disabled") {
+ withSQLConf(SQLConf.NESTED_SCHEMA_PRUNING_THROUGH_GENERATE_ENABLED.key ->
"false") {
+ val query = spark.table("contacts")
+ .select(explode(col("friends")).as("friend"))
+ .select(col("friend.first"), col("friend.last"))
+ checkScan(query,
"struct<friends:array<struct<first:string,middle:string,last:string>>>")
+ checkAnswer(query, Row("Susan", "Smith") :: Nil)
+ }
+ }
+
+ testSchemaPruning("SPARK-34956: a single field is still pushed into the
generator") {
+ val query = spark.table("contacts")
+ .select(explode(col("friends")).as("friend"))
+ .select(col("friend.first"))
+ checkScan(query, "struct<friends:array<struct<first:string>>>")
+ // GeneratorNestedColumnAliasing rewrote `explode(friends)` into
`explode(friends.first)`
+ // before schema pruning, which leaves nothing to prune through the
generator
+ val generatorInputTypes = query.queryExecution.optimizedPlan.collect {
+ case g: Generate => g.generator.children.head.dataType
+ }
+ assert(generatorInputTypes === Seq(ArrayType(StringType)))
+ checkAnswer(query, Row("Susan") :: Nil)
+ }
+
+ private val explodeOrderParts = "from orders lateral view outer
explode(items) t1 as item " +
+ "lateral view outer explode(item.parts) t2 as part"
+
+ testSchemaPruning("SPARK-47230: a field of the inner element of two chained
explodes") {
Review Comment:
Done in d9120db1: the 25 fixture-only cases are plain `test`s around
`withDataSourceTable`; the ones that also query `contacts` keep
`testSchemaPruning`.
--
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]
