LuciferYang commented on code in PR #55579: URL: https://github.com/apache/spark/pull/55579#discussion_r4489848485
########## sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala: ########## @@ -0,0 +1,798 @@ +/* + * 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.connector + +import java.io.File + +import org.apache.spark.SparkException +import org.apache.spark.sql.{DataFrame, QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, DynamicPruningExpression, DynamicPruningSubquery, EqualTo, Expression, Literal} +import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan} +import org.apache.spark.sql.connector.read.InputPartition +import org.apache.spark.sql.execution.{FilterExec, ScalarSubquery} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} +import org.apache.spark.sql.execution.datasources.{FilePartition, InMemoryFileIndex, PartitionDirectory} +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, FileScan, PushDownUtils, RuntimePartitionPruningFileScan} +import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan +import org.apache.spark.sql.execution.datasources.v2.text.TextScan +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, LongType} + +/** + * Tests for SPARK-30628: subquery partition pruning and DPP for V2 file sources. + * + * These tests exercise V2 file scans (e.g. ParquetScan) -- distinct from the in-memory V2 + * catalog used by [[org.apache.spark.sql.DynamicPartitionPruningV2Suite]] and from the V2 + * iterative-pushdown work covered by [[DataSourceV2EnhancedRuntimePartitionFilterSuite]]. + */ +abstract class DataSourceV2FileSourceDPPSuiteBase extends QueryTest + with SharedSparkSession + with AdaptiveSparkPlanHelper { + + // The tests register their temp views over temp directories deleted when the test ends. + override def afterEach(): Unit = { + try { + spark.sessionState.catalog.clearTempTables() + } finally { + super.afterEach() + } + } + + // Standard V2-only conf for DPP tests on file sources. AQE is set by the concrete subclasses, + // which mix in DisableAdaptiveExecutionSuite / EnableAdaptiveExecutionSuite: the DPP subquery is + // planned by PlanDynamicPruningFilters without AQE and by PlanAdaptiveDynamicPruningFilters with + // it, and those are two separate rules. + private def withDppV2Conf[T](thunk: => T): T = { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "", + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "2")(thunk) + } + + // Writes a partitioned `fact` and an unpartitioned `dim` table in `format` and registers them + // as temp views. `fact` has 100 rows across 10 partitions (10 rows each); `dim` has 10 rows + // with `dim_id` and `dim_val` both in [0, 10). The read schema is given explicitly because csv + // infers strings, and a cast on the join key changes whether DPP fires at all. + private def writeFactAndDim(dir: File, format: String = "parquet"): Unit = { + val factPath = new File(dir, "fact").getCanonicalPath + val dimPath = new File(dir, "dim").getCanonicalPath + spark.range(100) + .selectExpr("id", "id % 10 AS part") + .write.format(format).partitionBy("part").save(factPath) + spark.read.format(format).schema("id long, part int").load(factPath) + .createOrReplaceTempView("fact") + spark.range(10) + .selectExpr("id AS dim_id", "id AS dim_val") + .write.format(format).save(dimPath) + spark.read.format(format).schema("dim_id long, dim_val long").load(dimPath) + .createOrReplaceTempView("dim") + } + + // The BatchScanExec reading `fact`, the only partitioned table of the two. The lookup descends + // into AQE query stages, so it finds the scan whichever stage it ended up in. + private def factScanOf(df: DataFrame): BatchScanExec = { + val executedPlan = df.queryExecution.executedPlan + collectFirst(executedPlan) { + case b: BatchScanExec if b.scan.isInstanceOf[FileScan] && + b.scan.asInstanceOf[FileScan].readPartitionSchema.nonEmpty => b + }.getOrElse(fail("no fact BatchScanExec found in:\n" + executedPlan.treeString)) + } + + private def fileScanOf(df: DataFrame): RuntimePartitionPruningFileScan = + factScanOf(df).scan.asInstanceOf[RuntimePartitionPruningFileScan] + + // A copy of `scan` whose file index runs `onList` every time it lists files. + private def withListingHook(scan: ParquetScan)(onList: => Unit): ParquetScan = { + val index = scan.fileIndex + scan.copy(fileIndex = new InMemoryFileIndex(spark, index.rootPaths, Map.empty, None, + userSpecifiedPartitionSpec = Some(index.partitionSpec())) { + override def listFiles( + partitionFilters: Seq[Expression], + dataFilters: Seq[Expression]): Seq[PartitionDirectory] = { + onList + super.listFiles(partitionFilters, dataFilters) + } + }) + } + + // The splits of each partition, as (path, start, length). + private def layoutOf(parts: Array[InputPartition]): Seq[Seq[(String, Long, Long)]] = + parts.toSeq.map(_.asInstanceOf[FilePartition].files.toSeq.map { f => + (f.filePath.toString, f.start, f.length) + }) + + private def collectDppFilters(plan: LogicalPlan): Seq[Expression] = { + plan.collect { + case f: Filter if f.condition.exists(_.isInstanceOf[DynamicPruningSubquery]) => f.condition + } + } + + test("DPP inserts DynamicPruningSubquery for V2 parquet partitioned table") { + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql( + """SELECT f.id FROM fact f JOIN dim d + |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin) + val optimized = df.queryExecution.optimizedPlan + assert(collectDppFilters(optimized).nonEmpty, + "expected DynamicPruningSubquery in optimized plan over V2 parquet, got plan:\n" + + optimized.treeString) + } + } + } + + test("DPP does not fire when join key is a data column (non-partition)") { + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + // Join on f.id (data column), not f.part (partition column). `f.part` stays in the + // projection so the fact scan still reads a partition column: without it the scan declares + // nothing filterable for a reason that has nothing to do with the join key. + val df = sql( + """SELECT f.id, f.part FROM fact f JOIN dim d + |ON f.id = d.dim_id WHERE d.dim_val = 7""".stripMargin) + val optimized = df.queryExecution.optimizedPlan + assert(collectDppFilters(optimized).isEmpty, + "DPP should not fire on non-partition join keys, got plan:\n" + + optimized.treeString) + assert(df.collect().length === 1) + // And the scan is one that could have been pruned, had the join key been its partition + // column: a filterAttributes() reporting every read column would make DPP fire above, and + // the fact scan would read fewer than all 100 rows. + assert(fileScanOf(df).filterAttributes().flatMap(_.fieldNames).toSeq === Seq("part")) + val numOutputRows = factScanOf(df).metrics("numOutputRows").value + assert(numOutputRows === 100, + s"expected the fact scan to read every row, got $numOutputRows") + } + } + } + + test("DPP fires when partitioned fact is on the right side of the join") { + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + // dim on left, fact on right; join on partition column + val df = sql( + """SELECT f.id FROM dim d JOIN fact f + |ON d.dim_id = f.part WHERE d.dim_val = 7""".stripMargin) + val optimized = df.queryExecution.optimizedPlan + assert(collectDppFilters(optimized).nonEmpty, + "expected DPP on right-side partitioned fact, got plan:\n" + + optimized.treeString) + checkAnswer(df, (0 until 10).map(i => Row(i * 10 + 7))) + val numOutputRows = factScanOf(df).metrics("numOutputRows").value + assert(numOutputRows === 10, + s"expected the fact scan to read only partition 7, got $numOutputRows rows") + } + } + } + + test("DPP filter reaches BatchScanExec.runtimeFilters as DynamicPruningExpression") { + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql( + """SELECT f.id FROM fact f JOIN dim d + |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin) + df.collect() + val factScan = factScanOf(df) + // A DPP filter that degraded to true is also a DynamicPruningExpression, so match on the + // payload: only a subquery child means the filter can actually prune anything. + val dppFilters = factScan.runtimeFilters.collect { + case d @ DynamicPruningExpression(child) if !child.isInstanceOf[Literal] => d + } + assert(dppFilters.nonEmpty, + "expected the fact scan's runtimeFilters to carry a DynamicPruningExpression over a " + + s"subquery, got ${factScan.runtimeFilters}") + } + } + } + + test("DPP prunes input partitions at runtime for v2 parquet") { + // Fact has 100 rows across 10 partitions (10 rows each). DPP filter d.dim_val = 7 selects + // exactly partition 7 -> 10 rows. Without runtime pruning, all 100 rows are read. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql( + """SELECT f.id FROM fact f JOIN dim d + |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin) + val rows = df.collect() + // Correctness: 10 fact rows match (partition 7), each multiplied by 1 dim row. + assert(rows.length == 10, + s"expected 10 rows after join+filter, got ${rows.length}") + // Pruning: fact scan should have read only the 10 rows in partition 7. + val numOutputRows = factScanOf(df).metrics("numOutputRows").value + assert(numOutputRows == 10, + s"expected fact scan to read 10 rows after DPP, got $numOutputRows. plan:\n" + + df.queryExecution.executedPlan.treeString) + } + } + } + + test("DPP that prunes every partition leaves the scan with no input partition") { + // `fact5` has partitions 0..4 while the dim row selects 7, so the build side is not empty and + // the surviving partition set is. That combination is what reaches the zero-partition + // DataSourceRDD: BatchScanExec's EmptyRDDWithPartitions branch keys off SinglePartition, which + // a file scan never reports. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val factPath = new File(dir, "fact5").getCanonicalPath + spark.range(50).selectExpr("id", "id % 5 AS part") + .write.format("parquet").partitionBy("part").save(factPath) + spark.read.format("parquet").schema("id long, part int").load(factPath) + .createOrReplaceTempView("fact5") + val df = sql( + """SELECT f.id FROM fact5 f JOIN dim d + |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin) + assert(df.collect().isEmpty) + val factScan = factScanOf(df) + assert(factScan.filteredPartitions.flatten.isEmpty, + s"expected no input partition, got ${factScan.filteredPartitions.flatten.size}") + // What the assertion above does not see: BatchScanExec turning an empty partition set into + // an RDD with one empty partition still satisfies it. + assert(factScan.inputRDD.partitions.isEmpty, + s"expected an RDD with no partition, got ${factScan.inputRDD.partitions.length}") + val numOutputRows = factScan.metrics("numOutputRows").value + assert(numOutputRows == 0, + s"expected the fact scan to read nothing, got $numOutputRows rows") + } + } + } + + test("a DPP filter a union branch cannot apply is dropped instead of failing the query") { + // PartitionPruning resolves a union key through the first branch only, and the filter it + // inserts is then pushed positionally into every branch. So the unpartitioned branch's scan + // is offered a filter over a column it never declared filterable. Dropping it costs nothing, + // since that branch cannot prune on it and the join re-applies it; letting it through means an + // INTERNAL_ERROR from a query that runs today. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val otherPath = new File(dir, "fact_other").getCanonicalPath + spark.range(100).selectExpr("id", "id % 10 AS other") + .write.format("parquet").save(otherPath) + spark.read.format("parquet").schema("id long, other long").load(otherPath) + .createOrReplaceTempView("fact_other") + val df = sql( + """SELECT u.id FROM ( + | SELECT id, part FROM fact UNION ALL SELECT id, other FROM fact_other + |) u JOIN dim d ON u.part = d.dim_id WHERE d.dim_val = 7""".stripMargin) + val expected = (0 until 10).map(i => Row(i * 10 + 7)) + checkAnswer(df, expected ++ expected) + // The filter did reach the branch that cannot use it in the optimized plan, which is what + // makes this a regression test rather than a query that happens to work. + val optimizedBranchFilters = df.queryExecution.optimizedPlan.collect { + case f: Filter if f.condition.exists(_.isInstanceOf[DynamicPruningSubquery]) && + f.child.output.exists(_.name == "other") => f + } + assert(optimizedBranchFilters.size == 1, + "expected the DPP filter over the branch that cannot apply it, got plan:\n" + + df.queryExecution.optimizedPlan.treeString) + // The strategy dropped it at planning, so the scan of that branch has no runtime filter. + val otherScans = collect(df.queryExecution.executedPlan) { + case b: BatchScanExec if b.output.exists(_.name == "other") => b + } + assert(otherScans.size == 1, s"expected one scan of fact_other, got $otherScans") + assert(otherScans.head.runtimeFilters.isEmpty, + s"expected no runtime filter on fact_other, got ${otherScans.head.runtimeFilters}") + // And the branch that can use it still pruned. + val numOutputRows = factScanOf(df).metrics("numOutputRows").value + assert(numOutputRows == 10, + s"expected the partitioned branch to read 10 rows, got $numOutputRows") + } + } + } + + test("a runtime filter keeps a partition directory it cannot evaluate") { + // The join applies a DPP filter again, and the FilterExec above the scan a scalar subquery + // filter, so the scan keeps a directory whose value fails the filter's cast under ANSI. + withDppV2Conf { + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + withTempDir { dir => + writeFactAndDim(dir) + def readDirs(df: DataFrame): Seq[String] = factScanOf(df).filteredPartitions.flatten + .flatMap(_.asInstanceOf[FilePartition].files.map(_.partitionValues.getString(0))) + .distinct.sorted + val strPath = new File(dir, "fact_str").getCanonicalPath + spark.range(10) + .selectExpr("id", "CASE WHEN id = 3 THEN 'abc' ELSE CAST(id AS STRING) END AS part") + .write.format("parquet").partitionBy("part").save(strPath) + spark.read.format("parquet").schema("id long, part string").load(strPath) + .createOrReplaceTempView("fact_str") + // `f.id <> 3` drops the row of directory 'abc' before the join casts its key. + val dpp = sql( + """SELECT f.id FROM fact_str f JOIN dim d + |ON f.part = d.dim_id WHERE d.dim_val = 7 AND f.id <> 3""".stripMargin) + checkAnswer(dpp, Row(7)) + assert(factScanOf(dpp).runtimeFilters.exists(_.isInstanceOf[DynamicPruningExpression]), + s"expected a DPP filter, got ${factScanOf(dpp).runtimeFilters}") + assert(readDirs(dpp) === Seq("7", "abc")) + // Directory 'abc' holds a file with no row, so only the scan evaluates the filter on it. + val emptyPath = new File(dir, "fact_empty").getCanonicalPath + spark.range(10).selectExpr("id", "CAST(id AS STRING) AS part") + .write.format("parquet").partitionBy("part").save(emptyPath) + spark.range(0).write.parquet(new File(emptyPath, "part=abc").getCanonicalPath) + spark.read.format("parquet").schema("id long, part string").load(emptyPath) + .createOrReplaceTempView("fact_empty") + val scalar = sql( + """SELECT f.id FROM fact_empty f + |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)""".stripMargin) + checkAnswer(scalar, Row(7)) + val scalarFilters = factScanOf(scalar).runtimeFilters + assert(scalarFilters.exists(_.exists(_.isInstanceOf[ScalarSubquery])), + s"expected the scalar subquery filter on the scan, got $scalarFilters") + assert(readDirs(scalar) === Seq("7", "abc")) + } + } + } + } + + test("an empty runtime filter array returns the partitions planInputPartitions planned") { + // Spark passes an empty array when every runtime filter of a scan node is screened out, e.g. a + // DPP filter degraded to true. Packing the files again would repeat planning's work. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql("SELECT f.id, f.part FROM fact f") + df.collect() + val base = fileScanOf(df).asInstanceOf[ParquetScan] + var packings = 0 + val scan = new ParquetScan(base.sparkSession, base.hadoopConf, base.fileIndex, + base.dataSchema, base.readDataSchema, base.readPartitionSchema, base.pushedFilters, + base.options, base.pushedAggregate, base.partitionFilters, base.dataFilters, + base.pushedVariantExtractions) { + override protected def filePartitions( + selectedPartitions: Seq[PartitionDirectory]): Seq[FilePartition] = { + packings += 1 + super.filePartitions(selectedPartitions) + } + } + val planned = layoutOf(scan.planInputPartitions()) + assert(layoutOf(scan.planInputPartitionsWithRuntimeFilters(Array.empty)) === planned) + assert(packings === 1, s"expected the files to be packed once, got $packings") + } + } + } + + test("the plain and the runtime-filtered read paths share one listing") { + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql("SELECT f.id, f.part FROM fact f") + df.collect() + var listings = 0 + val scan = withListingHook(fileScanOf(df).asInstanceOf[ParquetScan]) { listings += 1 } + def partValues(parts: Array[InputPartition]): Seq[Int] = + parts.flatMap(_.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0))) + .distinct.sorted.toSeq + val partAttr = AttributeReference("part", IntegerType)() + assert(partValues(scan.planInputPartitions()) === (0 until 10)) + assert(partValues(scan.planInputPartitionsWithRuntimeFilters( + Array(EqualTo(partAttr, Literal(3))))) === Seq(3)) + assert(partValues(scan.planInputPartitions()) === (0 until 10)) + assert(listings === 1, s"expected one listing, got $listings") + } + } + } + + test("another thread can read the scan while it lists its files") { + // The thread forces a lazy val of the scan through `equals`. Listing under the scan's + // monitor would block it. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql("SELECT f.id, f.part FROM fact f") + df.collect() + val scan = fileScanOf(df).asInstanceOf[ParquetScan] + val partAttr = AttributeReference("part", IntegerType)() + Seq[ParquetScan => Array[InputPartition]]( + _.planInputPartitions(), + _.planInputPartitionsWithRuntimeFilters(Array(EqualTo(partAttr, Literal(3)))) + ).foreach { firstCall => + var probe: ParquetScan = null + probe = withListingHook(scan) { + val reader = new Thread(() => probe.equals(probe)) + reader.start() + reader.join(30000) + assert(!reader.isAlive, "the other thread is still blocked on the scan") + } + firstCall(probe) + } + } + } + } + + test("scalar subquery on partition column of v2 parquet prunes partitions") { + // Scalar subquery (SELECT max(dim_id) FROM dim WHERE dim_val = 7) evaluates to 7. + // WHERE f.part = 7 should select exactly partition 7 -> 10 rows. Without partition + // pruning, the scan reads all 100 rows and the post-filter keeps 10. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql( + """SELECT f.id FROM fact f + |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)""".stripMargin) + val rows = df.collect() + assert(rows.length == 10, + s"expected 10 rows after scalar-subquery filter, got ${rows.length}") + val numOutputRows = factScanOf(df).metrics("numOutputRows").value + assert(numOutputRows == 10, + s"expected fact scan to read 10 rows after scalar-subquery pruning, got " + + s"$numOutputRows. plan:\n" + df.queryExecution.executedPlan.treeString) + } + } + } + + test("filterAttributes reports the read partition columns") { + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql("SELECT f.id, f.part FROM fact f") + df.collect() + val declared = fileScanOf(df).filterAttributes().flatMap(_.fieldNames).toSeq + assert(declared == Seq("part"), + s"expected the partition column to be declared runtime-filterable, got $declared") + } + } + } + + test("filterAttributes reports nothing once a pushed aggregate drops the partition columns") { + // With an aggregate pushed down, ParquetScan.readSchema() returns only the aggregate schema, so + // the partition columns are not in the scan relation output. Spark resolves every declared + // runtime-filter attribute against that output and fails the query if one is missing, so + // filterAttributes() must not report them. + withDppV2Conf { + withSQLConf(SQLConf.PARQUET_AGGREGATE_PUSHDOWN_ENABLED.key -> "true") { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql("SELECT count(part) FROM fact") + checkAnswer(df, Row(100)) + val declared = fileScanOf(df).filterAttributes().flatMap(_.fieldNames).toSeq + assert(declared.isEmpty, + s"expected no runtime-filterable attribute under a pushed aggregate, got $declared") + } + } + } + } + + test("filterAttributes keeps a partition column a pushed aggregate groups by") { + // Aggregate pushdown only happens when the GROUP BY set covers every partition column, and + // those columns go into the pushed-down schema, so they stay in the scan relation output and + // stay runtime-filterable. This is the other side of the intersection from the test above. + withDppV2Conf { + withSQLConf(SQLConf.PARQUET_AGGREGATE_PUSHDOWN_ENABLED.key -> "true") { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql("SELECT part, count(*) FROM fact GROUP BY part") + checkAnswer(df, (0 until 10).map(p => Row(p, 10))) + val scan = fileScanOf(df) + // Without the pushdown the assertion below says nothing: `part` is in the read partition + // schema either way. Pin that the aggregate really went into the scan. + assert(scan.asInstanceOf[ParquetScan].pushedAggregate.isDefined, + s"expected the aggregate to be pushed into the scan, got ${scan.description()}") + val declared = scan.filterAttributes().flatMap(_.fieldNames).toSeq + assert(declared == Seq("part"), + s"expected the group-by partition column to stay declared, got $declared") + } + } + } + } + + test("a scalar subquery filter on a partition column is evaluated again after the scan") { + // The scan does not declare its partition columns fully pushed, so a plugin node that reads + // the scan's planned partitions instead of pruning them still returns the right rows, and one + // that goes through PushDownUtils.pushRuntimeFilters gets false back rather than an error. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql( + """SELECT f.id FROM fact f + |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)""".stripMargin) + checkAnswer(df, (0 until 10).map(i => Row(i * 10 + 7))) + assert(fileScanOf(df).fullyPushedFilterAttributes().isEmpty) + val executedPlan = df.queryExecution.executedPlan + assert(collect(executedPlan) { + case f: FilterExec if f.condition.exists(_.isInstanceOf[ScalarSubquery]) => f + }.nonEmpty, s"expected a FilterExec applying the scalar subquery, got:\n" + + executedPlan.treeString) + val factScan = factScanOf(df) + assert(factScan.runtimeFilters.exists(_.exists(_.isInstanceOf[ScalarSubquery])), + s"expected the scalar subquery filter on the scan, got ${factScan.runtimeFilters}") + assert(!PushDownUtils.pushRuntimeFilters( + factScan.scan, factScan.runtimeFilters, factScan.table, factScan.output)) + } + } + } + + test("a runtime filter is ANDed with the compile-time partition filters, not substituted") { + // The compile-time filter goes through `fileIndex.listFiles`, and the runtime filter prunes + // that listing, so the scan reads only what both keep. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + // min(dim_id) over dim_val > 6 is 7, so this keeps partitions 7 and 8: 7 from the runtime + // filter's lower bound, 8 from the compile-time upper bound. + val df = sql( + """SELECT f.id FROM fact f + |WHERE f.part <= 8 + | AND f.part >= (SELECT min(dim_id) FROM dim WHERE dim_val > 6)""".stripMargin) + checkAnswer(df, (0 until 10).flatMap(i => Seq(Row(i * 10 + 7), Row(i * 10 + 8)))) + val numOutputRows = factScanOf(df).metrics("numOutputRows").value + assert(numOutputRows == 20, + s"expected the fact scan to read only partitions 7 and 8, got $numOutputRows rows") + } + } + } + + test("runtime partition pruning survives a mixed-case partition column reference") { Review Comment: Done in a76b80e435f. Both tests are gone, and the format loop covers parquet, with the plan dump in its message. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala: ########## @@ -0,0 +1,798 @@ +/* + * 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.connector + +import java.io.File + +import org.apache.spark.SparkException +import org.apache.spark.sql.{DataFrame, QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, DynamicPruningExpression, DynamicPruningSubquery, EqualTo, Expression, Literal} +import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan} +import org.apache.spark.sql.connector.read.InputPartition +import org.apache.spark.sql.execution.{FilterExec, ScalarSubquery} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} +import org.apache.spark.sql.execution.datasources.{FilePartition, InMemoryFileIndex, PartitionDirectory} +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, FileScan, PushDownUtils, RuntimePartitionPruningFileScan} +import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan +import org.apache.spark.sql.execution.datasources.v2.text.TextScan +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, LongType} + +/** + * Tests for SPARK-30628: subquery partition pruning and DPP for V2 file sources. + * + * These tests exercise V2 file scans (e.g. ParquetScan) -- distinct from the in-memory V2 + * catalog used by [[org.apache.spark.sql.DynamicPartitionPruningV2Suite]] and from the V2 + * iterative-pushdown work covered by [[DataSourceV2EnhancedRuntimePartitionFilterSuite]]. + */ +abstract class DataSourceV2FileSourceDPPSuiteBase extends QueryTest + with SharedSparkSession + with AdaptiveSparkPlanHelper { + + // The tests register their temp views over temp directories deleted when the test ends. + override def afterEach(): Unit = { + try { + spark.sessionState.catalog.clearTempTables() + } finally { + super.afterEach() + } + } + + // Standard V2-only conf for DPP tests on file sources. AQE is set by the concrete subclasses, + // which mix in DisableAdaptiveExecutionSuite / EnableAdaptiveExecutionSuite: the DPP subquery is + // planned by PlanDynamicPruningFilters without AQE and by PlanAdaptiveDynamicPruningFilters with + // it, and those are two separate rules. + private def withDppV2Conf[T](thunk: => T): T = { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "", + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "2")(thunk) + } + + // Writes a partitioned `fact` and an unpartitioned `dim` table in `format` and registers them + // as temp views. `fact` has 100 rows across 10 partitions (10 rows each); `dim` has 10 rows + // with `dim_id` and `dim_val` both in [0, 10). The read schema is given explicitly because csv + // infers strings, and a cast on the join key changes whether DPP fires at all. + private def writeFactAndDim(dir: File, format: String = "parquet"): Unit = { + val factPath = new File(dir, "fact").getCanonicalPath + val dimPath = new File(dir, "dim").getCanonicalPath + spark.range(100) + .selectExpr("id", "id % 10 AS part") + .write.format(format).partitionBy("part").save(factPath) + spark.read.format(format).schema("id long, part int").load(factPath) + .createOrReplaceTempView("fact") + spark.range(10) + .selectExpr("id AS dim_id", "id AS dim_val") + .write.format(format).save(dimPath) + spark.read.format(format).schema("dim_id long, dim_val long").load(dimPath) + .createOrReplaceTempView("dim") + } + + // The BatchScanExec reading `fact`, the only partitioned table of the two. The lookup descends + // into AQE query stages, so it finds the scan whichever stage it ended up in. + private def factScanOf(df: DataFrame): BatchScanExec = { + val executedPlan = df.queryExecution.executedPlan + collectFirst(executedPlan) { + case b: BatchScanExec if b.scan.isInstanceOf[FileScan] && + b.scan.asInstanceOf[FileScan].readPartitionSchema.nonEmpty => b + }.getOrElse(fail("no fact BatchScanExec found in:\n" + executedPlan.treeString)) + } + + private def fileScanOf(df: DataFrame): RuntimePartitionPruningFileScan = + factScanOf(df).scan.asInstanceOf[RuntimePartitionPruningFileScan] + + // A copy of `scan` whose file index runs `onList` every time it lists files. + private def withListingHook(scan: ParquetScan)(onList: => Unit): ParquetScan = { + val index = scan.fileIndex + scan.copy(fileIndex = new InMemoryFileIndex(spark, index.rootPaths, Map.empty, None, + userSpecifiedPartitionSpec = Some(index.partitionSpec())) { + override def listFiles( + partitionFilters: Seq[Expression], + dataFilters: Seq[Expression]): Seq[PartitionDirectory] = { + onList + super.listFiles(partitionFilters, dataFilters) + } + }) + } + + // The splits of each partition, as (path, start, length). + private def layoutOf(parts: Array[InputPartition]): Seq[Seq[(String, Long, Long)]] = + parts.toSeq.map(_.asInstanceOf[FilePartition].files.toSeq.map { f => + (f.filePath.toString, f.start, f.length) + }) + + private def collectDppFilters(plan: LogicalPlan): Seq[Expression] = { + plan.collect { + case f: Filter if f.condition.exists(_.isInstanceOf[DynamicPruningSubquery]) => f.condition + } + } + + test("DPP inserts DynamicPruningSubquery for V2 parquet partitioned table") { + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql( + """SELECT f.id FROM fact f JOIN dim d + |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin) + val optimized = df.queryExecution.optimizedPlan + assert(collectDppFilters(optimized).nonEmpty, + "expected DynamicPruningSubquery in optimized plan over V2 parquet, got plan:\n" + + optimized.treeString) + } + } + } + + test("DPP does not fire when join key is a data column (non-partition)") { + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + // Join on f.id (data column), not f.part (partition column). `f.part` stays in the + // projection so the fact scan still reads a partition column: without it the scan declares + // nothing filterable for a reason that has nothing to do with the join key. + val df = sql( + """SELECT f.id, f.part FROM fact f JOIN dim d + |ON f.id = d.dim_id WHERE d.dim_val = 7""".stripMargin) + val optimized = df.queryExecution.optimizedPlan + assert(collectDppFilters(optimized).isEmpty, + "DPP should not fire on non-partition join keys, got plan:\n" + + optimized.treeString) + assert(df.collect().length === 1) + // And the scan is one that could have been pruned, had the join key been its partition + // column: a filterAttributes() reporting every read column would make DPP fire above, and + // the fact scan would read fewer than all 100 rows. + assert(fileScanOf(df).filterAttributes().flatMap(_.fieldNames).toSeq === Seq("part")) + val numOutputRows = factScanOf(df).metrics("numOutputRows").value + assert(numOutputRows === 100, + s"expected the fact scan to read every row, got $numOutputRows") + } + } + } + + test("DPP fires when partitioned fact is on the right side of the join") { + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + // dim on left, fact on right; join on partition column + val df = sql( + """SELECT f.id FROM dim d JOIN fact f + |ON d.dim_id = f.part WHERE d.dim_val = 7""".stripMargin) + val optimized = df.queryExecution.optimizedPlan + assert(collectDppFilters(optimized).nonEmpty, + "expected DPP on right-side partitioned fact, got plan:\n" + + optimized.treeString) + checkAnswer(df, (0 until 10).map(i => Row(i * 10 + 7))) + val numOutputRows = factScanOf(df).metrics("numOutputRows").value + assert(numOutputRows === 10, + s"expected the fact scan to read only partition 7, got $numOutputRows rows") + } + } + } + + test("DPP filter reaches BatchScanExec.runtimeFilters as DynamicPruningExpression") { + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql( + """SELECT f.id FROM fact f JOIN dim d + |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin) + df.collect() + val factScan = factScanOf(df) + // A DPP filter that degraded to true is also a DynamicPruningExpression, so match on the + // payload: only a subquery child means the filter can actually prune anything. + val dppFilters = factScan.runtimeFilters.collect { + case d @ DynamicPruningExpression(child) if !child.isInstanceOf[Literal] => d + } + assert(dppFilters.nonEmpty, + "expected the fact scan's runtimeFilters to carry a DynamicPruningExpression over a " + + s"subquery, got ${factScan.runtimeFilters}") + } + } + } + + test("DPP prunes input partitions at runtime for v2 parquet") { + // Fact has 100 rows across 10 partitions (10 rows each). DPP filter d.dim_val = 7 selects + // exactly partition 7 -> 10 rows. Without runtime pruning, all 100 rows are read. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql( + """SELECT f.id FROM fact f JOIN dim d + |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin) + val rows = df.collect() + // Correctness: 10 fact rows match (partition 7), each multiplied by 1 dim row. + assert(rows.length == 10, + s"expected 10 rows after join+filter, got ${rows.length}") + // Pruning: fact scan should have read only the 10 rows in partition 7. + val numOutputRows = factScanOf(df).metrics("numOutputRows").value + assert(numOutputRows == 10, + s"expected fact scan to read 10 rows after DPP, got $numOutputRows. plan:\n" + + df.queryExecution.executedPlan.treeString) + } + } + } + + test("DPP that prunes every partition leaves the scan with no input partition") { + // `fact5` has partitions 0..4 while the dim row selects 7, so the build side is not empty and + // the surviving partition set is. That combination is what reaches the zero-partition + // DataSourceRDD: BatchScanExec's EmptyRDDWithPartitions branch keys off SinglePartition, which + // a file scan never reports. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val factPath = new File(dir, "fact5").getCanonicalPath + spark.range(50).selectExpr("id", "id % 5 AS part") + .write.format("parquet").partitionBy("part").save(factPath) + spark.read.format("parquet").schema("id long, part int").load(factPath) + .createOrReplaceTempView("fact5") + val df = sql( + """SELECT f.id FROM fact5 f JOIN dim d + |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin) + assert(df.collect().isEmpty) + val factScan = factScanOf(df) + assert(factScan.filteredPartitions.flatten.isEmpty, + s"expected no input partition, got ${factScan.filteredPartitions.flatten.size}") + // What the assertion above does not see: BatchScanExec turning an empty partition set into + // an RDD with one empty partition still satisfies it. + assert(factScan.inputRDD.partitions.isEmpty, + s"expected an RDD with no partition, got ${factScan.inputRDD.partitions.length}") + val numOutputRows = factScan.metrics("numOutputRows").value + assert(numOutputRows == 0, + s"expected the fact scan to read nothing, got $numOutputRows rows") + } + } + } + + test("a DPP filter a union branch cannot apply is dropped instead of failing the query") { + // PartitionPruning resolves a union key through the first branch only, and the filter it + // inserts is then pushed positionally into every branch. So the unpartitioned branch's scan + // is offered a filter over a column it never declared filterable. Dropping it costs nothing, + // since that branch cannot prune on it and the join re-applies it; letting it through means an + // INTERNAL_ERROR from a query that runs today. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val otherPath = new File(dir, "fact_other").getCanonicalPath + spark.range(100).selectExpr("id", "id % 10 AS other") + .write.format("parquet").save(otherPath) + spark.read.format("parquet").schema("id long, other long").load(otherPath) + .createOrReplaceTempView("fact_other") + val df = sql( + """SELECT u.id FROM ( + | SELECT id, part FROM fact UNION ALL SELECT id, other FROM fact_other + |) u JOIN dim d ON u.part = d.dim_id WHERE d.dim_val = 7""".stripMargin) + val expected = (0 until 10).map(i => Row(i * 10 + 7)) + checkAnswer(df, expected ++ expected) + // The filter did reach the branch that cannot use it in the optimized plan, which is what + // makes this a regression test rather than a query that happens to work. + val optimizedBranchFilters = df.queryExecution.optimizedPlan.collect { + case f: Filter if f.condition.exists(_.isInstanceOf[DynamicPruningSubquery]) && + f.child.output.exists(_.name == "other") => f + } + assert(optimizedBranchFilters.size == 1, + "expected the DPP filter over the branch that cannot apply it, got plan:\n" + + df.queryExecution.optimizedPlan.treeString) + // The strategy dropped it at planning, so the scan of that branch has no runtime filter. + val otherScans = collect(df.queryExecution.executedPlan) { + case b: BatchScanExec if b.output.exists(_.name == "other") => b + } + assert(otherScans.size == 1, s"expected one scan of fact_other, got $otherScans") + assert(otherScans.head.runtimeFilters.isEmpty, + s"expected no runtime filter on fact_other, got ${otherScans.head.runtimeFilters}") + // And the branch that can use it still pruned. + val numOutputRows = factScanOf(df).metrics("numOutputRows").value + assert(numOutputRows == 10, + s"expected the partitioned branch to read 10 rows, got $numOutputRows") + } + } + } + + test("a runtime filter keeps a partition directory it cannot evaluate") { + // The join applies a DPP filter again, and the FilterExec above the scan a scalar subquery + // filter, so the scan keeps a directory whose value fails the filter's cast under ANSI. + withDppV2Conf { + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + withTempDir { dir => + writeFactAndDim(dir) + def readDirs(df: DataFrame): Seq[String] = factScanOf(df).filteredPartitions.flatten + .flatMap(_.asInstanceOf[FilePartition].files.map(_.partitionValues.getString(0))) + .distinct.sorted + val strPath = new File(dir, "fact_str").getCanonicalPath + spark.range(10) + .selectExpr("id", "CASE WHEN id = 3 THEN 'abc' ELSE CAST(id AS STRING) END AS part") + .write.format("parquet").partitionBy("part").save(strPath) + spark.read.format("parquet").schema("id long, part string").load(strPath) + .createOrReplaceTempView("fact_str") + // `f.id <> 3` drops the row of directory 'abc' before the join casts its key. + val dpp = sql( + """SELECT f.id FROM fact_str f JOIN dim d + |ON f.part = d.dim_id WHERE d.dim_val = 7 AND f.id <> 3""".stripMargin) + checkAnswer(dpp, Row(7)) + assert(factScanOf(dpp).runtimeFilters.exists(_.isInstanceOf[DynamicPruningExpression]), + s"expected a DPP filter, got ${factScanOf(dpp).runtimeFilters}") + assert(readDirs(dpp) === Seq("7", "abc")) + // Directory 'abc' holds a file with no row, so only the scan evaluates the filter on it. + val emptyPath = new File(dir, "fact_empty").getCanonicalPath + spark.range(10).selectExpr("id", "CAST(id AS STRING) AS part") + .write.format("parquet").partitionBy("part").save(emptyPath) + spark.range(0).write.parquet(new File(emptyPath, "part=abc").getCanonicalPath) + spark.read.format("parquet").schema("id long, part string").load(emptyPath) + .createOrReplaceTempView("fact_empty") + val scalar = sql( + """SELECT f.id FROM fact_empty f + |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)""".stripMargin) + checkAnswer(scalar, Row(7)) + val scalarFilters = factScanOf(scalar).runtimeFilters + assert(scalarFilters.exists(_.exists(_.isInstanceOf[ScalarSubquery])), + s"expected the scalar subquery filter on the scan, got $scalarFilters") + assert(readDirs(scalar) === Seq("7", "abc")) + } + } + } + } + + test("an empty runtime filter array returns the partitions planInputPartitions planned") { + // Spark passes an empty array when every runtime filter of a scan node is screened out, e.g. a + // DPP filter degraded to true. Packing the files again would repeat planning's work. + withDppV2Conf { + withTempDir { dir => + writeFactAndDim(dir) + val df = sql("SELECT f.id, f.part FROM fact f") + df.collect() + val base = fileScanOf(df).asInstanceOf[ParquetScan] + var packings = 0 + val scan = new ParquetScan(base.sparkSession, base.hadoopConf, base.fileIndex, Review Comment: Done in a76b80e435f. The test is now "a runtime filter that prunes nothing returns the partitions planInputPartitions planned" and uses your identity check. It covers both an empty array and a filter that keeps every directory. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala: ########## @@ -823,7 +1003,8 @@ private class BothRuntimeFilteringInterfacesScan override def filter(predicates: Array[Predicate]): Unit = {} - override def filter(expressions: Array[Expression]): Unit = {} + override def planInputPartitionsWithRuntimeFilters( Review Comment: Done in a76b80e435f. All five are removed. -- 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]
