sunchao commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4417133805
##########
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:
**Preserve separate filters when replaying `filtersStayUp`.**
`ScanOperation` keeps these filters separate when they cannot safely be
combined. A filter discards a row when its condition is NULL. With `NULL AND
rhs`, Spark evaluates `rhs`, so this rewrite can evaluate a cast the original
plan skipped.
I reproduced this on V2 Parquet with ANSI enabled. For `range(10)`, set
`(k,s)` to `(true,'1')` on `id=7` and `(NULL,'bad')` otherwise. Apply
`.filter("rand(3)>0").filter("k OR rand(7)<0.000001").filter("cast(s as
int)>0")` below an outer `explode`. Feature-off and accepted-rewrite plans
return the row for `id=7`. With `throughGenerate.rejectForTesting=true`, the
fallback throws `CAST_INVALID_INPUT` on `'bad'`. The previous head returned the
row in all three cases. I verified that folding `scan.filtersStayUp` over
`pushable` in order, building one `Filter` per condition, restores the expected
result. This reproducer uses forced rejection.
##########
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:
I verified that the original `posexplode(map_keys(m))` now returns both
positions. There is still a row-loss case when `map_keys` is evaluated above
the generator. With the ORC row reader, put a map with struct keys `(id,b,z) =
(1,23,17)` and `(2,23,17)` inside the single element of `items`, alongside `idx
= 4`. The query is
```sql
SELECT e.idx, map_keys(e.m).b
FROM t LATERAL VIEW OUTER explode(items) g AS e
```
The base and the head with generator pruning off return `(4,[23,23])`. On
`d9120db1` with pruning on, both V1 and V2 return `(4,[23])`. Filtering on
`array_size(filter(map_keys(e.m).b, b -> b = 23)) > 1` also changes one
matching row into zero. Here the generator input is just `items`, so
`input.exists(MapKeys)` never sees the map-key access in the projection or
filter. Please preserve the full key type when propagating these requirements
to the scan, including when they come from a higher generator. An assertion on
this filter’s row count would cover it.
##########
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:
I confirmed your master diagnosis. On base `4b6955eb` ordinary `SELECT
items.b,items.c` already hits the ordinal mismatch when nested pruning is
enabled. The two-field `explode` query on that same base succeeds in both
reader modes. On `d9120db1` it succeeds with generator pruning disabled and
throws `ORC_TYPE_CAST` with it enabled. I’d favor landing the ORC reader fix
first, as you offered. Guarding the new path for affected V1 scans until that
fix lands would also preserve the existing query behavior.
##########
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:
**Exercise the predicate combination that broke replay.** I ran this
unpartitioned fixture against `dec83e8fe4`, before the fix. Even with forced
rejection it already pushed `EqualTo(n,7)` and returned `[7,21]`, so the new
assertion would pass if the production fix were reverted. The original failure
needs both a fully handled partition predicate and a partially pushed data
predicate. Please partition the table by `p`, filter on `p = 1 AND n = 7`, and
check for the actual `EqualTo(n,7)` predicate. The current assertion also
accepts `IsNotNull(n)` alone. On the partitioned fixture I confirmed the
previous head drops the equality on rejection and the current head retains it.
##########
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:
Agreed that the positional reader problem predates the PR. The two-field
`explode` still returns `(22,33)` on base `4b6955eb` and `(11,22)` on
`d9120db1` with generator pruning enabled. I also reproduced this on V2 without
setting `orc.force.positional.evolution`. Write `SELECT
array(named_struct('a',11,'b',22,'c',33)) AS _col0` to ORC, read it with
`spark.read.orc(path)`, and run `SELECT e.b,e.c FROM t LATERAL VIEW OUTER
explode(_col0) g AS e`. Both reader modes silently read `(a,b)`. ORC selects
positional evolution for files with these names. Please land the reader fix
first or guard the new pruning path for both explicit and inferred positional
reads. Checking only the user option would miss this file.
--
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]
