peter-toth commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4411985187
##########
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:
**Finding 2.** A generator above the bottom one whose input is a separate
scan column never narrows that column. The narrowed root field does reach
`scanRootFields`, through the `scanRelation.outputSet.contains(root)` case
above. But the column also passes up through the bottom segment as a bare
attribute, because it is in the lower `Generate`'s `requiredChildOutput`. So
`inputs.last._2` keeps it, `identifyRootFields` reads it whole, and
`pruneSchema`'s merge widens it back.
I ran this on Parquet V1 and V2, with `a1: array<struct<p:int,q:string>>`
and `a2: array<struct<x:int,y:int,z:string>>`:
```sql
SELECT b.x, b.y FROM t LATERAL VIEW explode(a1) t1 AS a LATERAL VIEW
explode(a2) t2 AS b
```
With the flag on it reads
`a1:array<struct<p:int>>,a2:array<struct<x:int,y:int,z:string>>`. With it off,
`a1` is read whole too. The rows match.
The description's first step says a generator's input is "a scan column, or
the element of a generator below". Here that holds only for the bottom
generator. Two ways forward:
- Count a bottom-segment bare attribute as read where it is used, as the
intermediate segments already do. The whole reads of scan columns above then
need another way into `scanRootFields`.
- Or pin this shape as "not pruned through yet" in the query-shape suite and
list it with the other follow-ups.
##########
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:
**Finding 3.** The clash fixed in `scanRootFieldsOf` after
[r4360761797](https://github.com/apache/spark/pull/59706#discussion_r4360761797)
is still here. Only the element is renamed, so `identifyRootFields` can meet
two other attributes of the same name with different types. The optional-root
merge then throws, the `NonFatal` catch drops the pruning for the whole chain,
and a warning is logged.
I ran this on Parquet V1 and V2, with a scan column `pos:
struct<a:int,b:int>` and `items: array<struct<x:int,y:int,z:string>>`:
```sql
SELECT t1.pos, item.x, item.y, p2.pos IS NOT NULL AS f, p2.pos.a AS pa
FROM p2 LATERAL VIEW posexplode(items) t1 AS pos, item
```
- It logs `[CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE] Failed to merge
incompatible data types "STRUCT<a: INT, b: INT>" and "INT"` and reads `items`
whole.
- Without the two `p2.pos` reads, it prunes `items` to `<x, y>`.
Renaming the other attributes too fixes it, as `scanRootFieldsOf` now does
for the non-scan ones. The element's references keep the one canonical name,
which the "references in another case" test needs:
```suggestion
// 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}")
})
```
With it, the query above reads `items:array<struct<x:int,y:int>>` on V1 and
V2 and logs nothing. The 1817 tests of
`Parquet{V1,V2}SchemaPruningThroughGenerateSuite` and the Parquet and ORC V1/V2
`SchemaPruningSuite`s pass.
##########
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:
**Finding 5.** `pruneThroughGeneratorsImpl` also passes inputs over the
element of a generator below, e.g. `item.parts` for the upper generator of two
chained explodes:
```suggestion
* `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.
```
##########
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:
**Finding 4.** Each failed attempt logs a full stack trace, and one failure
repeats. Two things multiply it:
- The rule runs three times in `earlyScanPushDownRules`.
- After a rejection, `transformDown` retries at the lower chain nodes, the
re-matches in
[r4360850197](https://github.com/apache/spark/pull/59706#discussion_r4360850197).
In my run, the query in Finding 3 logged it 7 times on V1 and V2 together.
In tests, the catch also hides the drift between `SelectedField` and
`ProjectionOverSchema` that `rebuildAbove`'s doc warns about. A suite then
passes on the unpruned plan. I suggest:
- Log once, without the throwable, or at debug level.
- Rethrow when `Utils.isTesting`, so the pruning suites fail instead.
##########
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) {
Review Comment:
**Finding 7.** This check can never fail here. `pruneThroughGenerators` is
only reached after `GeneratorsOverScan` matched, and its `unapply` already
returns None when either flag is off. The `sized @ GeneratorsOverScan(...)`
match below then checks them a third time. Keeping the check in `unapply` only
is enough.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningThroughGenerateSuite.scala:
##########
@@ -0,0 +1,1027 @@
+/*
+ * 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
+
+import org.scalactic.Equality
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.{DataFrame, QueryTest, Row}
+import org.apache.spark.sql.catalyst.SchemaPruningTest
+import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.execution.FileSourceScanExec
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.tags.ExtendedSQLTest
+
+/**
+ * End-to-end query-shape coverage for nested schema pruning through
generators (SPARK-47230):
+ * each test runs a SQL query typical of analytics over arrays of structs -
explode chains,
+ * views, CTEs, windows, grouping sets, subqueries, higher-order functions -
and checks both
+ * the scan schema and that the results match the ones produced with
+ * [[SQLConf.NESTED_SCHEMA_PRUNING_THROUGH_GENERATE_ENABLED]] disabled.
+ *
+ * [[SchemaPruningSuite]] covers the mechanics of the pruning rules; this
suite covers the
+ * surrounding query shapes.
+ */
+abstract class SchemaPruningThroughGenerateSuite
+ extends QueryTest
+ with SchemaPruningTest
+ with SharedSparkSession {
+
+ // The first field of each struct is its smallest, which is the one the
`size(...) > 0`
+ // filters inferred for a non-outer explode read
+ case class Info(source: String, score: Double, hits: Array[Int])
+ case class Part(
+ inStock: Boolean, visible: Boolean, pid: Int, label: String, weight:
Double, price: Double,
+ tiers: Array[Int])
+ case class Item(active: Boolean, qty: Int, sku: String, note: String, parts:
Array[Part])
+ case class Order(
+ id: Int, country: String, memo: String, items: Array[Item], tags:
Map[String, Info])
+ case class Dim(id: Int, name: String)
+
+ // two exact-duplicate parts plus one differing only in `label`, and a null
element
+ private val dupParts = Array(
+ Part(true, true, 50, "dupA", 1.0, 1.0, Array(5)),
+ Part(true, true, 50, "dupB", 1.0, 1.0, Array(5)),
+ Part(true, true, 50, "dupA", 1.0, 1.0, Array(5)),
+ null)
+
+ val ordersData: Seq[Order] = Seq(
+ Order(1, "US", "m1", Array(
+ Item(true, 1, "chrome", "n1", Array(
+ Part(true, true, 10, "c10", 1.5, 1.2, Array(1, 2)),
+ Part(false, true, 11, "c11", 2.0, 0.4, Array.empty))),
+ Item(false, 2, "firefox", "n2", Array(
+ Part(true, false, 12, "f12", 0.5, 2.5, Array(3))))),
+ Map("a" -> Info("s1", 0.5, Array(7, 8)), "b" -> Info("s2", 1.5,
Array.empty))),
+ Order(2, "UK", "m2", Array(
+ Item(true, 3, "chrome", "n3", Array.empty)),
+ Map.empty),
+ Order(3, "US", "m3", Array.empty, Map("c" -> Info(null, 2.0, null))),
+ Order(4, "DE", null, null, null),
+ // order 5 makes whole-element semantics observable: items and parts that
differ only in
+ // a field the whole-element-operator tests below never read, so a dedup
or a sort over
+ // wrongly narrowed elements would change their rows
+ Order(5, "FR", "m5", Array(
+ Item(true, 9, "alpha", "n5a", dupParts),
+ Item(true, 9, "alpha", "n5x", dupParts),
+ Item(true, 1, "beta", "n5b", Array.empty)),
+ Map.empty))
+
+ val dimData: Seq[Dim] = Seq(Dim(1, "one"), Dim(2, "two"))
+
+ private val fullPartsType =
+ "array<struct<inStock:boolean,visible:boolean,pid:int,label:string," +
+ "weight:double,price:double,tiers:array<int>>>"
+ private val fullItemsType =
+
s"array<struct<active:boolean,qty:int,sku:string,note:string,parts:$fullPartsType>>"
+ private val fullTagsType =
"map<string,struct<source:string,score:double,hits:array<int>>>"
+
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(SQLConf.ANSI_ENABLED.key, "false")
+ // one deterministic plan per query, so the scans are collected in a
stable order
+ .set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false")
+
+ /**
+ * Runs `testThunk` against an `orders` parquet table of [[ordersData]] and
a `base` temp view
+ * exploding it twice, the common "flattening" view pattern this
optimization targets.
+ */
+ protected def testPruning(testName: String)(testThunk: => Unit): Unit = {
+ test(testName) {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark.createDataFrame(ordersData).write.format("parquet").save(path +
"/orders")
+ spark.createDataFrame(dimData).write.format("parquet").save(path +
"/dim")
+ withTempView("orders", "dim", "base") {
+ spark.read.format("parquet").load(path +
"/orders").createOrReplaceTempView("orders")
+ spark.read.format("parquet").load(path +
"/dim").createOrReplaceTempView("dim")
+ sql(
+ """CREATE OR REPLACE TEMP VIEW base AS
+ |SELECT * FROM orders
+ |LATERAL VIEW EXPLODE(items) t1 AS item
+ |LATERAL VIEW EXPLODE(item.parts) t2 AS part""".stripMargin)
+ testThunk
+ }
+ }
+ }
+ }
+
+ /**
+ * Checks that `sqlText` reads exactly `expectedSchemata` from its scans,
and returns the same
+ * rows as with pruning through generators disabled.
+ */
+ protected def checkPrunedQuery(sqlText: String, expectedSchemata: String*):
Unit = {
+ val key = SQLConf.NESTED_SCHEMA_PRUNING_THROUGH_GENERATE_ENABLED.key
+ val expectedRows = withSQLConf(key -> "false") {
+ sql(sqlText).collect().toSeq
+ }
+ withSQLConf(key -> "true") {
+ val query = sql(sqlText)
+ checkScanSchemata(query, expectedSchemata: _*)
+ checkAnswer(query, expectedRows)
+ }
+ }
+
+ protected val schemaEquality: Equality[StructType] = new
Equality[StructType] {
+ override def areEqual(a: StructType, b: Any): Boolean =
+ b match {
+ case otherType: StructType => DataTypeUtils.sameType(a, otherType)
+ case _ => false
+ }
+ }
+
+ protected def checkScanSchemata(df: DataFrame, expectedSchemaCatalogStrings:
String*): Unit = {
+ val fileSourceScanSchemata = df.queryExecution.executedPlan.collect {
+ case scan: FileSourceScanExec => scan.requiredSchema
+ }
+ assert(fileSourceScanSchemata.size === expectedSchemaCatalogStrings.size,
+ s"Found ${fileSourceScanSchemata.size} file sources in dataframe, " +
+ s"but expected $expectedSchemaCatalogStrings")
+ fileSourceScanSchemata.zip(expectedSchemaCatalogStrings).foreach {
+ case (scanSchema, expectedScanSchemaCatalogString) =>
+ val expectedScanSchema =
CatalystSqlParser.parseDataType(expectedScanSchemaCatalogString)
+ implicit val equality = schemaEquality
+ assert(scanSchema === expectedScanSchema)
+ }
+ }
+
+ testPruning("SPARK-47230: top-level fields only") {
+ checkPrunedQuery(
+ "select id, country from orders",
+ "struct<id:int,country:string>")
+ }
+
+ testPruning("SPARK-47230: double explode, fields from every level") {
+ checkPrunedQuery(
+ """select id, item.sku, part.inStock, part.pid
+ |from orders
+ |lateral view explode(items) t1 as item
+ |lateral view explode(item.parts) t2 as part""".stripMargin,
+ "struct<id:int,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<inStock:boolean,pid:int>>>>>")
+ }
+
+ testPruning("SPARK-47230: through a SELECT * view") {
+ // only `pid` is used from `part`, so GeneratorNestedColumnAliasing
rewrites the inner
+ // explode to explode(item.parts.pid), whose size() filter needs no extra
smallest field
+ checkPrunedQuery(
+ "select id, item.sku, part.pid from base",
+ "struct<id:int,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<pid:int>>>>>")
+ }
+
+ testPruning("SPARK-47230: filters on exploded fields") {
+ checkPrunedQuery(
+ """select id, item.sku, part.pid, part.price from base
+ |where item.active = true and part.inStock = true and part.price >
0""".stripMargin,
+ "struct<id:int,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<inStock:boolean,pid:int,price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: aggregation with order and limit") {
+ checkPrunedQuery(
+ """select country, item.sku, part.pid, count(*) as cnt,
+ | sum(case when part.inStock then 1 else 0 end) as stocked,
+ | avg(part.price) as avg_price
+ |from base
+ |where items is not null
+ |group by country, item.sku, part.pid
+ |order by cnt desc, part.pid
+ |limit 10""".stripMargin,
+ "struct<country:string,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<inStock:boolean,pid:int,price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: only inner-element fields") {
+ checkPrunedQuery(
+ "select part.pid, part.label, part.price from base",
+ "struct<items:array<struct<active:boolean," +
+
"parts:array<struct<inStock:boolean,pid:int,label:string,price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: CTE over the explode chain") {
+ // `id` and `pid` are selected by the CTE but unused above it, so they are
pruned away
+ checkPrunedQuery(
+ """with exploded as (
+ | select id, country, item.sku, part.pid, part.inStock
+ | from orders
+ | lateral view explode(items) t1 as item
+ | lateral view explode(item.parts) t2 as part)
+ |select country, sku, count(*) as total,
+ | sum(case when inStock then 1 else 0 end) as stocked
+ |from exploded
+ |group by country, sku""".stripMargin,
+ "struct<country:string,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<inStock:boolean>>>>>")
+ }
+
+ testPruning("SPARK-47230: outer explode chain") {
+ // no `size(...) > 0` filters are inferred for outer generators, so
nothing beyond the
+ // selected fields is read
+ checkPrunedQuery(
+ """select id, country, item.sku, part.pid
+ |from orders
+ |lateral view outer explode(items) t1 as item
+ |lateral view outer explode(item.parts) t2 as part""".stripMargin,
+ "struct<id:int,country:string,items:array<struct<sku:string," +
+ "parts:array<struct<pid:int>>>>>")
+ }
+
+ testPruning("SPARK-47230: posexplode chain with positions selected") {
+ checkPrunedQuery(
+ """select id, ipos, item.sku, ppos, part.pid, part.inStock
+ |from orders
+ |lateral view posexplode(items) t1 as ipos, item
+ |lateral view posexplode(item.parts) t2 as ppos, part""".stripMargin,
+ "struct<id:int,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<inStock:boolean,pid:int>>>>>")
+ }
+
+ testPruning("SPARK-47230: consecutive outer posexplodes under an aggregating
query") {
+ // the same aggregation over a posexploded CTE, with and without the
position columns
+ // selected: both prune the same way, and to the same results
+ def aggregateOverExploded(positions: String): String =
+ s"""with exploded as (
+ | select country, id, $positions item.sku, item.active,
+ | part.pid, part.inStock, part.price, part.visible
+ | from orders
+ | lateral view outer posexplode(items) t1 as ipos, item
+ | lateral view outer posexplode(item.parts) t2 as ppos, part)
+ |select country, sku, pid, count(*) as cnt,
+ | sum(case when inStock then 1 else 0 end) as stocked,
+ | sum(case when visible then 1 else 0 end) as vis,
+ | avg(price) as avg_price
+ |from exploded
+ |where active = true
+ |group by country, sku, pid
+ |order by cnt desc, pid
+ |limit 20""".stripMargin
+
+ val expectedSchema =
+ "struct<country:string,items:array<struct<active:boolean,sku:string," +
+
"parts:array<struct<inStock:boolean,visible:boolean,pid:int,price:double>>>>>"
+ checkPrunedQuery(aggregateOverExploded(positions = ""), expectedSchema)
+ checkPrunedQuery(aggregateOverExploded(positions = "ipos, ppos,"),
expectedSchema)
+ }
+
+ testPruning("SPARK-47230: window functions over exploded fields") {
+ checkPrunedQuery(
+ """select country, item.sku, part.pid,
+ | row_number() over (partition by country order by part.pid) as rn,
+ | sum(part.price) over (partition by country) as total_price
+ |from orders
+ |lateral view explode(items) t1 as item
+ |lateral view explode(item.parts) t2 as part""".stripMargin,
+ "struct<country:string,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<inStock:boolean,pid:int,price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: lag and lead over exploded fields") {
+ checkPrunedQuery(
+ """select country, part.pid,
+ | lag(part.price, 1) over (partition by country order by part.pid) as
prev_price,
+ | lead(part.price, 1) over (partition by country order by part.pid)
as next_price
+ |from orders
+ |lateral view explode(items) t1 as item
+ |lateral view explode(item.parts) t2 as part""".stripMargin,
+ "struct<country:string,items:array<struct<active:boolean," +
+ "parts:array<struct<inStock:boolean,pid:int,price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: transform lambda narrows the array without an
explode") {
+ checkPrunedQuery(
+ "select id, transform(items, i -> i.sku) as skus from orders",
+ "struct<id:int,items:array<struct<sku:string>>>")
+ }
+
+ testPruning("SPARK-47230: filter lambda keeps the elements whole") {
+ // `filter` returns the elements themselves, so reading the whole struct
is required,
+ // not a missed pruning
+ checkPrunedQuery(
+ "select id, filter(items, i -> i.active) as active_items from orders",
+ s"struct<id:int,items:$fullItemsType>")
+ }
+
+ testPruning("SPARK-47230: aggregate lambda is not narrowed") {
+ // Not narrowed today, independently of generators: nested pruning does
not look inside a
+ // higher-order `aggregate`, so `items` is read whole although only
`parts` sizes are needed
+ checkPrunedQuery(
+ "select id, aggregate(items, 0L, (acc, i) -> acc + size(i.parts)) as
total_parts " +
+ "from orders",
+ s"struct<id:int,items:$fullItemsType>")
+ }
+
+ testPruning("SPARK-47230: union of two explode chains prunes both scans") {
+ val branch =
+ """select country, item.sku, part.pid
+ |from orders
+ |lateral view explode(items) t1 as item
+ |lateral view explode(item.parts) t2 as part""".stripMargin
+ val branchSchema =
+ "struct<country:string,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<pid:int>>>>>"
+ checkPrunedQuery(
+ s"$branch\nwhere country = 'US'\nunion all\n$branch\nwhere country =
'UK'",
+ branchSchema, branchSchema)
+ }
+
+ testPruning("SPARK-47230: self-join on exploded data prunes each side
separately") {
+ checkPrunedQuery(
+ """with b as (
+ | select id, item.sku, part.pid, part.price from base)
+ |select b1.id, b1.sku, b1.pid, b2.price
+ |from b b1
+ |inner join b b2 on b1.pid = b2.pid""".stripMargin,
+ "struct<id:int,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<pid:int>>>>>",
+ "struct<items:array<struct<active:boolean," +
+ "parts:array<struct<inStock:boolean,pid:int,price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: exists subquery over an explode is not pruned") {
+ // Not pruned through yet (follow-up): the semi join an exists subquery is
rewritten into
+ // blocks pruning on both sides
+ checkPrunedQuery(
+ """select id, item.sku
+ |from orders
+ |lateral view explode(items) t1 as item
+ |where exists (
+ | select 1 from orders o2
+ | lateral view explode(o2.items) t2 as i2
+ | where i2.sku = item.sku)""".stripMargin,
+ s"struct<id:int,items:$fullItemsType>",
+ s"struct<items:$fullItemsType>")
+ }
+
+ testPruning("SPARK-47230: distinct over an exploded field") {
+ checkPrunedQuery(
+ "select distinct part.pid from base",
+
"struct<items:array<struct<active:boolean,parts:array<struct<pid:int>>>>>")
+ }
+
+ testPruning("SPARK-47230: an unused inner element reads a single field") {
+ checkPrunedQuery(
+ """select country, item.sku, count(*) as cnt
+ |from orders
+ |lateral view explode(items) t1 as item
+ |lateral view explode(item.parts) t2 as part
+ |group by country, item.sku
+ |having count(*) > 1""".stripMargin,
+ "struct<country:string,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<inStock:boolean>>>>>")
+ }
+
+ testPruning("SPARK-47230: rollup over exploded fields prunes only the outer
element") {
+ // Not pruned through yet (follow-up): NestedColumnAliasing could push
aggregate-input
+ // extractions below the Expand. It extracts `part.price` above the Expand
but grouping
+ // keys below it, so the bare `part` element flows through the Expand
whole and the full
+ // inner struct is read.
+ // In the cube test below the inner field is a grouping key instead, so
both
+ // generators are narrowed before the Expand matters
+ checkPrunedQuery(
+ """select country, item.sku, sum(part.price) as total_price
+ |from orders
+ |lateral view explode(items) t1 as item
+ |lateral view explode(item.parts) t2 as part
+ |group by rollup(country, item.sku)""".stripMargin,
+ "struct<country:string,items:array<struct<active:boolean,sku:string," +
+ s"parts:$fullPartsType>>>")
+ }
+
+ testPruning("SPARK-47230: cube over exploded fields") {
+ checkPrunedQuery(
+ """select country, part.pid, count(*) as cnt
+ |from orders
+ |lateral view explode(items) t1 as item
+ |lateral view explode(item.parts) t2 as part
+ |group by cube(country, part.pid)""".stripMargin,
+ "struct<country:string,items:array<struct<active:boolean," +
+ "parts:array<struct<pid:int>>>>>")
+ }
+
+ testPruning("SPARK-47230: case when over exploded fields") {
+ checkPrunedQuery(
+ """select country,
+ | case when part.price > 1.0 then 'high'
+ | when part.price > 0.5 then 'medium'
+ | else 'low' end as tier,
+ | part.pid
+ |from orders
+ |lateral view explode(items) t1 as item
+ |lateral view explode(item.parts) t2 as part""".stripMargin,
+ "struct<country:string,items:array<struct<active:boolean," +
+ "parts:array<struct<inStock:boolean,pid:int,price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: elements passed whole between CTEs") {
+ // `item` and `part` cross the CTE boundaries whole; only the fields the
final query
+ // uses are read
+ checkPrunedQuery(
+ """with items_exploded as (
+ | select id, country, item from orders
+ | lateral view explode(items) t1 as item),
+ |parts_exploded as (
+ | select id, country, item.sku, part from items_exploded
+ | lateral view explode(item.parts) t2 as part),
+ |aggregated as (
+ | select country, sku, avg(part.price) as avg_price
+ | from parts_exploded
+ | group by country, sku)
+ |select * from aggregated where avg_price > 0""".stripMargin,
+ "struct<country:string,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: coalesce over an outer explode chain") {
+ checkPrunedQuery(
+ """select id, coalesce(item.sku, 'unknown') as sku, coalesce(part.price,
0.0) as price
+ |from orders
+ |lateral view outer explode(items) t1 as item
+ |lateral view outer explode(item.parts) t2 as part""".stripMargin,
+
"struct<id:int,items:array<struct<sku:string,parts:array<struct<price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: string functions over an exploded field") {
+ checkPrunedQuery(
+ """select id, upper(item.sku) as sku_upper, length(item.sku) as sku_len,
+ | concat(country, '-', item.sku) as combined
+ |from orders
+ |lateral view explode(items) t1 as item""".stripMargin,
+ "struct<id:int,country:string,items:array<struct<sku:string>>>")
+ }
+
+ testPruning("SPARK-47230: collect_list and collect_set of exploded fields") {
+ checkPrunedQuery(
+ """select country, collect_list(item.sku) as skus
+ |from orders
+ |lateral view explode(items) t1 as item
+ |group by country""".stripMargin,
+ "struct<country:string,items:array<struct<sku:string>>>")
+
+ checkPrunedQuery(
+ "select country, collect_set(part.pid) as pids from base group by
country",
+ "struct<country:string,items:array<struct<active:boolean," +
+ "parts:array<struct<pid:int>>>>>")
+ }
+
+ testPruning("SPARK-47230: order by exploded fields with limit") {
+ checkPrunedQuery(
+ """select country, item.sku, part.pid, part.price
+ |from orders
+ |lateral view explode(items) t1 as item
+ |lateral view explode(item.parts) t2 as part
+ |order by country desc, part.price desc, part.pid
+ |limit 100""".stripMargin,
+ "struct<country:string,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<inStock:boolean,pid:int,price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: in subquery over an explode prunes both scans") {
+ checkPrunedQuery(
+ """select id, item.sku
+ |from orders
+ |lateral view explode(items) t1 as item
+ |where item.sku in (
+ | select i2.sku from orders o2
+ | lateral view explode(o2.items) t2 as i2
+ | where i2.active = true)""".stripMargin,
+ "struct<id:int,items:array<struct<active:boolean,sku:string>>>",
+ "struct<items:array<struct<active:boolean,sku:string>>>")
+ }
+
+ testPruning("SPARK-47230: SELECT * CTE with an outer posexplode chain") {
+ checkPrunedQuery(
+ """with exploded as (
+ | select * from orders
+ | lateral view outer posexplode(items) t1 as ipos, item
+ | lateral view outer posexplode(item.parts) t2 as ppos, part)
+ |select country, item.sku, part.pid, count(*) as cnt,
+ | sum(case when part.inStock then 1 else 0 end) as stocked,
+ | avg(part.price) as avg_price
+ |from exploded
+ |where item.active = true
+ |group by country, item.sku, part.pid""".stripMargin,
+ "struct<country:string,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<inStock:boolean,pid:int,price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: an explode of a map in the chain reads the map
whole") {
+ // Not pruned through yet (follow-up): a map is read whole through its
explode. The array
+ // explode in the same chain still prunes, and `item.sku` being the single
field used
+ // lets GeneratorNestedColumnAliasing narrow `explode(items)` itself, so
no size()
+ // field is needed either
+ checkPrunedQuery(
+ """select item.sku, k, v.score
+ |from orders
+ |lateral view explode(items) t1 as item
+ |lateral view explode(tags) t2 as k, v""".stripMargin,
+ s"struct<items:array<struct<sku:string>>,tags:$fullTagsType>")
+ }
+
+ testPruning("SPARK-47230: map_keys and map_values read the map whole") {
+ // Not narrowed today, independently of generators: the map's value struct
is read whole
+ // although only `score` is used
+ checkPrunedQuery(
+ "select id, map_keys(tags)[0] as k0, map_values(tags)[0].score as s0
from orders",
+ s"struct<id:int,tags:$fullTagsType>")
+ }
+
+ testPruning("SPARK-47230: left join of two exploded branches prunes each
side") {
+ checkPrunedQuery(
+ """with a as (
+ | select id, item.sku, part.pid from base),
+ |b as (
+ | select id, part.pid, part.price from base)
+ |select a.id, a.sku, a.pid, b.price
+ |from a left join b on a.pid = b.pid""".stripMargin,
+ "struct<id:int,items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<pid:int>>>>>",
+ "struct<items:array<struct<active:boolean," +
+ "parts:array<struct<inStock:boolean,pid:int,price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: a scala UDF over an exploded field") {
+ spark.udf.register("skuLen", (s: String) => if (s == null) -1 else
s.length)
Review Comment:
**Finding 8.** Some test state outlives its test:
- `skuLen` here and `partsCount` at line 773 are registered on the shared
session and never dropped.
- The streaming test's memory sink in `SchemaPruningSuite` (line 783) leaves
`spark_47230_stream` registered as a temp view.
`withUserDefinedFunction("skuLen" -> true) { ... }` and
`withTempView("spark_47230_stream") { ... }` clean them up.
--
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]
