ohad-definity commented on code in PR #59706: URL: https://github.com/apache/spark/pull/59706#discussion_r4416821299
########## 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: Done in d9120db1: `min(item)` instead of `first(item)`, and `sort_array` around the `collect_list` / `collect_set` results. ########## 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: Thanks, fixed in d9120db1 (`5L * numRows`); I'll regenerate the results with the benchmark workflow and push them. -- 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]
