IgorBerman commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4360761797
##########
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)
+ .map(rebuildAbove(op, generators.last.child, _))
+ .filter(newOp => newOp.resolved && newOp.schema == op.schema)
+ }
+
+ /** The root fields of scan columns that `expr` reads, ignoring generator
elements. */
+ private def scanRootFieldsOf(
Review Comment:
In `scanRootFieldsOf`, only generator *elements* are renamed before
`identifyRootFields`; the
`pos` output of `posexplode` (and the `key` of a map explode) keeps its
user-chosen name, and the
final filter keeps any root field whose *name* matches a scan column. When a
position output
shares its name with a real column and reaches the non-path-input fallback
(`inputReads` branch), this manufactures a bogus root field for the real
column, and
`pruneSchema`'s `reduceLeft(_ merge _)` throws on the type conflict.
Repro (succeeds with `throughGenerate.enabled=false`, crashes with it on):
```scala
case class P(a: Int, b: Int)
case class S(x: Int, y: Int)
case class Ord(pos: P, items: Array[S])
Seq(Ord(P(1, 2), Array(S(10, 20), S(30, 40)))).toDF
.write.parquet(path)
spark.read.parquet(path).createOrReplaceTempView("orders")
sql("""select orders.pos.a, t2.y
from orders
lateral view posexplode(items) t1 as pos, item
lateral view explode(array(t1.pos + item.x)) t2 as y""").collect()
```
```
org.apache.spark.SparkException: [CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE]
Failed to merge incompatible data types "STRUCT<a: INT>" and "INT".
SQLSTATE: 42825
```
If the real `pos` column is *not* read by the query, the pruned relation's
`dataSchema` silently
carries `pos:int` instead of the struct (latent, since the scan drops
unreferenced columns, but
the relation schema is wrong).
Suggested fix: rename every attribute that is not a scan output, not only
the elements:
```scala
val renamed = expr.transform {
case a: Attribute if !scanRelation.outputSet.contains(a) =>
a.withName(s"${a.name}#${a.exprId.id}")
}
```
(With the renamed control — `t1 as idx, item` — the same query prunes
beautifully:
`struct<pos:struct<a:int>,items:array<struct<x:int>>>`.)
---
--
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]
