sunchao commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4415883287
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SchemaPruning.scala:
##########
@@ -133,6 +134,63 @@ 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, or over the element of
a generator below
+ * it. 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, 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}")
+ })
+ val fields = identifyRootFields(Nil, renamed,
+ elementFields.map(f => RootField(f.copy(name = name), derivedFromAtt =
false)))
+ .filter(_.field.name == name)
+ val prunedElementType = (fields, elementType) match {
+ // not read: only the array sizes are needed
+ case (Seq(), StructType(structFields)) if structFields.length > 1 =>
+ StructType(Array(structFields.minBy(_.dataType.defaultSize)))
+ case (Seq(), _) => elementType
+ // used whole: the input is needed whole, which still narrows what it
is extracted from
+ case _ if fields.exists(_.derivedFromAtt) => elementType
+ case _ =>
+ pruneSchema(StructType(Array(StructField(name, elementType))),
fields).head.dataType
+ }
+ SelectedField.withDataType(input, ArrayType(prunedElementType,
containsNull))
Review Comment:
**[P1] Preserve the full map key when pruning through `map_keys`.** On ORC
V1 and V2 with `spark.sql.orc.enableVectorizedReader=false`, I wrote a map with
struct keys `(id,key,z) = (1,23,17)` and `(2,23,17)`. `SELECT pos, e.key, e.z
FROM t LATERAL VIEW OUTER posexplode(map_keys(m)) v AS pos,e` returns two rows
with positions 0 and 1 when this feature is off. With it on,
`SelectedField.withDataType` narrows the keys to `(key,z)`. Both keys are now
`(23,17)`, so the ORC row reader's `OrcMap` (`TreeMap`) collapses them before
Spark sees them and the row at position 1 disappears. Please retain key
identity through this path and add an assertion on the row count with the
non-vectorized reader.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SchemaPruning.scala:
##########
@@ -40,6 +49,19 @@ object SchemaPruning extends Rule[LogicalPlan] {
override def apply(plan: LogicalPlan): LogicalPlan =
plan transformDown {
+ case op @ GeneratorsOverScan(_, _,
+ l @ LogicalRelationWithTable(hadoopFsRelation: HadoopFsRelation, _))
=>
+ // Fields read from the generators' elements above them narrow the
scan below them.
+ pruneThroughGenerators(op, l) { (scan, otherProjects, inputRootFields)
=>
+ val allFilters = scan.filtersPushDown.reduceOption(And).toSeq ++
scan.filtersStayUp
+ prunePhysicalColumns(l, scan.projects, allFilters, hadoopFsRelation,
+ (prunedDataSchema, prunedMetadataSchema) => {
+ val prunedHadoopRelation = hadoopFsRelation.copy(
+ dataSchema = prunedDataSchema)(hadoopFsRelation.sparkSession)
Review Comment:
**[P1] Keep ORC V1 timestamp validation aligned with the file fields.**
Write an ORC table with `items = array(named_struct('a', TIMESTAMP '2024-01-02
01:02:03', 'b', TIMESTAMP_NTZ '2025-01-02 01:02:03', 'c', 22))`. `SELECT
e.b,e.c FROM t LATERAL VIEW OUTER explode(items) g AS e` succeeds with
generator pruning off. With it on, both V1 readers throw
`UNSUPPORTED_FEATURE.ORC_TYPE_CAST`. The relation now passes `<b,c>` as its
nested data schema. `OrcUtils.requestedColumnIds` zips that schema with
physical `<a,b,c>` by ordinal and compares physical `a` (`TIMESTAMP`) against
requested `b` (`TIMESTAMP_NTZ`). Please keep the physical schema available to
the compatibility check, or match the requested fields to the correct physical
fields before validating. This fails with default ORC schema-evolution settings.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -1015,8 +1016,56 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
Project(rewrittenProjectList, withFilter)
}
- def pruneColumns(plan: LogicalPlan): LogicalPlan = plan.transform {
+ def pruneColumns(plan: LogicalPlan, isSubquery: Boolean): LogicalPlan =
plan.transform {
+ case op @ SchemaPruning.GeneratorsOverScan(_, generators, sHolder:
ScanBuilderHolder)
+ if !hasBlockingPushdown(sHolder) =>
+ // Fields read from the generators' elements above them narrow the scan
below them. This
+ // prunes and builds the holder's builder before the plan above is
re-resolved. If that is
+ // then rejected, the builder must not be pruned again with a wider
schema (a builder may
+ // only narrow, e.g. Iceberg's), so the plain scan is built from a fresh
builder instead, as
+ // scan merging does, and the lower generators of the chain do not try
again.
+ var builderUsed = false
+ SchemaPruning.pruneThroughGenerators(op, sHolder) { (scan,
otherProjects, inputRootFields) =>
+ builderUsed = true
+ Some(pruneScan(scan.projects, scan.filtersStayUp,
scan.filtersPushDown, sHolder,
Review Comment:
**[P1] Preserve physical field positions for ORC V2 positional evolution.**
I wrote `items = array(named_struct('a',11,'b',22,'c',33))`, then read the file
through ORC V2 with `.option("orc.force.positional.evolution", "true")`.
`SELECT e.b,e.c FROM t LATERAL VIEW OUTER explode(items) v AS e` returns
`(22,33)` with generator pruning disabled and `(11,22)` with it enabled. The
scan requests `<b,c>` and the positional reader fills those slots from physical
`<a,b>`. Both the row and vectorized readers return the wrong values. Please
preserve the original positions or skip this pruning when the reader uses
positional evolution.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -1015,8 +1016,56 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
Project(rewrittenProjectList, withFilter)
}
- def pruneColumns(plan: LogicalPlan): LogicalPlan = plan.transform {
+ def pruneColumns(plan: LogicalPlan, isSubquery: Boolean): LogicalPlan =
plan.transform {
+ case op @ SchemaPruning.GeneratorsOverScan(_, generators, sHolder:
ScanBuilderHolder)
+ if !hasBlockingPushdown(sHolder) =>
+ // Fields read from the generators' elements above them narrow the scan
below them. This
+ // prunes and builds the holder's builder before the plan above is
re-resolved. If that is
+ // then rejected, the builder must not be pruned again with a wider
schema (a builder may
+ // only narrow, e.g. Iceberg's), so the plain scan is built from a fresh
builder instead, as
+ // scan merging does, and the lower generators of the chain do not try
again.
+ var builderUsed = false
+ SchemaPruning.pruneThroughGenerators(op, sHolder) { (scan,
otherProjects, inputRootFields) =>
+ builderUsed = true
+ Some(pruneScan(scan.projects, scan.filtersStayUp,
scan.filtersPushDown, sHolder,
+ Some(otherProjects), inputRootFields))
+ }.getOrElse {
+ if (builderUsed) {
+ val bottom = generators.last.child
+ val rebuilt = rebuildWithoutGenerators(bottom, sHolder, isSubquery)
+ op.transformDown { case p if p eq bottom => rebuilt }
+ } else {
+ op
+ }
+ }
+
case ScanOperation(project, filtersStayUp, filtersPushDown, sHolder:
ScanBuilderHolder) =>
+ pruneScan(project, filtersStayUp, filtersPushDown, sHolder, None, Nil)
+ }
+
+ /**
+ * Plans `bottom`, the projects and filters over `sHolder` below a chain of
generators, from a
+ * fresh `ScanBuilder` of the same relation, re-applying the filters pushed
to `sHolder`, like
+ * [[rebuildScan]]. Only valid when `sHolder` has no pushdown such a rebuild
cannot reproduce.
+ */
+ private def rebuildWithoutGenerators(
+ bottom: LogicalPlan,
+ sHolder: ScanBuilderHolder,
+ isSubquery: Boolean): LogicalPlan = {
+ val relation = sHolder.relation
+ val withPushedFilters = sHolder.pushedFilterExpressions.reduceOption(And)
+ .map(Filter(_, relation)).getOrElse(relation)
Review Comment:
**[P2] Replay partially pushed data filters on the fresh builder.** On V2
Parquet partitioned by `p`, I ran an outer `explode(items)` with `WHERE p = 1
AND n = 7`. The ordinary scan and the accepted rewrite both report
`PushedFilters: [IsNotNull(n), EqualTo(n,7)]`. Forcing rejection with
`throughGenerate.rejectForTesting` leaves `PushedFilters: []`. All three return
the same row, but the fallback loses Parquet row-group skipping on `n`. The
replay inserts a filter for the fully handled partition predicates directly
above the relation, leaving the residual data filter one level too high for
`pushDownFilters` to offer it to the new builder. Please replay both eligible
predicate sets together and assert the pushed filters as well as the result.
The narrow-only-builder test could also assert that rejection actually used a
second builder. My reproducer uses the test switch to exercise rejection.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/NestedSchemaPruningBenchmark.scala:
##########
@@ -177,6 +177,44 @@ abstract class NestedSchemaPruningBenchmark extends
SqlBasedBenchmark {
}
}
+ protected def explodeBenchmark(numRows: Int, numIters: Int): Unit = {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ // col4 ARRAY<STRUCT<_1: BIGINT, _2: BIGINT, _3: STRING>>: _1 and _2 are
read through an
+ // explode, _3 mimics the burden for the other fields
+ spark.range(numRows)
+ .map(x => (0 until 5).map(i => (x + i, x * i, s"$x".repeat(100))))
+ .toDF("col4")
+ .write.format(dataSourceName).save(path + "/4")
+ spark.read.format(dataSourceName).load(path +
"/4").createOrReplaceTempView("t4")
+
+ // col5 ARRAY<STRUCT<_1: BIGINT, _2: ARRAY<STRUCT<_1: BIGINT, _2:
STRING>>>>: _2._1 is read
+ // through two chained explodes, _2._2 mimics the burden for the other
fields
+ spark.range(numRows / 5)
+ .map(x => (0 until 5).map(i => (x + i, (0 until 5).map(j => (x * j,
s"$x".repeat(20))))))
+ .toDF("col5")
+ .write.format(dataSourceName).save(path + "/5")
+ spark.read.format(dataSourceName).load(path +
"/5").createOrReplaceTempView("t5")
+
+ val benchmark = new Benchmark(s"Explode", numRows, numIters, output =
output)
Review Comment:
**Use a consistent row count for the benchmark rates.** Both cases emit `5 *
numRows` rows. The first scans `numRows` input rows and explodes five elements
each. The second scans `numRows / 5` input rows and explodes 25 elements each.
Passing `numRows` to `Benchmark` makes the reported output rate five times too
small and the nanoseconds per output row five times too large. For example, the
checked-in 1,001 ms Parquet run emits 5M rows, about 5.0M rows/s. Please use
`5L * numRows` for output-row units, or separate the cases and label the
input-row units. The elapsed times and flag-on/off speedups are still
comparable.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala:
##########
@@ -372,6 +408,525 @@ abstract class SchemaPruningSuite
}
}
+ testSchemaPruning("SPARK-34956: select multiple fields of an exploded array
of struct") {
+ val exploded =
spark.table("contacts").select(explode(col("friends")).as("friend"))
+
+ val query1 = exploded.select(col("friend.first"), col("friend.last"))
+ checkScan(query1,
"struct<friends:array<struct<first:string,last:string>>>")
+ checkAnswer(query1, Row("Susan", "Smith") :: Nil)
+
+ // a filter between the generator and the projection reads its fields too
+ // (`first` is also read by the `size(friends) > 0` filter inferred for
the explode)
+ val query2 =
exploded.where(col("friend.middle").isNotNull).select(col("friend.first"))
+ checkScan(query2,
"struct<friends:array<struct<first:string,middle:string>>>")
+ checkAnswer(query2, Row("Susan") :: Nil)
+
+ val query3 = sql(
+ "select id, pos, friend.first, friend.last from contacts " +
+ "lateral view outer posexplode(friends) t as pos, friend")
+ checkScan(query3,
"struct<id:int,friends:array<struct<first:string,last:string>>>")
+ // contacts 2 and 3 have no friends column, so the outer generator gives
them a null row too
+ checkAnswer(query3,
+ Row(0, 0, "Susan", "Smith") :: Row(1, null, null, null) ::
+ Row(2, null, null, null) :: Row(3, null, null, null) :: Nil)
+ }
+
+ testSchemaPruning("SPARK-34956: an exploded element used whole is not
pruned") {
+ val query = spark.table("contacts")
+ .select(explode(col("friends")).as("friend"))
+ .select(col("friend"), col("friend.first"))
+ checkScan(query,
"struct<friends:array<struct<first:string,middle:string,last:string>>>")
+ checkAnswer(query, Row(Row("Susan", "Z.", "Smith"), "Susan") :: Nil)
+ }
+
+ testSchemaPruning("SPARK-34956: an unused exploded element reads a single
field") {
+ // only the array sizes are needed; the `size(friends) > 0` filter
inferred for the inner
+ // explode already reads the same (smallest) field
+ val query1 = sql("select id from contacts lateral view explode(friends) t
as friend")
+ checkScan(query1, "struct<id:int,friends:array<struct<first:string>>>")
+ checkAnswer(query1, Row(0) :: Nil)
+
+ val query2 = sql("select id from contacts lateral view outer
explode(friends) t as friend")
+ checkScan(query2, "struct<id:int,friends:array<struct<first:string>>>")
+ checkAnswer(query2, Row(0) :: Row(1) :: Row(2) :: Row(3) :: Nil)
+ }
+
+ testSchemaPruning("SPARK-34956: aggregate over fields of an exploded array
of struct") {
+ val query = sql(
+ "select id, count(friend.first), max(friend.last) " +
+ "from contacts lateral view explode(friends) t as friend group by id")
+ checkScan(query,
"struct<id:int,friends:array<struct<first:string,last:string>>>")
+ checkAnswer(query, Row(0, 1L, "Smith") :: Nil)
+ }
+
+ testSchemaPruning("SPARK-34956: fields of an exploded array of struct -
case-insensitivity") {
+ withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") {
+ val query = sql(
+ "select friend.FIRST, friend.Last from contacts lateral view
explode(friends) t as friend")
+ checkScan(query,
"struct<friends:array<struct<first:string,last:string>>>")
+ checkAnswer(query, Row("Susan", "Smith") :: Nil)
+ }
+ }
+
+ testSchemaPruning("SPARK-34956: pruning through a generator can be
disabled") {
+ withSQLConf(SQLConf.NESTED_SCHEMA_PRUNING_THROUGH_GENERATE_ENABLED.key ->
"false") {
+ val query = spark.table("contacts")
+ .select(explode(col("friends")).as("friend"))
+ .select(col("friend.first"), col("friend.last"))
+ checkScan(query,
"struct<friends:array<struct<first:string,middle:string,last:string>>>")
+ checkAnswer(query, Row("Susan", "Smith") :: Nil)
+ }
+ }
+
+ testSchemaPruning("SPARK-34956: a single field is still pushed into the
generator") {
+ val query = spark.table("contacts")
+ .select(explode(col("friends")).as("friend"))
+ .select(col("friend.first"))
+ checkScan(query, "struct<friends:array<struct<first:string>>>")
+ // GeneratorNestedColumnAliasing rewrote `explode(friends)` into
`explode(friends.first)`
+ // before schema pruning, which leaves nothing to prune through the
generator
+ val generatorInputTypes = query.queryExecution.optimizedPlan.collect {
+ case g: Generate => g.generator.children.head.dataType
+ }
+ assert(generatorInputTypes === Seq(ArrayType(StringType)))
+ checkAnswer(query, Row("Susan") :: Nil)
+ }
+
+ private val explodeOrderParts = "from orders lateral view outer
explode(items) t1 as item " +
+ "lateral view outer explode(item.parts) t2 as part"
+
+ testSchemaPruning("SPARK-47230: a field of the inner element of two chained
explodes") {
Review Comment:
**Avoid running the same reader matrix four times.** This test and 24 other
new cases combine `testSchemaPruning` with `withDataSourceTable` even though
they only query their own fixture. The outer helper registers four cases and
writes contacts and departments for each. The inner helper then overrides the
reader settings and runs the same three configurations each time. That produces
12 executions per case per suite for three distinct configurations. A plain
`test` around `withDataSourceTable` would retain the three reader modes and
remove the unused fixture writes. The mixed case that also queries `contacts`
still needs its outer fixture.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningThroughGenerateSuite.scala:
##########
@@ -0,0 +1,1031 @@
+/*
+ * 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") {
+ withUserDefinedFunction("skuLen" -> true) {
+ spark.udf.register("skuLen", (s: String) => if (s == null) -1 else
s.length)
+ checkPrunedQuery(
+ """select id, skuLen(item.sku) as sku_len
+ |from orders
+ |lateral view explode(items) t1 as item""".stripMargin,
+ "struct<id:int,items:array<struct<sku:string>>>")
+ }
+ }
+
+ testPruning("SPARK-47230: aliases shadowing real column names") {
+ // `country` and `memo` are real columns; using their names as aliases of
exploded
+ // fields must neither confuse the pruning nor read the shadowed columns
+ checkPrunedQuery(
+ "select part.pid as country, item.sku as memo from base",
+ "struct<items:array<struct<active:boolean,sku:string," +
+ "parts:array<struct<pid:int>>>>>")
+ }
+
+ testPruning("SPARK-47230: an explode over an aliased element field across
CTEs") {
+ // the inner generator's input is an alias (`re_parts`) of a field of the
outer element,
+ // defined in a different CTE; `sku_alias` is selected there but unused in
the end
+ checkPrunedQuery(
+ """with item_level as (
+ | select id as order_id, item.sku as sku_alias, item.parts as re_parts
+ | from orders
+ | lateral view outer explode(items) t1 as item
+ | where item.parts is not null),
+ |part_level as (
+ | select order_id, sku_alias, p.pid as part_id, p.price as boost
+ | from item_level
+ | lateral view outer explode(re_parts) t2 as p
+ | where p.price is not null)
+ |select order_id, part_id,
+ | cast(sum(coalesce(boost, 0)) as decimal(18,4)) as sum_boost,
+ | count(case when boost is not null then 1 end) as num_boost
+ |from part_level
+ |group by order_id, part_id""".stripMargin,
+
"struct<id:int,items:array<struct<parts:array<struct<pid:int,price:double>>>>>")
+ }
+
+ testPruning("SPARK-47230: query shapes using the element whole are not
pruned") {
+ // ordering by, grouping by, or aggregating the element itself needs every
field of it
+ checkPrunedQuery(
+ """select item.sku
+ |from orders
+ |lateral view explode(items) t1 as item
+ |order by item""".stripMargin,
+ s"struct<items:$fullItemsType>")
+
+ checkPrunedQuery(
+ """select max(item.sku) as msku, count(*) as cnt
+ |from orders
+ |lateral view explode(items) t1 as item
+ |group by item""".stripMargin,
+ s"struct<items:$fullItemsType>")
+
+ checkPrunedQuery(
+ """select item.qty, first(item) as it
Review Comment:
**Make the flag-off result deterministic before comparing it with a second
execution.** `first(item)` after `GROUP BY item.qty` can select either the
`chrome` or the `beta` item for `qty = 1` in this fixture. A change in shuffle
merge order can fail the equality check even when both plans return valid
results. `min(item)` still requires the whole struct and gives the comparison a
stable expected value. The `collect_list` and `collect_set` cases in this suite
have the same issue with array order and can use `sort_array` around the
aggregate. The queries establish this risk. I did not observe a flaky run.
--
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]
