ohad-definity commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4425793439
##########
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:
Thanks, fixed in 00f72e9a: map key types are now never narrowed by this
path, wherever the requirement comes from (reads above the generator, a higher
generator, or an input that is not a field path). Your query is a new test
checking `(4, [23, 23])` and the filter's row count.
A correction to my earlier reply: master does narrow map keys in ordinary
nested pruning. On the 2026-10-29 nightly, `SELECT map_keys(m).key FROM t` (no
generator) reads `m: map<struct<key:int>, string>` and returns `[23]` instead
of `[23, 23]` with the ORC row reader on V1 and V2 (the vectorized reader
returns both). My first check also selected `size(map_keys(m))`, which read the
map whole. That one is a separate master issue.
##########
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:
Agreed. In 00f72e9a this PR no longer prunes ORC (V1 or V2) through
generators, so the two-field `explode` reads as on the base until the ORC
reader fix lands. New test "ORC queries through generators read as before"
covers your timestamp file in both reader modes; the ORC suites also check that
every generator query scans the same as with this path disabled. I'll file the
reader fix (match nested fields by name before the compatibility check)
separately.
##########
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:
Thanks for the inferred case. Since the positional choice is made per file
at read time, the guard in 00f72e9a covers all ORC scans (V1 and V2) rather
than the option alone, so both the explicit and the inferred positional reads
keep the base behavior. The new ORC test includes your `_col0` file in both
reader modes and expects `(22, 33)`. Enabling ORC can follow once the
positional reader handles narrowed nested schemas.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -1015,8 +1016,58 @@ 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(segments, 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(segments.last, 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 `scan`, the projects and filters over `sHolder` below a chain of
generators, from a
+ * fresh `ScanBuilder` of the same relation, like [[rebuildScan]]: the
filters pushed to
+ * `sHolder` and the pushable ones of `scan` are offered to the new builder
together, directly
+ * over the relation. Only valid when `sHolder` has no pushdown such a
rebuild cannot reproduce.
+ */
+ private def rebuildWithoutGenerators(
+ scan: SchemaPruning.Segment,
+ sHolder: ScanBuilderHolder,
+ isSubquery: Boolean): LogicalPlan = {
+ val relation = sHolder.relation
+ val pushable = (sHolder.pushedFilterExpressions ++
scan.filtersPushDown).reduceOption(And)
+ .map(Filter(_, relation)).getOrElse(relation)
+ val filtered = scan.filtersStayUp.reduceOption(And)
Review Comment:
Fixed in 00f72e9a as you verified: the filters that stay up are folded over
the pushable ones in order, one `Filter` per condition.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/SchemaPruningThroughGenerateV2BuilderSuite.scala:
##########
@@ -0,0 +1,157 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.v2
+
+import java.util
+import java.util.concurrent.atomic.AtomicInteger
+
+import scala.collection.mutable.ArrayBuffer
+
+import org.apache.spark.sql.{QueryTest, Row}
+import org.apache.spark.sql.connector.catalog.{SupportsRead, Table,
TableCapability}
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownRequiredColumns}
+import
org.apache.spark.sql.execution.datasources.v2.parquet.{ParquetDataSourceV2,
ParquetScan}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{ArrayType, DataType, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Pruning nested fields through generators on a DataSource V2 scan whose
builder only narrows:
+ * a later `pruneColumns` call cannot widen what an earlier one kept, as in
Iceberg.
+ */
+class SchemaPruningThroughGenerateV2BuilderSuite extends QueryTest with
SharedSparkSession {
+
+ test("SPARK-47230: a rejected rewrite does not prune a narrow-only builder
again") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ sql(
+ """select id, array(
+ | named_struct('a', id, 'b', id * 2, 'c', id * 3, 's', cast(id * 4
as int)),
+ | named_struct('a', id * 10, 'b', id * 20, 'c', id * 30, 's',
cast(id * 40 as int))
+ |) as items from range(1, 3)""".stripMargin).write.parquet(path)
+ val query = "select item.a, item.c from t lateral view explode(items) e
as item"
+ // V1 Parquet as the baseline
+ var expected: Seq[Row] = Nil
+ withTempView("t") {
+ spark.read.parquet(path).createOrReplaceTempView("t")
+ expected = sql(query).collect().toSeq
+ }
+ assert(expected.size == 4)
+
+ Seq(false, true).foreach { reject =>
+ withSQLConf(
+
SQLConf.NESTED_SCHEMA_PRUNING_THROUGH_GENERATE_REJECT_FOR_TESTING.key ->
+ reject.toString) {
+ NarrowOnlyParquetSource.pruneCalls.clear()
+ withTempView("t") {
+
spark.read.format(classOf[NarrowOnlyParquetSource].getName).load(path)
+ .createOrReplaceTempView("t")
+ checkAnswer(sql(query), expected)
+ }
+ val calls = NarrowOnlyParquetSource.pruneCalls.map(_.get).toSeq
+ assert(calls.nonEmpty)
+ assert(calls.forall(_ <= 1), s"pruneColumns calls per builder:
$calls")
+ // a rejected rewrite plans the scan from a second, fresh builder
+ if (reject) assert(calls.count(_ == 1) >= 2, s"pruneColumns calls
per builder: $calls")
+ }
+ }
+ }
+ }
+
+ test("SPARK-47230: a rejected rewrite offers the data filters to the fresh
builder") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ sql("select id as n, array(named_struct('a', id, 'b', id * 2, 'c', id *
3)) as items " +
+ "from range(10)").write.parquet(path)
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") {
+ Seq(false, true).foreach { reject =>
+ withSQLConf(
+
SQLConf.NESTED_SCHEMA_PRUNING_THROUGH_GENERATE_REJECT_FOR_TESTING.key ->
+ reject.toString) {
+ withTempView("t") {
+ spark.read.parquet(path).createOrReplaceTempView("t")
+ val query = sql("select item.a, item.c from t " +
+ "lateral view outer explode(items) e as item where n = 7")
+ val pushed = query.queryExecution.optimizedPlan.collect {
+ case r: DataSourceV2ScanRelation => r.scan
+ }.collect { case s: ParquetScan => s.pushedFilters.toSeq
}.flatten
+ assert(pushed.exists(_.references.contains("n")),
s"reject=$reject: $pushed")
Review Comment:
Done in 00f72e9a: the fixture is partitioned by `p`, the query filters `p =
1 AND n = 7`, and the assertion requires `EqualTo(n, 7)` itself in the pushed
filters, with and without forced rejection.
--
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]
