ohad-definity commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4363112007
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala:
##########
@@ -51,23 +71,241 @@ 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 {
Review Comment:
Matches a stack of `explode`/`posexplode` over a scan, collecting the
`Project`/`Filter` segment above each generator with `ScanOperation`. Anything
else in between (a join, an aggregate, an `Expand`) ends the match, so those
shapes keep today's behavior; the query-shape suite pins them as "not pruned
through yet".
##########
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) =>
Review Comment:
The V2 path reuses the same matcher and requirements; only the scan pruning
(`pruneScan`, previously inlined in this case) differs. See "Notes for
reviewers" in the description about `pruneColumns` being called again if the
rewrite is rejected.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala:
##########
@@ -51,23 +71,241 @@ 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) =>
Review Comment:
The core of the top-down pass: for generator i, the reads are every
expression above it, except bare columns that an inner segment only passes up
(they are read where they are used) and the projections that only feed the
generator below (their reads become that generator's requirement instead).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala:
##########
@@ -51,23 +71,241 @@ 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))
+ (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)
+ }
+ }
+ pruneScan(segments.last, inputs.last._2, scanRootFields.toSeq)
+ .map(rebuildAbove(op, generators.last.child, _))
+ .filter(newOp => newOp.resolved && newOp.schema == op.schema)
+ }
+
+ /** The root fields of scan columns that `expr` reads. */
+ private def scanRootFieldsOf(
+ expr: Expression,
+ scanRelation: LogicalPlan): Seq[RootField] = {
+ // Give every other attribute, e.g. a generator's element or the position
of a posexplode,
+ // a name no scan column has, so that it is not mistaken for a column of
its name.
+ val renamed = expr.transform {
+ case a: Attribute if !scanRelation.outputSet.contains(a) =>
+ a.withName(s"${a.name}#${a.exprId.id}")
+ }
+ val scanColumns = scanRelation.output.map(_.name).toSet
+ identifyRootFields(Nil, Seq(renamed)).filter(f =>
scanColumns.contains(f.field.name))
+ }
+
+ /**
+ * Replaces `bottom` in `plan` with `newBottom`, whose output types are
narrower, and
+ * re-resolves every operator from there up: references take their child's
new type, nested
+ * field accesses are re-resolved by name and generator outputs take the new
element type.
+ *
+ * Every access collected through [[SelectedField]] must be rewritten here by
+ * [[ProjectionOverSchema]]. A missed rewrite usually fails resolution or
the output schema
+ * check, but a stale ordinal that lands on a sibling field of the same type
would not, which
+ * is why references are matched by exprId rather than by name.
+ */
+ private def rebuildAbove(
Review Comment:
One attribute at a time with `ProjectionOverSchema`, because two generators'
outputs can share a name (e.g. two `col`s). References are matched by exprId
and renamed first, otherwise a reference written in another case keeps its
stale ordinal.
##########
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:
`SelectedField.withDataType` (SPARK-57022) does the work of expressing "this
input, narrowed to this element type" as a root field, including through
`items.parts`-style paths, so no new access-path code is needed here.
--
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]
