IgorBerman commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4360850197
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala:
##########
@@ -51,23 +69,223 @@ object SchemaPruning extends Rule[LogicalPlan] {
}).getOrElse(op)
}
+ /** The projects and filters [[ScanOperation]] collects between two
operators. */
+ private[sql] case class Segment(
+ projects: Seq[NamedExpression],
+ filtersStayUp: Seq[Expression],
+ filtersPushDown: Seq[Expression]) {
+ def filters: Seq[Expression] = filtersPushDown ++ filtersStayUp
+ }
+
+ /**
+ * Matches `explode` / `posexplode` generators stacked over a leaf, each one
optionally under
+ * `Project`s and `Filter`s. Returns the segments from the top (the one
above the top generator
+ * first, the one over the leaf last), the generators from the top, and the
leaf.
+ */
+ private[sql] object GeneratorsOverScan {
+ def unapply(plan: LogicalPlan): Option[(Seq[Segment], Seq[Generate],
LogicalPlan)] = {
Review Comment:
perf: early-exit the matcher on generate-free plans
> **File:**
`sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala`,
`GeneratorsOverScan.unapply`
`GeneratorsOverScan.unapply` runs `ScanOperation` collapsing (alias
inlining, allocations) at
every node of every plan, and V1 `SchemaPruning` appears three times in
`earlyScanPushDownRules`, plus the V2 `pruneColumns` call — even when
`throughGenerate.enabled` is false, because the conf is only checked inside
`pruneThroughGenerators` after the match. The usual idiom is an O(1) bitset
test at the top:
```scala
def unapply(plan: LogicalPlan): Option[...] = {
if (!plan.containsPattern(GENERATE)) return None
...
}
```
Related: after a successful rewrite, `transformDown` re-matches every
sub-chain node and re-runs
the full analysis. This is semantically safe (a bare element at a sub-chain
root reads whole, so
`countLeaves` finds nothing more to prune and the schema guard protects the
rest) but is O(n^2)
work per chain; an early exit when the leaf is already pruned would avoid it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala:
##########
@@ -51,23 +69,223 @@ object SchemaPruning extends Rule[LogicalPlan] {
}).getOrElse(op)
}
+ /** The projects and filters [[ScanOperation]] collects between two
operators. */
+ private[sql] case class Segment(
+ projects: Seq[NamedExpression],
+ filtersStayUp: Seq[Expression],
+ filtersPushDown: Seq[Expression]) {
+ def filters: Seq[Expression] = filtersPushDown ++ filtersStayUp
+ }
+
+ /**
+ * Matches `explode` / `posexplode` generators stacked over a leaf, each one
optionally under
+ * `Project`s and `Filter`s. Returns the segments from the top (the one
above the top generator
+ * first, the one over the leaf last), the generators from the top, and the
leaf.
+ */
+ private[sql] object GeneratorsOverScan {
+ def unapply(plan: LogicalPlan): Option[(Seq[Segment], Seq[Generate],
LogicalPlan)] = {
+ @tailrec
+ def collect(p: LogicalPlan, segments: Seq[Segment], generators:
Seq[Generate])
+ : Option[(Seq[Segment], Seq[Generate], LogicalPlan)] = p match {
+ case ScanOperation(projects, filtersStayUp, filtersPushDown,
+ g @ Generate(_: ExplodeBase, _, _, _, _, _)) =>
+ collect(g.child, segments :+ Segment(projects, filtersStayUp,
filtersPushDown),
+ generators :+ g)
+ case ScanOperation(projects, filtersStayUp, filtersPushDown, leaf:
LeafNode)
+ if generators.nonEmpty =>
+ Some((segments :+ Segment(projects, filtersStayUp, filtersPushDown),
generators, leaf))
+ case _ => None
+ }
+ collect(plan, Nil, Nil)
+ }
+ }
+
+ /**
+ * Prunes the scan below a chain of `explode` / `posexplode` generators of
arrays to the fields
+ * that the plan above each generator, up to `op`, reads from its element,
then re-resolves the
+ * chain against the narrower types. Going down from the top generator, the
fields read from a
+ * generator's element become fields read from its input: a column of the
scan, or the element
+ * of a generator below it.
+ *
+ * `pruneScan` prunes the scan given the projects over the scan that do not
only feed the
+ * bottom generator and the root fields the generators need, and returns the
new plan over the
+ * scan. Returns None if nothing is pruned, or if the output of `op` would
change.
+ */
+ private[sql] def pruneThroughGenerators(
+ op: LogicalPlan,
+ scanRelation: LogicalPlan)(
+ pruneScan: (Segment, Seq[NamedExpression], Seq[RootField]) =>
Option[LogicalPlan])
+ : Option[LogicalPlan] = {
+ if (!conf.nestedSchemaPruningEnabled ||
!conf.nestedSchemaPruningThroughGenerateEnabled) {
+ None
+ } else {
+ // e.g. the `size(arr) > 0` filters inferred for a non-outer explode
+ op.transformAllExpressions {
Review Comment:
question: the `size()` rewrite persists in the final plan
When pruning succeeds, the `lengthOnly` rewrite persists in all
matched-chain expressions,
including user-visible projections: `SELECT size(items)` shows as
`size(items.qty)` in explain
output. It is value-equivalent (array length preserved by
`GetArrayStructFields`, null arrays
give -1/NULL identically, under ANSI too) and consistent with
`RewriteSizeOfArrayStruct`, but
worth stating as intended in the PR description.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -990,7 +991,26 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
}
def pruneColumns(plan: LogicalPlan): LogicalPlan = plan.transform {
+ case op @ SchemaPruning.GeneratorsOverScan(_, _, sHolder:
ScanBuilderHolder) =>
+ // Fields read from the generators' elements above them narrow the scan
below them. If the
+ // plan above cannot be re-resolved, the scan is pruned again on its own
below, which calls
+ // `pruneColumns` on the scan builder a second time.
Review Comment:
In the generator case, the connector's `ScanBuilder` is mutated
(`pruneColumns` + `build()`)
*before* the rebuilt plan is validated. When the guard rejects the rewrite,
`transformDown`
descends to the plain `ScanOperation` case and both calls repeat with a
wider schema. The
`SupportsPushDownRequiredColumns` javadoc doesn't define repeated-call
semantics; connectors have
historically seen exactly one prune+build cycle. A builder that intersects
successive schemas, or
caches its first `Scan`, would keep the narrow schema while the plan above
expects the wide one.
The in-code comment acknowledges the double call — but is "last call wins"
an assumption we can
state on the interface (javadoc note on `SupportsPushDownRequiredColumns`),
or should the
generator path avoid mutating the builder until the rewrite is validated? At
minimum it seems
worth a sentence in the PR description. (V1 is fine: nothing is mutated
until a validated
replacement exists.)
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningThroughGenerateSuite.scala:
##########
@@ -0,0 +1,883 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
Review Comment:
The new suite pins ANSI off (fixtures rely on `size(null) = -1`). Since the
`lengthOnly`
rewrite must also be value-preserving under ANSI (`size(null)` is NULL on
both sides), one
ANSI-on smoke test of a non-outer chain would pin that.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala:
##########
@@ -51,23 +69,223 @@ object SchemaPruning extends Rule[LogicalPlan] {
}).getOrElse(op)
}
+ /** The projects and filters [[ScanOperation]] collects between two
operators. */
+ private[sql] case class Segment(
+ projects: Seq[NamedExpression],
+ filtersStayUp: Seq[Expression],
+ filtersPushDown: Seq[Expression]) {
+ def filters: Seq[Expression] = filtersPushDown ++ filtersStayUp
+ }
+
+ /**
+ * Matches `explode` / `posexplode` generators stacked over a leaf, each one
optionally under
+ * `Project`s and `Filter`s. Returns the segments from the top (the one
above the top generator
+ * first, the one over the leaf last), the generators from the top, and the
leaf.
+ */
+ private[sql] object GeneratorsOverScan {
+ def unapply(plan: LogicalPlan): Option[(Seq[Segment], Seq[Generate],
LogicalPlan)] = {
+ @tailrec
+ def collect(p: LogicalPlan, segments: Seq[Segment], generators:
Seq[Generate])
+ : Option[(Seq[Segment], Seq[Generate], LogicalPlan)] = p match {
+ case ScanOperation(projects, filtersStayUp, filtersPushDown,
+ g @ Generate(_: ExplodeBase, _, _, _, _, _)) =>
+ collect(g.child, segments :+ Segment(projects, filtersStayUp,
filtersPushDown),
+ generators :+ g)
+ case ScanOperation(projects, filtersStayUp, filtersPushDown, leaf:
LeafNode)
+ if generators.nonEmpty =>
+ Some((segments :+ Segment(projects, filtersStayUp, filtersPushDown),
generators, leaf))
+ case _ => None
+ }
+ collect(plan, Nil, Nil)
+ }
+ }
+
+ /**
+ * Prunes the scan below a chain of `explode` / `posexplode` generators of
arrays to the fields
+ * that the plan above each generator, up to `op`, reads from its element,
then re-resolves the
+ * chain against the narrower types. Going down from the top generator, the
fields read from a
+ * generator's element become fields read from its input: a column of the
scan, or the element
+ * of a generator below it.
+ *
+ * `pruneScan` prunes the scan given the projects over the scan that do not
only feed the
+ * bottom generator and the root fields the generators need, and returns the
new plan over the
+ * scan. Returns None if nothing is pruned, or if the output of `op` would
change.
+ */
+ private[sql] def pruneThroughGenerators(
+ op: LogicalPlan,
+ scanRelation: LogicalPlan)(
+ pruneScan: (Segment, Seq[NamedExpression], Seq[RootField]) =>
Option[LogicalPlan])
+ : Option[LogicalPlan] = {
+ if (!conf.nestedSchemaPruningEnabled ||
!conf.nestedSchemaPruningThroughGenerateEnabled) {
+ None
+ } else {
+ // e.g. the `size(arr) > 0` filters inferred for a non-outer explode
+ op.transformAllExpressions {
+ case s @ Size(child, _) => s.copy(child = lengthOnly(child))
+ } match {
+ case sized @ GeneratorsOverScan(segments, generators, _) =>
+ pruneThroughGeneratorsImpl(sized, segments, generators,
scanRelation)(pruneScan)
+ case _ => None
+ }
+ }
+ }
+
+ /**
+ * The size of an array only needs its length. Returns an array of the same
length that reads
+ * as little as possible: the smallest field of an array of structs, like
+ * [[org.apache.spark.sql.catalyst.optimizer.RewriteSizeOfArrayStruct]] does
for scan columns,
+ * and of the outer array for a nested field of an array of structs, which
has the same length.
+ */
+ private def lengthOnly(array: Expression): Expression = {
+ def smallestField(array: Expression): Option[Expression] = array.dataType
match {
+ case ArrayType(StructType(fields), containsNull) if fields.length > 1 =>
+ val (field, ordinal) =
fields.zipWithIndex.minBy(_._1.dataType.defaultSize)
+ Some(GetArrayStructFields(
+ array, field, ordinal, fields.length, containsNull ||
field.nullable))
+ case _ => None
+ }
+ array match {
+ // a field that is itself nested, e.g. `items.parts`: read a field of
`items` instead
+ case g @ GetArrayStructFields(child, field, _, _, _)
+ if !field.dataType.isInstanceOf[AtomicType] =>
+ smallestField(child).orElse(smallestField(g)).getOrElse(g)
+ case _ => smallestField(array).getOrElse(array)
+ }
+ }
+
+ private def pruneThroughGeneratorsImpl(
+ op: LogicalPlan,
+ segments: Seq[Segment],
+ generators: Seq[Generate],
+ scanRelation: LogicalPlan)(
+ pruneScan: (Segment, Seq[NamedExpression], Seq[RootField]) =>
Option[LogicalPlan])
+ : Option[LogicalPlan] = {
+ val attrNameMap = scanRelation.output.map(att => (att.exprId,
att.name)).toMap
+ // Each generator's input in terms of the segment below it, and the
projects of that segment
+ // which only feed the generator: their reads are the generator's input
root fields instead.
+ val inputs = generators.zip(segments.tail).map { case (g, below) =>
+ val aliases = below.projects.collect { case a: Alias => a.exprId ->
a.child }.toMap
+ val input = g.generator.children.head.transformUp {
+ case a: Attribute if aliases.contains(a.exprId) => aliases(a.exprId)
+ }
+ val feeding = g.generator.references -- g.requiredChildOutput
+ (normalizeAttributeRefNames(attrNameMap, Seq(input)).head,
+ below.projects.filterNot(p => feeding.contains(p.toAttribute)))
+ }
+ val elements = generators.map(_.generatorOutput.last)
+ val lowerElements = elements.map(_.exprId).toSet
+ // Reads above each generator. Columns that a lower segment only passes up
are read where
+ // they are used, so only the top segment counts them as read whole.
+ val readsAbove =
segments.indices.dropRight(1).scanLeft(Seq.empty[Expression]) { (reads, i) =>
+ val projects = if (i == 0) segments(i).projects else inputs(i -
1)._2.filter {
+ case _: Attribute => false
+ case _ => true
+ }
+ reads ++ projects ++ segments(i).filters
+ }.tail
+ val elementFields = mutable.Map.empty[ExprId,
Seq[StructField]].withDefaultValue(Nil)
+ val scanRootFields = mutable.ArrayBuffer.empty[RootField]
+ // Inputs that are not a nested field access path, e.g.
`explode(array(item.sku))`, are read
+ // by the generators below like any other expression above them.
+ val inputReads = mutable.ArrayBuffer.empty[Expression]
+ generators.indices.foreach { i =>
+ val (input, _) = inputs(i)
+ val rootField = getGeneratorInputRootField(
+ input, elements(i), readsAbove(i) ++ inputReads,
elementFields(elements(i).exprId))
+ (rootField, input.references.toSeq) match {
+ case (Some(field), Seq(root)) if lowerElements.contains(root.exprId) =>
+ elementFields(root.exprId) :+= field.field
+ case (Some(field), Seq(root)) if scanRelation.outputSet.contains(root)
=>
+ scanRootFields += field
+ case _ =>
+ inputReads += input
+ scanRootFields ++= scanRootFieldsOf(input, scanRelation,
lowerElements)
+ }
+ }
+ pruneScan(segments.last, inputs.last._2, scanRootFields.toSeq)
Review Comment:
The final guard `.filter(newOp => newOp.resolved && newOp.schema ==
op.schema)` only protects
rebuilds that *complete*. If the read collection ever misses an access,
`ProjectionOverSchema`
throws before the guard runs (`StructType.fieldIndex` → `FIELD_NOT_FOUND`;
`SparkException.internalError` for `GetArrayStructFields`), so a collection
bug becomes a
query-killing error instead of a fallback to the unpruned plan — the crash
in the previous
comment is one instance (there it throws during collection, in
`pruneSchema`'s merge).
Since the bookkeeping here (`readsAbove` / `elementFields` / `inputReads`
across segments) is
substantially more intricate than the existing scan-local pruning, consider
wrapping the
collect-prune-rebuild in
```scala
try { ... } catch { case NonFatal(e) => logWarning(..., e); None }
```
so any residual bug degrades to a missed optimization. There is precedent
nearby:
`V2ScanRelationPushDown.pruneScan` catches `FIELD_NOT_FOUND` when remapping
inferred filters.
This would also strengthen the case for shipping the flag default-true.
One residual blind spot worth a code comment in `rebuildAbove`: a missed
rewrite whose stale
ordinal happens to land on a *same-typed sibling* field passes both
`resolved` and the schema
check and silently reads the wrong column. The per-exprId `withName`
normalization is what
prevents the one known instance (references written in another case), so the
"keep
`SelectedField` and `ProjectionOverSchema` access paths in sync" invariant
now extends through
generator elements and deserves to be written down.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruning.scala:
##########
@@ -133,6 +134,53 @@ 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. Returns None if `input`
is not an access
+ * path [[SelectedField]] can narrow.
+ */
+ def getGeneratorInputRootField(
Review Comment:
follow-up optimization: length-only reads could reuse an already-required
field
> **File:** `sql/catalyst/.../expressions/SchemaPruning.scala`,
`getGeneratorInputRootField` /
> `sql/core/.../datasources/SchemaPruning.scala`, `lengthOnly`
When only the array length is needed, the smallest field is kept even when
another field of the
same struct is already required, so the scan reads one extra column. Visible
in the new suite:
in "a filter on the exploded array column itself" the scan reads
`struct<items:array<struct<qty:int,sku:string>>>` where `sku` alone would
determine `size(items)`.
Choosing the smallest *already-required* field when one exists would close
this. Fine as a
follow-up; maybe worth a `// TODO` where the smallest field is chosen.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningThroughGenerateSuite.scala:
##########
@@ -0,0 +1,883 @@
+/*
Review Comment:
PR description: mention the query-shape suite in "How was this patch tested?"
The description's test section covers `SchemaPruningSuite` and the benchmark
but not
`SchemaPruningThroughGenerateSuite` — the largest file in the PR (58 query
shapes x V1/V2, each
checking the exact scan schema *and* comparing results against a flag-off
baseline). Reviewers
read that section closely; a paragraph on it would strengthen the PR.
--
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]
