ohad-definity commented on code in PR #59706:
URL: https://github.com/apache/spark/pull/59706#discussion_r4414288634
##########
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:
Removed in fd44a39d; the check stays in `GeneratorsOverScan.unapply`.
##########
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:
Done in fd44a39d: `withUserDefinedFunction` for both UDFs and `withTempView`
around the streaming test.
--
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]
