ohad-definity commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4414286120
##########
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:
Thanks for measuring this, done in b5869812 along your points: the generator
case is skipped when `hasBlockingPushdown(sHolder)`; if an attempt that used
the builder is rejected, the projects and filters below the generators are
planned right there from a fresh builder of the same relation (re-applying
`pushedFilterExpressions`, like `rebuildScan`), so `transformDown` finds no
`ScanBuilderHolder` left for the lower generators; and `isSubquery` is now
passed to `pruneColumns` for the fresh builder's `ScanContext`. New
`SchemaPruningThroughGenerateV2BuilderSuite` reproduces your setup: a Parquet
V2 builder that only narrows, your query, with and without an internal
test-only flag that rejects every rewrite after the scan is pruned. Rows match
V1 Parquet, and every builder sees at most one `pruneColumns`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala:
##########
@@ -51,23 +71,261 @@ 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)] = {
+ if (!plan.containsPattern(GENERATE) || !conf.nestedSchemaPruningEnabled
||
+ !conf.nestedSchemaPruningThroughGenerateEnabled) {
+ return None
+ }
+ @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 {
+ try {
+ // 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
+ }
+ } catch {
+ // A missed optimization, not a failed query: keep the plan as it is.
+ case NonFatal(e) =>
+ logWarning(log"Failed to prune nested fields through generators; " +
+ log"the scan is pruned without them", e)
+ 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))
+ // the input also reads what a filter predicate or sort comparator
inside it reads
+ val inputFields = rootField.map(_ +:
getGeneratorInputLambdaRootFields(input))
+ (inputFields, input.references.toSeq) match {
+ case (Some(fields), Seq(root)) if lowerElements.contains(root.exprId)
=>
+ elementFields(root.exprId) ++= fields.map(_.field)
+ case (Some(fields), Seq(root)) if
scanRelation.outputSet.contains(root) =>
+ scanRootFields ++= fields
+ case _ =>
+ inputReads += input
+ scanRootFields ++= scanRootFieldsOf(input, scanRelation)
+ }
+ }
+ pruneScan(segments.last, inputs.last._2, scanRootFields.toSeq)
Review Comment:
Pinned it for now, as you suggested: "an explode of a second scan column is
not pruned through yet" in `SchemaPruningSuite` (your shape, reading `a2`
whole), and added to the follow-ups in the description. Counting bottom-segment
pass-through columns where they are used is the right fix, but it touches how
whole reads of scan columns reach `scanRootFields`, so I would rather do it as
a follow-up.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruning.scala:
##########
@@ -133,6 +134,61 @@ 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(
+ input: Expression,
+ element: Attribute,
+ exprs: Seq[Expression],
+ elementFields: Seq[StructField] = Nil): Option[RootField] =
input.dataType match {
+ case ArrayType(elementType, containsNull) =>
+ // Rename the element so that its fields can be told apart from other
columns by name.
+ val name = s"${element.name}#${element.exprId.id}"
+ val renamed = exprs.map(_.transform {
+ case a: Attribute if a.exprId == element.exprId => a.withName(name)
+ })
Review Comment:
Applied your suggestion in fd44a39d, thanks. Your query is a new test in
`SchemaPruningSuite`; `items` now reads `<x, y>` (the `pos` column itself
passes through the posexplode and is read whole, as on master).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala:
##########
@@ -51,23 +71,261 @@ 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)] = {
+ if (!plan.containsPattern(GENERATE) || !conf.nestedSchemaPruningEnabled
||
+ !conf.nestedSchemaPruningThroughGenerateEnabled) {
+ return None
+ }
+ @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 {
+ try {
+ // 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
+ }
+ } catch {
+ // A missed optimization, not a failed query: keep the plan as it is.
+ case NonFatal(e) =>
+ logWarning(log"Failed to prune nested fields through generators; " +
Review Comment:
Done in fd44a39d: it rethrows when `Utils.isTesting`, and otherwise logs at
debug level without the warning. With the fresh-builder fallback there are also
no more retries at the lower chain nodes after a V2 rejection.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruning.scala:
##########
@@ -133,6 +134,61 @@ 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.
Review Comment:
Applied in fd44a39d.
--
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]
