You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "LuciferYang (via GitHub)" <gi...@apache.org> on 2026/04/28 08:59:03 UTC
[PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
LuciferYang opened a new pull request, #55579:
URL: https://github.com/apache/spark/pull/55579
### What changes were proposed in this pull request?
Extends Dynamic Partition Pruning (DPP) and scalar-subquery partition pruning to V2 file sources (Parquet, ORC, CSV, JSON, Avro, Text). Before this change, V2 file scans were not recognized as DPP-eligible by `PartitionPruning.getFilterableTableScan` because they do not implement `SupportsRuntimeV2Filtering`, and scalar-subquery filters on partition columns were left as post-scan `FilterExec` without any file-level pruning.
Implementation summary (4 files modified, ~50 lines):
- `PartitionPruning.scala`: extend `getFilterableTableScan` with a case matching `DataSourceV2ScanRelation(scan: FileScan)`, resolving partition columns from `FileScan.readPartitionSchema`. The DPP rule then inserts the existing `DynamicPruningSubquery`; `PlanDynamicPruningFilters` and `CleanupDynamicPruningFilters` flow unchanged.
- `FileScan.scala`: factor partition listing into `partitionsImpl(allPartitionFilters)`. Add public `planInputPartitionsWithRuntimeFilters(extraFilters)` that re-lists files with `partitionFilters ++ extraFilters`. **Immutable** — no mutable state added on the scan.
- `BatchScanExec.scala`: `filteredPartitions` special-cases `FileScan` with non-empty `runtimeFilters` by calling the new method. Non-file scans continue to use `PushDownUtils.pushRuntimeFilters` as before, **preserving the `SupportsRuntimeV2Filtering` iterative-pushdown contract for upstream connectors (Iceberg, Delta).**
- `DataSourceV2Strategy.scala`: lift scalar-subquery filters whose references are a subset of partition columns into `BatchScanExec.runtimeFilters` via an `effectiveRuntimeFilterAttrs` computation. Filters stay in `postScanFilters` for correctness; runtime layer applies them for pruning. **Scalar subqueries are NOT wrapped in `DynamicPruningExpression`** — that type remains reserved for join-derived pruning, so the runtime-filter classifier in `DataSourceV2EnhancedRuntimePartitionFilterSuite` stays correct.
This addresses the same gap as the previously-stale [PR #37514](https://github.com/apache/spark/pull/37514) but with three deliberate design improvements:
1. Immutable runtime-filter slot (no mutating `var` on `FileScan`).
2. No new public trait — `PushDownUtils` ordering and the `SupportsRuntimeV2Filtering` contract for non-file V2 sources are untouched.
3. Scalar-subquery routing without `DynamicPruningExpression` wrapping (preserves the test-classifier semantics in the SPARK-56521 / SPARK-55596 work).
### Why are the changes needed?
V2 file sources are widely used (especially V2 Parquet), but they currently miss two important optimizations that V1 file sources have via `HadoopFsRelation`:
1. **DPP**: a star-schema query like `SELECT * FROM fact JOIN dim ON fact.part = dim.id WHERE dim.x = 'k'` reads all partitions of `fact` even when `fact.part` matches only a few values from the filtered `dim`.
2. **Scalar-subquery partition pruning**: `WHERE part_col = (SELECT max(x) FROM dim)` reads all partitions even when the subquery resolves to a single value.
This PR closes both gaps using a minimal, immutable design.
### Does this PR introduce _any_ user-facing change?
No. The change is internal to the optimizer/planner. Users see the same query results but, where applicable, with fewer files read at runtime.
### How was this patch tested?
A new test suite `DataSourceV2FileSourceDPPSuite` covers seven cases:
- DPP inserts `DynamicPruningSubquery` for V2 parquet partitioned table (positive)
- DPP does NOT fire when join key is a non-partition column (negative)
- DPP fires when partitioned fact is on the right side of the join
- DPP filter reaches `BatchScanExec.runtimeFilters` as `DynamicPruningExpression` (physical-plan layer)
- DPP prunes input partitions at runtime (verified via `numOutputRows` metric on the fact scan)
- Scalar subquery on partition column of V2 parquet prunes partitions
- Scalar-subquery filter does NOT wrap in `DynamicPruning` (regression guard)
Existing test suites confirmed green:
- `DataSourceV2EnhancedRuntimePartitionFilterSuite` (12 cases) — guards SPARK-56521 + SPARK-55596 iterative pushdown for non-file V2 sources
- `OrcV2AggregatePushDownSuite` and `ParquetV2AggregatePushDownSuite` — V2 aggregate pushdown unaffected
- `DataSourceV2SQLSuiteV2Filter` (incl. SPARK-56467) — scalar-subquery classification unchanged for non-file V2 catalogs
- `DynamicPartitionPruningV1SuiteAEOff/On` and `DynamicPartitionPruningV2SuiteAEOff/On` — V1 and in-memory V2 catalog DPP unaffected
- `SubquerySuite` — scalar-subquery handling unchanged
Total broader regression: **437 passed, 0 failed**.
Local validation:
```bash
build/sbt 'sql/testOnly *DataSourceV2FileSourceDPPSuite *DataSourceV2EnhancedRuntimePartitionFilterSuite *OrcV2AggregatePushDownSuite *ParquetV2AggregatePushDownSuite *DataSourceV2SQLSuiteV2Filter *DynamicPartitionPruningV1SuiteAEOff *DynamicPartitionPruningV1SuiteAEOn *DynamicPartitionPruningV2SuiteAEOff *DynamicPartitionPruningV2SuiteAEOn *SubquerySuite'
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 4.7
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4439286054
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizeName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizeName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizeName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizeName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute => BoundReference(fieldIndex(normalizeName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
Review Comment:
Done in bfc057d. A `FileScan` that does not extend a built-in scan keeps master's behavior, so its `planInputPartitions()` override is never bypassed. `DataSourceV2FileScanSubclassSuite` runs a DPP join and a scalar subquery on a scan with its own partition type.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -43,6 +45,7 @@ trait FileScan extends Scan
with Batch
with SupportsReportStatistics
with SupportsMetadata
+ with SupportsRuntimeCatalystFiltering
Review Comment:
Done in bfc057d. `FileScan` no longer mixes in `SupportsRuntimeCatalystFiltering`, so a subclass implementing `SupportsRuntimeFiltering` is not a dual-interface scan. `DataSourceV2FileScanSubclassSuite` runs a plain `SELECT`, a DPP join and a scalar subquery on two such scans, and checks that their own `filter` still prunes the DPP join.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizeName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizeName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizeName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizeName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute => BoundReference(fieldIndex(normalizeName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
+ .flatMap(_.files.filter(file => predicate.eval(file.partitionValues)))
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+ // Pack the survivors instead of handing back the planned partitions with files removed:
+ // `partitions` sized those bins for the whole file set, so a scan that pruned most of its files
+ // would run the same few tasks over a fraction of the data. V1 packs the pruned set the same
+ // way, from `dynamicallySelectedPartitions`.
+ val openCostInBytes = conf.filesOpenCostInBytes
+ val maxSplitBytes = FilePartition.maxSplitBytes(
+ sparkSession, keptFiles.map(_.length + openCostInBytes).sum)
+ FilePartition.getFilePartitions(sparkSession, keptFiles, maxSplitBytes).toArray
Review Comment:
With per-format opt-in in bfc057d, only the built-in scans take this path, and none of them reports a partitioning or an ordering. The new trait's scaladoc says a scan that overrides `planInputPartitions()` or reports either should not mix it in.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4489851864
##########
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
Review Comment:
Done in a76b80e435f, with your wording.
##########
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.
Review Comment:
Done in a76b80e435f, with your wording.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4169230615
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizeName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizeName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizeName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizeName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute => BoundReference(fieldIndex(normalizeName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
Review Comment:
**This path re-plans from `partitions` only and never consults an overridden `planInputPartitions()`, so my [round-2 comment](https://github.com/apache/spark/pull/55579#discussion_r4048745686) is only half addressed.** It said the old path "bypasses `partitions` / `planInputPartitions()` entirely". An override of `partitions` is now honored, but `BatchScanExec` still reads the static path through `batch.planInputPartitions()` and always takes the reader factory from `batch.createReaderFactory()`, while this method returns plain repacked `FilePartition`s.
This already happens in the wild. Apache Sedona's [`ShapefileScan`](https://github.com/apache/sedona/blob/3118f725c7e59ca0aaa534720a79f0e316e9ca6c/spark/common/src/main/scala/org/apache/sedona/sql/datasources/shapefile/ShapefileScan.scala#L79-L124) extends `FileScan` and overrides `planInputPartitions()` to group each shapefile's .shp/.shx/.dbf files into a `ShapefilePartition`, and [`ShapefilePartitionReaderFactory`](https://github.com/apache/sedona/blob/3118f725c7e59ca0aaa534720a79f0e316e9ca6c/spark/common/src/main/scala/org/apache/sedona/sql/datasources/shapefile/ShapefilePartitionReaderFactory.scala#L106-L112) throws `Unexpected partition type` for anything else. Its `RasterScan` does the same with `RasterInputPartition`. Sedona [tests hive-partitioned directories](https://github.com/apache/sedona/blob/3118f725c7e59ca0aaa534720a79f0e316e9ca6c/spark/common/src/test/scala/org/apache/sedona/sql/ShapefileTests.scala#L523), so with this PR `SELECT * FROM shp WHERE part = (SELECT ma
x(id) FROM t)`, or a DPP join on `part`, drops the `FilterExec`, routes the filter here, and every task fails. The same queries work today.
The choice has to be made at planning, because the `FilterExec` is already gone by execution. Could the opt-in be explicit for the built-in formats (or behind an overridable hook), as the alternative in that comment suggested, instead of the default for every `FileScan`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -230,48 +231,64 @@ object PushDownUtils extends Logging {
translatedFiltersPushed || partPredicatesPushed
- case catalystScan: SupportsRuntimeCatalystFiltering if pushableFilters.nonEmpty =>
- // A runtime filter is normally evaluated twice: the source prunes with it, and the
- // FilterExec above the scan applies it again. The two have to agree, so this screen
- // pushes only predicates the source can be trusted to evaluate in Spark's place.
- //
- // But pushing a non-deterministic one would let the source prune on its own coin flip and
- // Spark flip again for the rows that survive, so we handle that specially.
- //
- // Not pushing a filter is safe only while its FilterExec still evaluates it. A fully
- // pushed filter has none, so the source is its only evaluator. That is why
- // DataSourceV2Strategy deletes a FilterExec at planning time only for a filter we push
- // below, and we use the same method (isPushablePartitionFilter) to determine if we
- // should push it.
- //
- // A DPP filter degrades to TrueLiteral once its subquery is pruned away, so it matches
- // every row. The V2 path above drops these implicitly, since translateRuntimeFilterV2
- // returns None; here we push Catalyst expressions directly, so we remove them explicitly.
- val catalystFilters = pushableFilters
- .flatMap(unwrapRuntimeFilterExpression)
- .filterNot(_ == Literal.TrueLiteral)
- if (catalystFilters.nonEmpty) {
- catalystScan.filter(catalystFilters.toArray)
- true
- } else {
- false
- }
-
case _ =>
Review Comment:
**`pushRuntimeFilters` now silently returns `false` for a Catalyst-only scan, which is every V2 file scan, while `DataSourceV2Strategy` has already removed the `FilterExec` of a fully pushed scalar-subquery filter.** Before this PR a Catalyst scan reached its own branch here and was filtered. External code that still follows the old "push, then `scan.toBatch.planInputPartitions()`" pattern keeps compiling, reads every partition, and has nothing left to filter the rows. Apache Gluten's [`AbstractBatchScanExec`](https://github.com/apache/gluten/blob/39fdc5f584668ad58178714644ba199d37787c2c/shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AbstractBatchScanExec.scala#L69-L75) shim does exactly that, so a shim derived from it for a Spark with this PR would return wrong rows for `WHERE part = (SELECT max(...) ...)`. The scaladoc at L175 ("does not go through this method") is also inaccurate, since `replanWithRuntimeFilters` routes every Catalyst scan without ru
ntime filters through here.
Could this fail closed instead, e.g. throw an internal error that points to `replanWithRuntimeFilters` when a `SupportsRuntimeCatalystFiltering` scan reaches here with pushable filters?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -43,6 +45,7 @@ trait FileScan extends Scan
with Batch
with SupportsReportStatistics
with SupportsMetadata
+ with SupportsRuntimeCatalystFiltering
Review Comment:
**Mixing `SupportsRuntimeCatalystFiltering` into the trait turns every subclass that already implements `SupportsRuntimeFiltering` into a dual-interface scan, and such a scan now fails every query, not only DPP queries.** DataSourceV2Strategy.scala:182 evaluates `relation.runtimeFilterAttrs` for every scan relation, which runs `checkRuntimeFilteringInterfaces()` and throws `INTERNAL_ERROR`.
The PR description lists this with "No in-tree subclass does", but out-of-tree ones exist. Apache Kyuubi's Spark Hive connector has [`HiveScan(...) extends FileScan with SupportsRuntimeFiltering`](https://github.com/apache/kyuubi/blob/cebcfd7eb4a4d1651defc998c16054cd8fd4f1e4/extensions/spark/kyuubi-spark-connector-hive/src/main/scala/org/apache/kyuubi/spark/connector/hive/read/HiveScan.scala#L42-L52), and YTsaurus SPYT has [`YtScan ... extends FileScan with SupportsReportPartitioning with SupportsRuntimeFiltering`](https://github.com/ytsaurus/ytsaurus-spyt/blob/7c2db2b5808ef3001b5a0ce86bcbd53babe9565c/data-source/src/main/scala/org/apache/spark/sql/v2/YtScan.scala#L44-L46). On a Spark with this PR, a plain `SELECT * FROM t` through the Kyuubi connector fails at planning. The subclass cannot opt out of an interface it inherits from the trait, and dropping its own `SupportsRuntimeFiltering` would give up its re-listing through `fileIndex.listHiveFiles(partitionFilters ++ runtimeFilt
ers, ...)`.
Could the built-in scans opt in individually instead of the trait, so that these connectors keep working?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala:
##########
@@ -40,9 +39,15 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
/**
* Returns attributes this scan can be filtered by at runtime.
*
- * Spark will call [[filter]] if it can derive a runtime filter for any of these attributes.
- * Each reference must resolve against the scan relation output when Spark builds it. Attributes
- * pruned out of [[Scan.readSchema]] fail to resolve.
+ * Spark will call [[planInputPartitionsWithRuntimeFilters]] if it can derive a runtime filter
+ * for any of these attributes. Each reference must resolve against the scan relation output
+ * when Spark builds it. Attributes pruned out of [[Scan.readSchema]] fail to resolve.
+ *
+ * A dynamic partition pruning filter has no post-scan evaluator at all: Spark removes it from
+ * the post-scan filters whether or not the attribute is also in [[fullyPushedFilterAttributes]],
+ * because the join it was derived from already implies it. So an attribute exposed here must be
+ * one this scan can filter on exactly, not approximately. [[fullyPushedFilterAttributes]]
Review Comment:
Minor: the conclusion here does not follow from its premise. If "the join it was derived from already implies it", a scan that keeps extra rows, by pruning conservatively or not at all, is still correct; only dropping a row that satisfies the filter is wrong. That is what PushDownUtils.scala:247-248 says too ("Dropping it costs pruning and not correctness"), and what the in-memory fixture does: it keeps a partition on evaluation failure, and `InMemoryCatalystRuntimeFilterTable` declares `part` filterable under `days(part)` without ever pruning on it. As written, an implementor that can only prune conservatively, e.g. on a transform or bucket source column or a nested field, would leave the column out and lose DPP for no correctness gain.
Could this say instead that the scan may return rows that fail the filter, because Spark re-applies it (by the join, the rewrite, or the `FilterExec` unless fully pushed), but must never drop a row that satisfies it?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,658 @@
+/*
+ * 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, SparkThrowable}
+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.execution.FilterExec
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, FileScan}
+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.{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 {
+
+ // 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): FileScan = factScanOf(df).scan.asInstanceOf[FileScan]
+
+ 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 handed a filter over a column it never declared filterable. Dropping it costs pruning
+ // on that branch and nothing else, since 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: that is what makes this a regression
+ // test rather than a query that happens to work. Both the unpartitioned branch and `dim`
+ // are scans with no partition column, so ask for the one that was handed a live DPP filter.
+ val unusable = collect(df.queryExecution.executedPlan) {
+ case b: BatchScanExec if b.scan.isInstanceOf[FileScan] &&
+ b.scan.asInstanceOf[FileScan].readPartitionSchema.isEmpty &&
+ b.runtimeFilters.exists {
+ case DynamicPruningExpression(child) => !child.isInstanceOf[Literal]
+ case _ => false
+ } => b
+ }
+ assert(unusable.size == 1,
+ "expected the DPP filter to reach exactly the branch that cannot apply it, got " +
+ unusable.map(_.scan.description()).mkString("\n"))
+ // 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 partition value the runtime filter cannot evaluate fails the query") {
+ // A partition value the filter cannot be evaluated against has to fail the query: keeping the
+ // file would return rows that nothing filters, since the filter is declared fully pushed. What
+ // this pins is that the failure is not swallowed. The exposure is the same one a compile-time
+ // partition filter over the same column has, and it is why this path must not gain the
+ // eval-failure tolerance the iterative PartitionPredicate path has. That no FilterExec is left
+ // above the scan is pinned separately.
+ withDppV2Conf {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ 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")
+ val e = intercept[SparkThrowable] {
+ sql(
+ """SELECT f.id FROM fact_str f
+ |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)""".stripMargin)
+ .collect()
+ }
+ assert(e.getCondition === "CAST_INVALID_INPUT",
+ s"expected the failing cast to surface, got ${e.getCondition}")
+ }
+ }
+ }
+ }
+
+ test("a subclass that narrows the file set in partitions keeps it under a runtime filter") {
+ // The runtime-filter path filters what `partitions` returned, so an override of it is honored.
+ // Listing the files again instead would hand the excluded ones back as soon as a filter fired,
+ // and with the filter declared fully pushed nothing above the scan would remove them.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT f.id, f.part FROM fact f")
+ df.collect()
+ val scan = fileScanOf(df).asInstanceOf[ParquetScan]
+ // Keeps the files of one partition directory, which the second filter below does not
+ // select, and counts how often the scan asks for its partitions.
+ var partitionsCalls = 0
+ val narrowed = new ParquetScan(scan.sparkSession, scan.hadoopConf, scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema, scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters, scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ partitionsCalls += 1
+ super.partitions.map { part =>
+ part.copy(files = part.files.filter(_.partitionValues.getInt(0) == 3))
+ }.filter(_.files.nonEmpty)
+ }
+ }
+ val partAttr = AttributeReference("part", IntegerType)()
+ val kept = narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(3))))
+ assert(kept.flatMap(_.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0)))
+ .distinct.toSeq === Seq(3))
+ // And a filter selecting a directory the override dropped comes back with nothing.
+ assert(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(7)))).isEmpty)
+ // Both calls filtered one listing, which is what keeps the plan and the execution on one
+ // snapshot of the file index.
+ assert(partitionsCalls === 1, s"expected one listing, got $partitionsCalls")
+ }
+ }
+ }
+
+ 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("scalar subquery on a partition column is not re-evaluated after the scan") {
+ // FileScan declares its partition columns in fullyPushedFilterAttributes: applying such a
+ // filter selects partition directories, so every row the scan returns satisfies it and Spark
+ // drops the post-scan FilterExec, the same treatment a compile-time partition filter gets.
+ 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, got ${rows.length}")
+ val executedPlan = df.queryExecution.executedPlan
+ assert(collect(executedPlan) { case f: FilterExec => f }.isEmpty,
+ "a fully pushed partition filter should leave no FilterExec above the scan, got:\n" +
+ executedPlan.treeString)
+ }
+ }
+ }
+
+ test("a runtime filter is ANDed with the compile-time partition filters, not substituted") {
+ // The two halves are applied in different places: `partitions` puts `partitionFilters` through
+ // `fileIndex.listFiles`, and `planInputPartitionsWithRuntimeFilters` filters what that planned.
+ // Both are the scan's only evaluator: FileScanBuilder.pushFilters returns just the data filters
+ // as post-scan filters, and a partition-column runtime filter is declared fully pushed, so
+ // dropping either half returns rows that nothing filters out.
+ 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") {
+ // `f.PART` is resolved to the schema's `part` before it ever becomes a runtime filter, so this
+ // pins the end-to-end path rather than any name matching: what the scan receives is identical
+ // to the lower-case query's. The test below is the one that pins matching by exact name.
Review Comment:
Nit: "The test below is the one that pins matching by exact name" is stale. Since the guard moved to `normalizeName`, the test below asserts that `PART` binds and returns partition 7; only `id` hits the guard. The PR description has the same leftover in its test section ("or with the partition column spelled differently, hits the scan's guard"). Could we reword both, e.g. "pins how the scan matches the name" here?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizeName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizeName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
Review Comment:
**With the partition columns declared fully pushed, a scalar-subquery partition filter disappears from `EXPLAIN FORMATTED` and from the SQL UI, which uses the formatted mode by default.** Before this PR the predicate showed as the `FilterExec` condition. That node is gone now, and the formatted output never prints `BatchScanExec.runtimeFilters`: the tree line comes from `simpleStringWithNodeId`, and the details come from `DataSourceV2ScanExecBase.verboseStringWithOperatorId`, where `PartitionFilters` (L139) lists only compile-time filters. For `SELECT id FROM fact WHERE part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)`, the plan shows the hosting subquery but no `part = ...` anywhere. V1 shows the same filter under `PartitionFilters`.
Could `BatchScanExec` override `verboseStringWithOperatorId` to add `RuntimeFilters: [...]` when non-empty, with an `ExplainSuite` case for a V2 file source?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,658 @@
+/*
+ * 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, SparkThrowable}
+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.execution.FilterExec
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, FileScan}
+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.{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 {
+
+ // 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)
Review Comment:
**Every fixture has the single partition column `part`, so projected `partitionValues` are never combined with a runtime filter.** `partitions` projects the values to `readPartitionSchema` only when a strict subset of the partition columns is read (FileScan.scala:159-160), and the new binding depends on `readPartitionSchema` ordinals (FileScan.scala:234-235). Binding by `fileIndex.partitionSchema` ordinals, as V1 and an earlier revision of this method did against unprojected values, would read the wrong slot here, and every current test would still pass because the ordinal is 0 in both schemas. With no `FilterExec` left, that means wrong rows. For example, with `partitionBy("p1", "p2", "p3")` and `SELECT id, p3 FROM t WHERE p2 = (SELECT ...)`, `p2` would read the `p3` slot.
Could we add that case, plus a DPP join on `p2`, asserting `readPartitionSchema.fieldNames.toSeq === Seq("p2", "p3")` so the projection stays engaged?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizeName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizeName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizeName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizeName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute => BoundReference(fieldIndex(normalizeName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
+ .flatMap(_.files.filter(file => predicate.eval(file.partitionValues)))
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+ // Pack the survivors instead of handing back the planned partitions with files removed:
+ // `partitions` sized those bins for the whole file set, so a scan that pruned most of its files
+ // would run the same few tasks over a fraction of the data. V1 packs the pruned set the same
+ // way, from `dynamicallySelectedPartitions`.
+ val openCostInBytes = conf.filesOpenCostInBytes
+ val maxSplitBytes = FilePartition.maxSplitBytes(
+ sparkSession, keptFiles.map(_.length + openCostInBytes).sum)
+ FilePartition.getFilePartitions(sparkSession, keptFiles, maxSplitBytes).toArray
Review Comment:
**Correcting my own [round-3 suggestion](https://github.com/apache/spark/pull/55579#issuecomment-5768680813): the repack keeps only the file set of a `partitions` override, not its grouping or partition type.** There I wrote that both could be kept, but `getFilePartitions` re-bins the survivors into plain `FilePartition`s. A subclass that mixes in `SupportsReportPartitioning` and returns `HasPartitionKey` splits from `partitions` would hit "Data source must have preserved the original partitioning during runtime filtering" (PushDownUtils.scala:334-337) as soon as a DPP or scalar-subquery filter fires; YTsaurus SPYT's `YtScan` already mixes `SupportsReportPartitioning` into a `FileScan`. A subclass that reports an ordering with one sorted file per split loses it once Next Fit packs several files together. The PR description's "keeps whatever a subclass did in `partitions`" and the test comment at DataSourceV2FileSourceDPPSuite.scala:303 overstate this as well.
Could `filterAttributes()` return nothing when the scan reports partitioning or ordering (or the opt-in be per format), so that the repack only applies to scans whose layout Spark owns?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala:
##########
@@ -72,35 +80,34 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty
/**
- * Filters this scan using runtime Catalyst expressions.
+ * Plans this scan's input partitions with runtime Catalyst expressions applied.
+ *
+ * The provided expressions must be interpreted as a set of predicates that are ANDed together,
+ * and applied on top of the filters already pushed at query compilation. Implementations may use
+ * them to plan fewer [[InputPartition]]s than [[org.apache.spark.sql.connector.read.Batch]]
+ * would without them.
*
- * The provided expressions must be interpreted as a set of predicates that are ANDed together.
- * Implementations may use the expressions to prune initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s.
+ * Spark calls this method at execution time, and only when at least one runtime filter survives
+ * screening; the partitions it returns replace the ones that scan node would otherwise have
+ * read. Spark may also call `Batch.planInputPartitions()` on the same scan, so neither method may
+ * assume it is the only one used. One scan instance can back several scan nodes (e.g. the two
+ * branches of a group-based UPDATE), each with its own runtime filters, so an implementation
Review Comment:
Minor: with the old mutating `filter`, a group-based (`ReplaceData`) scan narrowed itself, so its write could ask the scan what had been read. This contract now forbids carrying state from one call to the next, but it does not say how such a write learns which groups were read. The in-memory fixture needed the new `readPartitions`, a union of every call's result, for exactly this reason. Since every `FileScan` now also enables `RowLevelOperationRuntimeGroupFiltering` for a row-level operation built on it, a connector whose commit replaces the files of the scan it configured would delete groups that were never read. Could this paragraph say that a group-based scan has to record what each call returned, and treat a node that read through `planInputPartitions()` as having read everything?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
Review Comment:
Minor: as a Scala 2.13 `lazy val`, this runs `partitions`, including a subclass override, while holding the scan instance's monitor, which `planInputPartitions()` never did before. An override that hands the work to another thread and waits, e.g. `Await.result(Future(super.partitions), timeout)`, deadlocks planning as soon as that thread first touches another lazy val of the same scan, such as `codecFactory` through `isSplitable` in a `TextBasedFileScan`. Nothing in-tree does this. Could the memo be lock-free (e.g. an `AtomicReference` set with `compareAndSet`), or at least documented on `partitions`?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,658 @@
+/*
+ * 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, SparkThrowable}
+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.execution.FilterExec
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, FileScan}
+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.{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 {
+
+ // 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): FileScan = factScanOf(df).scan.asInstanceOf[FileScan]
+
+ 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 handed a filter over a column it never declared filterable. Dropping it costs pruning
+ // on that branch and nothing else, since 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: that is what makes this a regression
+ // test rather than a query that happens to work. Both the unpartitioned branch and `dim`
+ // are scans with no partition column, so ask for the one that was handed a live DPP filter.
+ val unusable = collect(df.queryExecution.executedPlan) {
+ case b: BatchScanExec if b.scan.isInstanceOf[FileScan] &&
+ b.scan.asInstanceOf[FileScan].readPartitionSchema.isEmpty &&
+ b.runtimeFilters.exists {
+ case DynamicPruningExpression(child) => !child.isInstanceOf[Literal]
+ case _ => false
+ } => b
+ }
+ assert(unusable.size == 1,
+ "expected the DPP filter to reach exactly the branch that cannot apply it, got " +
+ unusable.map(_.scan.description()).mkString("\n"))
+ // 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 partition value the runtime filter cannot evaluate fails the query") {
+ // A partition value the filter cannot be evaluated against has to fail the query: keeping the
+ // file would return rows that nothing filters, since the filter is declared fully pushed. What
+ // this pins is that the failure is not swallowed. The exposure is the same one a compile-time
+ // partition filter over the same column has, and it is why this path must not gain the
+ // eval-failure tolerance the iterative PartitionPredicate path has. That no FilterExec is left
+ // above the scan is pinned separately.
+ withDppV2Conf {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ 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")
+ val e = intercept[SparkThrowable] {
+ sql(
+ """SELECT f.id FROM fact_str f
+ |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)""".stripMargin)
+ .collect()
+ }
+ assert(e.getCondition === "CAST_INVALID_INPUT",
+ s"expected the failing cast to surface, got ${e.getCondition}")
+ }
+ }
+ }
+ }
+
+ test("a subclass that narrows the file set in partitions keeps it under a runtime filter") {
+ // The runtime-filter path filters what `partitions` returned, so an override of it is honored.
+ // Listing the files again instead would hand the excluded ones back as soon as a filter fired,
+ // and with the filter declared fully pushed nothing above the scan would remove them.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT f.id, f.part FROM fact f")
+ df.collect()
+ val scan = fileScanOf(df).asInstanceOf[ParquetScan]
+ // Keeps the files of one partition directory, which the second filter below does not
+ // select, and counts how often the scan asks for its partitions.
+ var partitionsCalls = 0
+ val narrowed = new ParquetScan(scan.sparkSession, scan.hadoopConf, scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema, scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters, scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ partitionsCalls += 1
+ super.partitions.map { part =>
+ part.copy(files = part.files.filter(_.partitionValues.getInt(0) == 3))
+ }.filter(_.files.nonEmpty)
+ }
+ }
+ val partAttr = AttributeReference("part", IntegerType)()
+ val kept = narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(3))))
+ assert(kept.flatMap(_.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0)))
+ .distinct.toSeq === Seq(3))
+ // And a filter selecting a directory the override dropped comes back with nothing.
+ assert(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(7)))).isEmpty)
+ // Both calls filtered one listing, which is what keeps the plan and the execution on one
+ // snapshot of the file index.
+ assert(partitionsCalls === 1, s"expected one listing, got $partitionsCalls")
Review Comment:
**This test does not pin the single listing across both paths that its comment and the PR description describe.** It calls `planInputPartitionsWithRuntimeFilters` twice and never `narrowed.planInputPartitions()`, so reverting FileScan.scala:197 to the pre-PR `partitions.toArray` keeps it green, even though a real query would then list, split and pack twice: at planning through `supportsColumnar`, and again at execution. The second filter (`part = 7`, a directory the override already dropped) also cannot catch state carried over from the first call, because narrowing the memo or ANDing successive calls returns empty too.
Could the test call `narrowed.planInputPartitions()` first, then filter `part = 7` (expect empty) and `part = 3` (expect `Seq(3)`), keeping `partitionsCalls === 1`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -230,48 +231,64 @@ object PushDownUtils extends Logging {
translatedFiltersPushed || partPredicatesPushed
- case catalystScan: SupportsRuntimeCatalystFiltering if pushableFilters.nonEmpty =>
- // A runtime filter is normally evaluated twice: the source prunes with it, and the
- // FilterExec above the scan applies it again. The two have to agree, so this screen
- // pushes only predicates the source can be trusted to evaluate in Spark's place.
- //
- // But pushing a non-deterministic one would let the source prune on its own coin flip and
- // Spark flip again for the rows that survive, so we handle that specially.
- //
- // Not pushing a filter is safe only while its FilterExec still evaluates it. A fully
- // pushed filter has none, so the source is its only evaluator. That is why
- // DataSourceV2Strategy deletes a FilterExec at planning time only for a filter we push
- // below, and we use the same method (isPushablePartitionFilter) to determine if we
- // should push it.
- //
- // A DPP filter degrades to TrueLiteral once its subquery is pruned away, so it matches
- // every row. The V2 path above drops these implicitly, since translateRuntimeFilterV2
- // returns None; here we push Catalyst expressions directly, so we remove them explicitly.
- val catalystFilters = pushableFilters
- .flatMap(unwrapRuntimeFilterExpression)
- .filterNot(_ == Literal.TrueLiteral)
- if (catalystFilters.nonEmpty) {
- catalystScan.filter(catalystFilters.toArray)
- true
- } else {
- false
- }
-
case _ =>
false
}
}
+ /**
+ * Screens runtime filters for a [[SupportsRuntimeCatalystFiltering]] scan and unwraps them to the
+ * Catalyst predicates the scan should plan with. No translation to connector predicates happens
+ * here: what a scan can apply is the scan's own business.
+ *
+ * Only filters over attributes the scan declared filterable are kept. `DataSourceV2Strategy`
+ * routes a dynamic pruning filter without checking that declaration, and the key of a filter
+ * inserted over one union branch is pushed positionally into the others, so a scan can be handed
+ * a filter over a column it never declared. Dropping it costs pruning and not correctness: its
+ * rows are filtered anyway, by the join a DPP filter was derived from, and by the rewrite
+ * re-applying its own condition for a row-level operation's group filter. For a filter derived
+ * from a scalar subquery this drops nothing: the strategy already required its references to be a
+ * subset of the same declaration.
+ *
+ * The determinism screen is the other half. A scalar subquery filter is evaluated twice, by the
+ * source and by the FilterExec above the scan, so pushing a non-deterministic one would let the
+ * two disagree. `DataSourceV2Strategy` deletes that FilterExec only for a filter this method
+ * keeps, which is why both sites decide with `isPushablePartitionFilter`.
+ *
+ * A DPP filter degrades to TrueLiteral once its subquery is pruned away, so it matches every row.
+ * The V2 path drops these implicitly, since translateRuntimeFilterV2 returns None; here the
+ * expressions reach the scan directly, so we remove them explicitly.
+ */
+ private def catalystRuntimeFilters(
+ scan: SupportsRuntimeCatalystFiltering,
+ runtimeFilters: Seq[Expression],
+ output: Seq[AttributeReference]): Seq[Expression] = {
+ val filterAttrs = V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs(
+ scan.filterAttributes(),
+ output,
+ "filterAttributes()",
+ scan.getClass.getName)
+ runtimeFilters
+ .filter(f => f.references.subsetOf(filterAttrs))
Review Comment:
Minor: the union case is handled here at execution, but DataSourceV2Strategy.scala:201 could screen it at planning, the same way DataSourceV2Strategy.scala:187 already screens scalar-subquery filters by `relation.runtimeFilterAttrs`. As it is, the `fact_other` scan in the new union test keeps a live `dynamicpruningexpression(other IN dynamicpruning#..)`. It shows in EXPLAIN on a scan with no partition columns, it takes part in `equals`/`doCanonicalize` so the node no longer matches an otherwise identical scan for reuse, and its subquery is still planned. The new test then pins that the filter reaches the branch (`unusable.size == 1`). Could the strategy filter `dynamicFilters` by `relation.runtimeFilterAttrs` for `SupportsRuntimeCatalystFiltering` scans, keeping this line as a defensive check and flipping that assertion?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
Review Comment:
Minor: with partition columns declared here, `PartitionPruning` now inserts DPP for V2 file scans, but the SPARK-38148 rule in `CleanupDynamicPruningFilters` cannot see a static `part = <literal>` on the same key, because `FileScanBuilder` pushes it fully into the scan instead of leaving it in the `Filter`. So `SELECT /*+ broadcast(d) */ f.id FROM fact f JOIN dim d ON f.part = d.dim_id WHERE f.part = 3 AND d.dim_val = 7` keeps a DPP that can prune nothing, where V1 removes it. That costs a `SubqueryBroadcastExec` and a re-filter and re-pack of partition 3, and, with `reuseBroadcastOnly=false` on a non-broadcast join, a second run of the dim side. Could the rule's `DataSourceV2ScanRelation` case also read equalities from `DataSourceV2ScanRelation.pushedFilters`, with the SPARK-38148 cases run against V2 parquet?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala:
##########
@@ -862,21 +883,17 @@ private class PartitioningBreakingScan(
afterFilter: Seq[InputPartition])
extends Scan with Batch with SupportsRuntimeCatalystFiltering {
- private var filtered = false
-
override def readSchema(): StructType = new StructType().add("part", IntegerType)
override def toBatch: Batch = this
- override def planInputPartitions(): Array[InputPartition] =
- if (filtered) afterFilter.toArray else initialPartitions.toArray
+ override def planInputPartitions(): Array[InputPartition] = initialPartitions.toArray
Review Comment:
Nit: after this change nothing reaches `initialPartitions`, the `Batch` mix-in, `toBatch`, `planInputPartitions()` or `createReaderFactory()`. The only caller goes through the Catalyst branch of `replanWithRuntimeFilters`, which calls only `filterAttributes()` and `planInputPartitionsWithRuntimeFilters`, and passes `originalPartitions = Seq.empty`. Could this become `PartitioningBreakingScan(afterFilter: Seq[InputPartition]) extends SupportsRuntimeCatalystFiltering` like the other scans in this file, dropping `Batch` and `PartitionReaderFactory` from the imports?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
Review Comment:
Minor: the new DPP also changes which runtime filter V2 file scans get on shuffle joins. `PartitionPruning` inserts it with `onlyInBroadcast = true` by default, `InjectRuntimeFilter` then skips the bloom filter for that key (`!hasDynamicPruningSubquery(...)`), and on a sort-merge join `PlanDynamicPruningFilters` turns the DPP into `true`. A V2 query that used to get `might_contain(bloom, xxhash64(part))` before the shuffle now gets no runtime filter at all. This matches V1, and the real fix belongs in `InjectRuntimeFilter`, but could the user-facing section of the description mention it?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4173544005
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizeName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizeName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizeName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizeName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute => BoundReference(fieldIndex(normalizeName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
Review Comment:
This and item 2 come down to per-format opt-in versus an overridable hook. I compared the two against downstream-shaped scans in https://github.com/apache/spark/pull/55579#issuecomment-5966973332 and will change this once we agree on one.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -43,6 +45,7 @@ trait FileScan extends Scan
with Batch
with SupportsReportStatistics
with SupportsMetadata
+ with SupportsRuntimeCatalystFiltering
Review Comment:
Same as item 1. The comparison is in https://github.com/apache/spark/pull/55579#issuecomment-5966973332, along with a question for pan3793 about how Kyuubi would move.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -230,48 +231,64 @@ object PushDownUtils extends Logging {
translatedFiltersPushed || partPredicatesPushed
- case catalystScan: SupportsRuntimeCatalystFiltering if pushableFilters.nonEmpty =>
- // A runtime filter is normally evaluated twice: the source prunes with it, and the
- // FilterExec above the scan applies it again. The two have to agree, so this screen
- // pushes only predicates the source can be trusted to evaluate in Spark's place.
- //
- // But pushing a non-deterministic one would let the source prune on its own coin flip and
- // Spark flip again for the rows that survive, so we handle that specially.
- //
- // Not pushing a filter is safe only while its FilterExec still evaluates it. A fully
- // pushed filter has none, so the source is its only evaluator. That is why
- // DataSourceV2Strategy deletes a FilterExec at planning time only for a filter we push
- // below, and we use the same method (isPushablePartitionFilter) to determine if we
- // should push it.
- //
- // A DPP filter degrades to TrueLiteral once its subquery is pruned away, so it matches
- // every row. The V2 path above drops these implicitly, since translateRuntimeFilterV2
- // returns None; here we push Catalyst expressions directly, so we remove them explicitly.
- val catalystFilters = pushableFilters
- .flatMap(unwrapRuntimeFilterExpression)
- .filterNot(_ == Literal.TrueLiteral)
- if (catalystFilters.nonEmpty) {
- catalystScan.filter(catalystFilters.toArray)
- true
- } else {
- false
- }
-
case _ =>
Review Comment:
Done in 2acbef9. `pushRuntimeFilters` now throws an internal error pointing to `replanWithRuntimeFilters` when a `SupportsRuntimeCatalystFiltering` scan reaches it with a pushable filter. With nothing to push it still returns false, which is the call `replanWithRuntimeFilters` makes for a Catalyst scan with no runtime filters, and the scaladoc now says so. `Catalyst scan handed to pushRuntimeFilters with a pushable filter -> rejected` covers it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizeName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizeName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizeName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizeName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute => BoundReference(fieldIndex(normalizeName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
+ .flatMap(_.files.filter(file => predicate.eval(file.partitionValues)))
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+ // Pack the survivors instead of handing back the planned partitions with files removed:
+ // `partitions` sized those bins for the whole file set, so a scan that pruned most of its files
+ // would run the same few tasks over a fraction of the data. V1 packs the pruned set the same
+ // way, from `dynamicallySelectedPartitions`.
+ val openCostInBytes = conf.filesOpenCostInBytes
+ val maxSplitBytes = FilePartition.maxSplitBytes(
+ sparkSession, keptFiles.map(_.length + openCostInBytes).sum)
+ FilePartition.getFilePartitions(sparkSession, keptFiles, maxSplitBytes).toArray
Review Comment:
Agreed, the repack keeps only the file set. Whether to gate it on reported partitioning or ordering depends on items 1 and 2: with per-format opt-in only the built-in scans take this path, and none of them reports either. In 2acbef9 the test comment now claims only that a file the override drops stays dropped. I'll fix the description's wording when I next update it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizeName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizeName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
Review Comment:
Done in 2acbef9. `BatchScanExec` overrides `verboseStringWithOperatorId` to add a redacted `RuntimeFilters: [...]` line when it has runtime filters. A new `ExplainSuite` case checks that line for a scalar subquery on a V2 parquet partition column, and checks that no `Filter` node is left.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,658 @@
+/*
+ * 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, SparkThrowable}
+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.execution.FilterExec
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, FileScan}
+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.{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 {
+
+ // 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): FileScan = factScanOf(df).scan.asInstanceOf[FileScan]
+
+ 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 handed a filter over a column it never declared filterable. Dropping it costs pruning
+ // on that branch and nothing else, since 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: that is what makes this a regression
+ // test rather than a query that happens to work. Both the unpartitioned branch and `dim`
+ // are scans with no partition column, so ask for the one that was handed a live DPP filter.
+ val unusable = collect(df.queryExecution.executedPlan) {
+ case b: BatchScanExec if b.scan.isInstanceOf[FileScan] &&
+ b.scan.asInstanceOf[FileScan].readPartitionSchema.isEmpty &&
+ b.runtimeFilters.exists {
+ case DynamicPruningExpression(child) => !child.isInstanceOf[Literal]
+ case _ => false
+ } => b
+ }
+ assert(unusable.size == 1,
+ "expected the DPP filter to reach exactly the branch that cannot apply it, got " +
+ unusable.map(_.scan.description()).mkString("\n"))
+ // 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 partition value the runtime filter cannot evaluate fails the query") {
+ // A partition value the filter cannot be evaluated against has to fail the query: keeping the
+ // file would return rows that nothing filters, since the filter is declared fully pushed. What
+ // this pins is that the failure is not swallowed. The exposure is the same one a compile-time
+ // partition filter over the same column has, and it is why this path must not gain the
+ // eval-failure tolerance the iterative PartitionPredicate path has. That no FilterExec is left
+ // above the scan is pinned separately.
+ withDppV2Conf {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ 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")
+ val e = intercept[SparkThrowable] {
+ sql(
+ """SELECT f.id FROM fact_str f
+ |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)""".stripMargin)
+ .collect()
+ }
+ assert(e.getCondition === "CAST_INVALID_INPUT",
+ s"expected the failing cast to surface, got ${e.getCondition}")
+ }
+ }
+ }
+ }
+
+ test("a subclass that narrows the file set in partitions keeps it under a runtime filter") {
+ // The runtime-filter path filters what `partitions` returned, so an override of it is honored.
+ // Listing the files again instead would hand the excluded ones back as soon as a filter fired,
+ // and with the filter declared fully pushed nothing above the scan would remove them.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT f.id, f.part FROM fact f")
+ df.collect()
+ val scan = fileScanOf(df).asInstanceOf[ParquetScan]
+ // Keeps the files of one partition directory, which the second filter below does not
+ // select, and counts how often the scan asks for its partitions.
+ var partitionsCalls = 0
+ val narrowed = new ParquetScan(scan.sparkSession, scan.hadoopConf, scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema, scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters, scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ partitionsCalls += 1
+ super.partitions.map { part =>
+ part.copy(files = part.files.filter(_.partitionValues.getInt(0) == 3))
+ }.filter(_.files.nonEmpty)
+ }
+ }
+ val partAttr = AttributeReference("part", IntegerType)()
+ val kept = narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(3))))
+ assert(kept.flatMap(_.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0)))
+ .distinct.toSeq === Seq(3))
+ // And a filter selecting a directory the override dropped comes back with nothing.
+ assert(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(7)))).isEmpty)
+ // Both calls filtered one listing, which is what keeps the plan and the execution on one
+ // snapshot of the file index.
+ assert(partitionsCalls === 1, s"expected one listing, got $partitionsCalls")
Review Comment:
Done in 2acbef9. The test now calls `planInputPartitions()`, filters on `part = 7` (empty), calls `planInputPartitions()` again, then filters on `part = 3` (`Seq(3)`), and still expects `partitionsCalls === 1`. Changing `planInputPartitions()` back to `partitions.toArray` fails it.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4056395112
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -186,6 +209,47 @@ trait FileScan extends Scan
partitions.toArray
}
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is applied by selecting partition directories in `buildPartitions`,
+ * the same treatment a compile-time `partitionFilters` entry gets, so the scan evaluates it in
+ * full and Spark does not evaluate it again after the scan -- `FileScanBuilder.pushFilters`
+ * already keeps compile-time partition filters out of the post-scan filters for that reason.
+ * Directory selection matches a predicate's references against `fileIndex.partitionSchema` by
+ * name, so the names reported here have to be the ones that schema uses; both sides derive from
+ * it today, and a divergence would silently drop the predicate instead of failing.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizeName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizeName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ // Directory selection is the only thing that applies these, and it silently ignores a predicate
+ // whose references are not all partition columns -- which, for attributes declared fully
+ // pushed, would leave the predicate evaluated nowhere. Spark screens for that before it gets
+ // here; fail loudly rather than return wrong rows if that ever stops being true. Compare the
+ // way the file index does, by exact name against its own partition schema, so this guard cannot
+ // pass something the index will then drop.
+ val partitionNames = fileIndex.partitionSchema.fieldNames.toSet
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => partitionNames.contains(a.name)))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"partition columns ${fileIndex.partitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ buildPartitions(partitionFilters ++ expressions).toArray
Review Comment:
You are right, and my earlier answer in this thread was wrong in the case that matters. It said an override of `partitions` keeps working because the compile-time path still calls it, which only holds while no runtime filter fires, and a runtime filter firing is the whole point of the feature.
Taken, in the shape you proposed. `planInputPartitionsWithRuntimeFilters` now filters what `partitions` already planned instead of listing again: it binds the expressions against `readPartitionSchema`, drops each `PartitionedFile` whose `partitionValues` do not satisfy them, and drops the partitions left empty.
```scala
val predicate = Predicate.createInterpreted(bound)
partitions.flatMap { part =>
val kept = part.files.filter(file => predicate.eval(file.partitionValues))
if (kept.isEmpty) None else Some(part.copy(files = kept))
}.toArray
```
Three things fall out of it. An override of `partitions` is honored, since its files are the only ones reachable from here. `partitions` is memoized, so both read paths share one listing and the duplicate driver-side listing and packing is gone, which takes the first item in the description's list of things this PR does not address with it. And the `buildPartitions` seam this PR had added is no longer needed, so `partitions` goes back to what it was on master.
The guard stays, and now compares against `readPartitionSchema`, which is what the binding uses. It matches names the way the rest of the scan does, so a mixed-case reference binds rather than being rejected.
New test: `a subclass that narrows the file set in partitions keeps it under a runtime filter` builds a `ParquetScan` subclass whose `partitions` keeps one directory's files, then asks it for a filter selecting a different directory. Reverting the method to the listing-based version turns it red by handing back exactly the files the override excluded:
```
Array(FilePartition(0, Array(... /fact/part=7/part-00000-....parquet, partition values: [7])),
FilePartition(1, Array(... /fact/part=7/part-00001-....parquet, partition values: [7]))) was not empty
```
The same test asserts the scan asks for its partitions once across both calls; dropping the memoization turns that into `2 did not equal 1`.
The description's user-facing section now also says that the scan is the only evaluator of such a filter, and that a file whose partition value cannot be evaluated fails the query.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "pan3793 (via GitHub) via reviews" <re...@spark.apache.org>.
pan3793 commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-6649385617
Thanks @LuciferYang. Re-verified at `33a96ff` with the same scratch tests as my [review](https://github.com/apache/spark/pull/55579#pullrequestreview-5733830527):
1. **Split size:** fixed. Same setup (10 partitions, about 4 MB each, `minPartitionNum=20`):
| Query | Tasks | Largest split |
|---|---|---|
| V2 runtime-pruned, scalar subquery or DPP | 21 | 200,722 B |
| V2, same partition selected at compile time | 21 | 200,722 B |
| V1, scalar subquery or DPP | 21 | |
2. **Eval failure:** fixed. With ANSI on, both the DPP join and the scalar subquery return the correct 10 rows. The DPP scan keeps the `abc` directory and the join drops its rows.
3. **Empty array:** fixed. It returns the same partitions as `planInputPartitions()`.
4. **Screened-out filters:** fixed. A DPP filter degraded to `true` now reaches the scan (1 call, was 0).
5. **`CleanupDynamicPruningFilters`:** fixed. A non-file DSv2 relation keeps its DPP filter (1 before, 1 after).
6. **Description:** PR description is stale.
7. **Nits:** fixed.
Your suites plus the DSv2 DPP suites pass locally: 365 tests, 0 failures.
CI fails, merging master should fix it.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r3220519836
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala:
##########
@@ -63,10 +63,19 @@ case class BatchScanExec(
@transient private[sql] lazy val filteredPartitions: Seq[Option[InputPartition]] = {
val originalPartitioning = outputPartitioning
- val filtered = PushDownUtils.pushRuntimeFilters(scan, runtimeFilters, table, output)
+ // SPARK-30628: V2 file scans don't implement SupportsRuntimeV2Filtering, so route
+ // runtime filters through FileScan.planInputPartitionsWithRuntimeFilters as Catalyst
+ // expressions. This preserves DynamicPruningExpression and scalar-subquery semantics
+ // that V2-Predicate translation would drop.
+ val (filtered, newPartitions) = scan match {
+ case fs: FileScan if runtimeFilters.nonEmpty =>
Review Comment:
For the following, do you think we can add a guard explicitly?
> SPARK-30628: V2 file scans don't implement SupportsRuntimeV2Filtering
```scala
- case fs: FileScan if runtimeFilters.nonEmpty =>
+ case fs: FileScan if runtimeFilters.nonEmpty && !fs.isInstanceOf[SupportsRuntimeV2Filtering] =>
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-6602663552
@pan3793 thanks for the review. Addressed in 7158934 and e721ae7:
1. The memo now holds the file listing, and `FileScan.partitions` is split into a `filePartitions(selectedPartitions)` step. The runtime path filters the listed directories first, then splits and packs the survivors for their own size, as V1 does. "runtime pruning splits the surviving files for their own size" uses your `minPartitionNum=20` setup on small files and checks that a scalar-subquery scan and a DPP scan get the same splits as selecting those partitions at compile time.
2. A runtime filter that fails to evaluate on a directory now keeps it. After dongjoon's review the file scans no longer declare their partition columns fully pushed, so a scalar-subquery filter keeps its `FilterExec` and is evaluated again like a DPP filter. One test covers both kinds with ANSI on.
3. An empty array now returns the partitions `planInputPartitions()` planned, without packing the files again.
4. `replanWithRuntimeFilters` now calls `planInputPartitionsWithRuntimeFilters` with an empty array when all of a node's runtime filters are screened out, so a group-based scan sees every node that has runtime filters. I updated the interface doc and the fixture comment, and added a test where a DPP filter degraded to `true` still reaches the scan.
5. Split out to SPARK-60840 (#60053), with the in-memory table repro, for master and branch-4.x. This PR keeps the same change until that one merges, then I'll merge master here.
6. I'll rewrite the description once CI is green.
7. I shortened the comments in the three files and removed the history wording. The `holdsLock` test is replaced by one where another thread reads the scan while its files are listed. That thread blocks if the listing holds the scan's monitor.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4505536077
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala:
##########
@@ -233,6 +233,10 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join
private def hasSelectivePredicate(plan: LogicalPlan): Boolean = plan.exists {
case f: Filter => isLikelySelective(f.condition)
+ // A built-in V2 file scan takes its partition filters in without leaving a `Filter` above it,
+ // where V1 keeps one.
+ case r: DataSourceV2ScanRelation if r.scan.isInstanceOf[RuntimePartitionPruningFileScan] =>
Review Comment:
Done in c1ba7eacfd7. `hasSelectivePredicate` now reads the `pushedFilters` of any `DataSourceV2ScanRelation`, and the description says that a connector dimension can now trigger DPP. "DPP for a dimension whose filter a connector scan took in fully" runs a dimension on the in-memory catalog, whose scan takes `dp IN (1, 3)` in fully, and expects DPP and a fact scan that reads only the matching partitions.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/CleanupDynamicPruningFilters.scala:
##########
@@ -79,9 +90,14 @@ object CleanupDynamicPruningFilters extends Rule[LogicalPlan] with PredicateHelp
case p @ NodeWithOnlyDeterministicProjectAndFilter(
HiveTableRelation(_, _, _, _, _)) =>
removeUnnecessaryDynamicPruningSubquery(p)
- case p @ NodeWithOnlyDeterministicProjectAndFilter(
- _: DataSourceV2ScanRelation) =>
- removeUnnecessaryDynamicPruningSubquery(p)
+ // Only a built-in file scan is known to prune what it plans with a partition filter it fully
+ // took. Another DSv2 source may apply such a filter row by row and leave the pruning to DPP.
Review Comment:
Done in c1ba7eacfd7. The comment keeps the reason without contrasting it with the `Filter` path: another DSv2 source may apply such a filter row by row and leave the pruning to DPP, so its pushed filters are left out.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -230,48 +238,92 @@ object PushDownUtils extends Logging {
translatedFiltersPushed || partPredicatesPushed
- case catalystScan: SupportsRuntimeCatalystFiltering if pushableFilters.nonEmpty =>
- // A runtime filter is normally evaluated twice: the source prunes with it, and the
- // FilterExec above the scan applies it again. The two have to agree, so this screen
- // pushes only predicates the source can be trusted to evaluate in Spark's place.
- //
- // But pushing a non-deterministic one would let the source prune on its own coin flip and
- // Spark flip again for the rows that survive, so we handle that specially.
- //
- // Not pushing a filter is safe only while its FilterExec still evaluates it. A fully
- // pushed filter has none, so the source is its only evaluator. That is why
- // DataSourceV2Strategy deletes a FilterExec at planning time only for a filter we push
- // below, and we use the same method (isPushablePartitionFilter) to determine if we
- // should push it.
- //
- // A DPP filter degrades to TrueLiteral once its subquery is pruned away, so it matches
- // every row. The V2 path above drops these implicitly, since translateRuntimeFilterV2
- // returns None; here we push Catalyst expressions directly, so we remove them explicitly.
- val catalystFilters = pushableFilters
- .flatMap(unwrapRuntimeFilterExpression)
- .filterNot(_ == Literal.TrueLiteral)
- if (catalystFilters.nonEmpty) {
- catalystScan.filter(catalystFilters.toArray)
- true
- } else {
- false
+ case catalystScan: SupportsRuntimeCatalystFiltering
+ if runtimeFilters.nonEmpty && implementsOnlyDeprecatedFilter(catalystScan) =>
+ filterWithDeprecatedMethod(catalystScan, runtimeFilters, output)
+
+ case catalystScan: SupportsRuntimeCatalystFiltering =>
+ // Resolved only when a filter could be fully pushed, so a node with no such filter does
+ // not ask the scan.
+ lazy val fullyPushedAttrs = V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs(
+ catalystScan.fullyPushedFilterAttributes(),
+ output,
+ "fullyPushedFilterAttributes()",
+ catalystScan.getClass.getName)
+ val evaluatedOnlyByScan =
+ runtimeFilters.filter(isFullyPushedRuntimeFilter(_, fullyPushedAttrs))
+ if (evaluatedOnlyByScan.nonEmpty) {
+ throw SparkException.internalError(
+ s"${scan.getClass.getName} fully evaluates the runtime filters " +
+ s"${evaluatedOnlyByScan.mkString(", ")}, so they must go through " +
+ "PushDownUtils.replanWithRuntimeFilters, not pushRuntimeFilters.")
}
+ false
case _ =>
false
}
}
+ /**
+ * Screens runtime filters for a [[SupportsRuntimeCatalystFiltering]] scan and unwraps them to the
+ * Catalyst predicates the scan plans with, untranslated.
+ *
+ * A filter is kept only if it references attributes the scan declared filterable and passes
+ * `isPushablePartitionFilter`, the check `DataSourceV2Strategy` uses to drop a pushed filter's
+ * `FilterExec`. A DPP filter degraded to true is removed, as the V2 path does by not translating
+ * it.
+ */
+ private def catalystRuntimeFilters(
+ scan: SupportsRuntimeCatalystFiltering,
+ runtimeFilters: Seq[Expression],
+ output: Seq[AttributeReference]): Seq[Expression] = {
+ val filterAttrs = V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs(
+ scan.filterAttributes(),
+ output,
+ "filterAttributes()",
+ scan.getClass.getName)
+ runtimeFilters
+ .filter(f => f.references.subsetOf(filterAttrs))
Review Comment:
Done in c1ba7eacfd7, with the scaladoc. It says that `DataSourceV2Strategy` already plans only filters over the declared attributes, so the check matters only for a scan node built elsewhere.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
+ fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private val allPartitions = new TransientBestEffortLazyVal(() => filePartitions(listing()))
+
+ // Final so that a subclass of a built-in scan cannot change one read path without the other.
+ final override protected def partitions: Seq[FilePartition] = allPartitions()
+
+ final override def planInputPartitions(): Array[InputPartition] = allPartitions().toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ if (expressions.isEmpty) return allPartitions().toArray
+ // Directory values follow the table's partition schema, not the read one, and bind as V1
+ // binds them: with that schema's types, nullable.
+ val partitionSchema = fileIndex.partitionSchema
+ val fieldIndex = partitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ val predicates = expressions.toSeq.map { e =>
+ Predicate.createInterpreted(e.transform {
+ case a: Attribute =>
+ // PushDownUtils passes only filters over `filterAttributes()`.
+ val i = fieldIndex.getOrElse(normalizePartitionName(a.name),
+ throw SparkException.internalError("A file scan can only apply a runtime filter " +
+ s"over its partition columns ${partitionSchema.fieldNames.mkString("[", ", ", "]")}" +
+ s", got $e"))
+ BoundReference(i, partitionSchema(i).dataType, nullable = true)
+ })
+ }
+ val all = listing()
+ val kept = all.filter { dir =>
+ predicates.forall { predicate =>
+ try predicate.eval(dir.values) catch { case NonFatal(_) => true }
+ }
+ }
+ if (kept.length == all.length) allPartitions().toArray else filePartitions(kept).toArray
+ }
+
+ // Read when the scan is built, as `FileScan` does.
+ private val isCaseSensitive = conf.caseSensitiveAnalysis
+
+ private def normalizePartitionName(name: String): String = {
Review Comment:
Done in c1ba7eacfd7, with your helper. `FileScan.normalizePartitionColumnName` is `private[v2]` and calls `normalizeName`, and the trait's copy of the rule is gone.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-5970130854
> @dongjoon-hyun Items 1 and 2 set the direction for several of the others, so before picking one I ran both alternatives against `FileScan` subclasses modeled on the ones you linked.
>
> The probe is a throwaway suite with three subclasses: one shaped like Kyuubi's `HiveScan` (`FileScan with SupportsRuntimeFiltering`, re-listing with the runtime filters in `partitions`), one like SPYT's `YtScan` (the same, except it caches `partitions` and clears the cache in `filter`), and one like Sedona's `ShapefileScan` (its own partition type from `planInputPartitions()`, read by a factory that rejects any other type; Sedona's `RasterScan`, `NetCdfMetadataScan` and `GeoTiffMetadataScan` have the same shape). Each runs a plain `SELECT`, a DPP join and a scalar subquery on the partition column, next to built-in parquet.
>
> built-in parquet Kyuubi-like SPYT-like Sedona-like
> master no pruning DPP prunes DPP prunes works, no pruning
> this PR DPP and subquery prune every query fails at planning every query fails at planning DPP and subquery fail: `Unexpected partition type: FilePartition`
> A: per-format opt-in DPP and subquery prune same as master same as master same as master
> B: overridable hook DPP and subquery prune every query fails at planning every query fails at planning same as master
> B + precedence rule, memo behind the flag DPP and subquery prune same as master same as master same as master
> In the table, every query that ran returned the right rows. The planning failure is the `INTERNAL_ERROR` from the dual-interface check.
>
> Option A moves `SupportsRuntimeCatalystFiltering` and its three methods off `FileScan` into a trait that `ParquetScan`, `OrcScan`, `AvroScan` and `TextBasedFileScan` (CSV, JSON, text) mix in, so any other `FileScan` subclass sees what it sees on master. The memo has to move into that trait as well. When I left it in `FileScan.planInputPartitions()`, the SPYT-like scan's DPP stopped pruning without any error: its fact scan read 100 rows instead of 10, and only the join kept the answer right. `supportsColumnar` fills the memo at planning, and the `filter()` that later clears YtScan's own cache never reaches it. With A, items 4 and 12 also apply only to the scans that mix the trait in.
>
> Option B keeps the interface on `FileScan` and returns no filter attributes unless a subclass turns on an overridable flag, which the built-in scans do. On its own that fixes Sedona but not Kyuubi or SPYT: both dual-interface checks (DataSourceV2Relation.scala:309, PushDownUtils.scala:190) match on the type, and a subclass cannot drop a trait it inherits. Dropping the checks returns wrong rows: the Kyuubi-like and SPYT-like scans then inherit a `fullyPushedFilterAttributes()` that returns their own partition columns, Spark removes the `FilterExec`, and their `filter()` ignores the scalar subquery, so that query returned 100 rows instead of 10. What worked in the probe is the last row: a precedence rule, where a dual-interface scan takes the V2 path and Spark ignores its `fullyPushedFilterAttributes()`, plus the memo behind the same flag. Without the memo change the SPYT-like fact scan read 100 rows under DPP, the same silent loss as leaving the memo in `FileScan` under A. The pr
ecedence rule replaces the one-interface rule [SPARK-58523](https://issues.apache.org/jira/browse/SPARK-58523) set, and the probe covers only the shapes above. The other route is that each such connector moves to the Catalyst interface.
>
> @pan3793, for Kyuubi the second route means `HiveScan`, `KyuubiParquetScan` and `KyuubiOrcScan` drop `SupportsRuntimeFiltering` and each override `planInputPartitionsWithRuntimeFilters`. All three re-list with the runtime filters today, `HiveScan` through `listHiveFiles(partitionFilters ++ runtimeFilters)` and the Parquet and ORC wrappers through a sibling scan built with `partitionFilters ++ runtimeFilters`, while the inherited default filters a listing made without them. The filters would arrive as Catalyst expressions rather than `In` filters, and since `FileScan` declares the partition columns fully pushed, each override has to apply every expression it gets, or return nothing from `fullyPushedFilterAttributes()`. The interface is already in 4.3, with a mutating `filter(Array[Expression])` that this PR replaces, and the connector is one binary for Spark 3.5 through 4.2 (compiled against 3.5, tested on 4.0 to 4.2), so these scans would need separately compiled classes for 4.4
, as this PR is meant for branch-4.x too. Does that match how you see it?
>
> My preference is A: among `FileScan` subclasses it changes behavior only for the built-in scans, and it leaves the one-interface rule alone. A third-party `FileScan` that wants runtime pruning can mix the trait in itself. I'd like to hear what you both think before I change it. I'll answer the other comments inline.
also cc @cloud-fan @peter-toth
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4454062991
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala:
##########
@@ -72,35 +80,38 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty
/**
- * Filters this scan using runtime Catalyst expressions.
+ * Plans this scan's input partitions with runtime Catalyst expressions applied.
+ *
+ * The provided expressions must be interpreted as a set of predicates that are ANDed together,
+ * and applied on top of the filters already pushed at query compilation. Implementations may use
+ * them to plan fewer [[InputPartition]]s than [[org.apache.spark.sql.connector.read.Batch]]
+ * would without them.
*
- * The provided expressions must be interpreted as a set of predicates that are ANDed together.
- * Implementations may use the expressions to prune initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s.
+ * Spark calls this method at execution time, and only when at least one runtime filter survives
+ * screening; the partitions it returns replace the ones that scan node would otherwise have
+ * read. Spark may also call `Batch.planInputPartitions()` on the same scan, so neither method may
+ * assume it is the only one used. One scan instance can back several scan nodes (e.g. the two
+ * branches of a group-based UPDATE), each with its own runtime filters, so an implementation
+ * must derive its result from the given expressions alone and must not carry state over from a
+ * previous call. A group-based scan whose write replaces the groups that were read has to record
+ * what each call returned for that write, and count a scan node that read through
+ * `planInputPartitions()` as having read every group that method returns. Spark may also call
+ * that method while planning a scan node, so such a call alone does not show that the node read
+ * through it.
*
* Spark currently tracks runtime-filter eligibility by root attribute. If [[filterAttributes]]
* returns a nested reference, an expression may access another nested field under the same root.
* The scan must match each access against its own partition layout and use only expressions it
* can apply.
*
- * Spark may call this method more than once for the same scan instance: a plan can hold several
- * scan nodes sharing one scan (e.g. the two branches of a group-based UPDATE), and each pushes
- * its own copy of the runtime filters. Implementations must treat successive calls as additive,
- * ANDing the new expressions with those already pushed rather than replacing them.
- *
* If the scan also implements
* [[org.apache.spark.sql.connector.read.SupportsReportPartitioning]], it must preserve
- * the originally reported partitioning during runtime filtering. While applying runtime
- * predicates, the scan may detect that some
- * [[org.apache.spark.sql.connector.read.InputPartition]]s have no matching data, in which
- * case it can either replace the initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s that have no matching data with
- * empty [[org.apache.spark.sql.connector.read.InputPartition]]s, or report only a subset of
- * the original partition values (omitting those with no data) via
- * [[org.apache.spark.sql.connector.read.Batch#planInputPartitions]]. The scan must not
- * report new partition values that were not present in the original partitioning.
- *
- * Note that Spark will call [[Scan.toBatch]] again after filtering the scan at runtime.
+ * the originally reported partitioning. While applying runtime predicates, the scan may detect
+ * that some [[InputPartition]]s have no matching data, in which case it can either replace
+ * them with empty [[InputPartition]]s, or report only a subset of the original partition
+ * values (omitting those with no data). The scan must not report new partition values that
+ * were not present in the original partitioning.
*/
- def filter(expressions: Array[Expression]): Unit
+ def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition]
Review Comment:
Done in e721ae7. `planInputPartitionsWithRuntimeFilters` now has a default that calls `filter` when the array is not empty and then returns `toBatch.planInputPartitions()`, and `filter` is deprecated since 4.4.0 with a default that throws an internal error. The empty-array check is there because 4.3.0 never called `filter` without a filter to push. A caller that pushes through `PushDownUtils.pushRuntimeFilters` and then plans the scan, as 4.3.0 did, also still prunes such a scan. "scan implementing only the deprecated filter -> still pruned, as in Spark 4.3" covers both paths with a scan that implements only `filter`. With a minimal trait under scalac 2.13.18, I also checked that a class compiled against the old shape reaches the new default at runtime. A Scala scan recompiled against this needs `override` on its `filter`, since `filter` now has a body.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
Review Comment:
Done in e721ae7. The file scans no longer declare `fullyPushedFilterAttributes()`, so a scalar-subquery filter keeps its `FilterExec`, and a node that replaces `BatchScanExec` without pruning still returns the right rows. Since every runtime filter is now evaluated again above the scan, a filter that fails to evaluate on a directory keeps it, whether it is a DPP filter or a scalar-subquery filter. "a scalar subquery filter on a partition column is evaluated again after the scan" checks the `FilterExec`. `PushDownUtils.pushRuntimeFilters` now also returns false for these scans instead of failing, so a node that pushes through it and then plans the scan reads every partition and the join or the `FilterExec` drops the extra rows.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/TextBasedFileScan.scala:
##########
@@ -27,7 +27,7 @@ import org.apache.spark.util.Utils
abstract class TextBasedFileScan(
sparkSession: SparkSession,
- options: CaseInsensitiveStringMap) extends FileScan {
+ options: CaseInsensitiveStringMap) extends RuntimePartitionPruningFileScan {
Review Comment:
Done in e721ae7. `TextBasedFileScan` is back to what master has, and `CSVScan`, `JsonScan` and `TextScan` mix in the trait themselves.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
+ .flatMap(_.files.filter(file => predicate.eval(file.partitionValues)))
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+ // Pack the survivors instead of handing back the planned partitions with files removed:
+ // `partitions` sized those bins for the whole file set, so a scan that pruned most of its files
+ // would run the same few tasks over a fraction of the data. V1 packs the pruned set the same
+ // way, from `dynamicallySelectedPartitions`.
+ val openCostInBytes = conf.filesOpenCostInBytes
+ val maxSplitBytes = FilePartition.maxSplitBytes(
+ sparkSession, keptFiles.map(_.length + openCostInBytes).sum)
Review Comment:
This went away in 7158934. The runtime path now filters the listed directories and passes them to `filePartitions`, which sums the length and open cost per file, as V1 does.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType, a.nullable)
Review Comment:
Done in e721ae7, with one change from your snippet. Since 7158934 the index is into `fileIndex.partitionSchema`, the schema the directory values follow, so the type comes from there rather than from `readPartitionSchema`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
+ .flatMap(_.files.filter(file => predicate.eval(file.partitionValues)))
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+ // Pack the survivors instead of handing back the planned partitions with files removed:
+ // `partitions` sized those bins for the whole file set, so a scan that pruned most of its files
+ // would run the same few tasks over a fraction of the data. V1 packs the pruned set the same
+ // way, from `dynamicallySelectedPartitions`.
+ val openCostInBytes = conf.filesOpenCostInBytes
+ val maxSplitBytes = FilePartition.maxSplitBytes(
+ sparkSession, keptFiles.map(_.length + openCostInBytes).sum)
+ FilePartition.getFilePartitions(sparkSession, keptFiles, maxSplitBytes).toArray
+ }
+
+ // Read when the scan is built, as `FileScan` reads it for the names `partitions` matches, so the
+ // runtime filters bind the same way even if the session's setting changes before they arrive.
+ private val isCaseSensitive = conf.caseSensitiveAnalysis
+
+ private def normalizePartitionName(name: String): String = {
+ if (isCaseSensitive) name else name.toLowerCase(Locale.ROOT)
+ }
Review Comment:
I'd rather keep the copy. Kyuubi's `HiveScan` extends `FileScan` and defines its own `private def normalizeName`. If the one in `FileScan` becomes `protected` or package-private, scalac 2.13 rejects that private method with "weaker access privileges in overriding", so Kyuubi would stop compiling.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,787 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+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
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, FileScan, RuntimePartitionPruningFileScan}
+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.{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 {
+
+ // 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]
+
+ 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 partition value the runtime filter cannot evaluate fails the query") {
+ // A partition value the filter cannot be evaluated against has to fail the query: keeping the
+ // file would return rows that nothing filters, since the filter is declared fully pushed. What
+ // this pins is that the failure is not swallowed. The exposure is the same one a compile-time
+ // partition filter over the same column has, and it is why this path must not gain the
+ // eval-failure tolerance the iterative PartitionPredicate path has. That no FilterExec is left
+ // above the scan is pinned separately.
+ withDppV2Conf {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ 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")
+ val e = intercept[SparkThrowable] {
+ sql(
+ """SELECT f.id FROM fact_str f
+ |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)""".stripMargin)
+ .collect()
+ }
+ assert(e.getCondition === "CAST_INVALID_INPUT",
+ s"expected the failing cast to surface, got ${e.getCondition}")
+ }
+ }
+ }
+ }
+
+ test("a subclass that narrows the file set in partitions keeps it under a runtime filter") {
+ // The runtime-filter path filters the files `partitions` returned, so a file an override of it
+ // drops stays dropped. Listing the files again instead would hand it back as soon as a filter
+ // fired, and with the filter declared fully pushed nothing above the scan would remove 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]
+ // Keeps only partition directory 3's files and counts how often the scan asks for its
+ // partitions.
+ var partitionsCalls = 0
+ val narrowed = new ParquetScan(scan.sparkSession, scan.hadoopConf, scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema, scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters, scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ partitionsCalls += 1
+ super.partitions.map { part =>
+ part.copy(files = part.files.filter(_.partitionValues.getInt(0) == 3))
+ }.filter(_.files.nonEmpty)
+ }
+ }
+ def partValues(parts: Array[InputPartition]): Seq[Int] =
+ parts.flatMap(_.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0)))
+ .distinct.toSeq
+ val partAttr = AttributeReference("part", IntegerType)()
+ // The plain read path first, as planning does through `supportsColumnar`.
+ assert(partValues(narrowed.planInputPartitions()) === Seq(3))
+ // A filter selecting a directory the override dropped comes back with nothing, and the
+ // calls after it, plain and filtered, still find the directory it kept: no call narrows
+ // what the next one starts from.
+ assert(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(7)))).isEmpty)
+ assert(partValues(narrowed.planInputPartitions()) === Seq(3))
+ assert(partValues(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(3))))) === Seq(3))
+ // All four calls above shared one listing, which is what keeps the plan and the execution
+ // on one snapshot of the file index.
+ assert(partitionsCalls === 1, s"expected one listing, got $partitionsCalls")
+ }
+ }
+ }
+
+ test("partitions runs without holding the scan's monitor") {
+ // An override of `partitions` may hand work to another thread and wait for it. If the memo
+ // held the scan's monitor meanwhile, that thread would block when it tried to initialize one
+ // of the scan's lazy vals.
+ 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 =>
+ val heldLock = mutable.ArrayBuffer.empty[Boolean]
+ val probe = new ParquetScan(scan.sparkSession, scan.hadoopConf, scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema, scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters, scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ heldLock += Thread.holdsLock(this)
+ super.partitions
+ }
+ }
+ firstCall(probe)
+ assert(heldLock.toSeq === Seq(false))
+ }
+ }
+ }
+ }
+
+ 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("scalar subquery on a partition column is not re-evaluated after the scan") {
+ // The file scan declares its partition columns in fullyPushedFilterAttributes: applying such a
+ // filter selects partition directories, so every row the scan returns satisfies it and Spark
+ // drops the post-scan FilterExec, the same treatment a compile-time partition filter gets.
+ 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, got ${rows.length}")
+ val executedPlan = df.queryExecution.executedPlan
+ assert(collect(executedPlan) { case f: FilterExec => f }.isEmpty,
+ "a fully pushed partition filter should leave no FilterExec above the scan, got:\n" +
+ executedPlan.treeString)
+ }
+ }
+ }
+
+ test("a runtime filter is ANDed with the compile-time partition filters, not substituted") {
+ // The two halves are applied in different places: `partitions` puts `partitionFilters` through
+ // `fileIndex.listFiles`, and `planInputPartitionsWithRuntimeFilters` filters what that planned.
+ // Both are the scan's only evaluator: FileScanBuilder.pushFilters returns just the data filters
+ // as post-scan filters, and a partition-column runtime filter is declared fully pushed, so
+ // dropping either half returns rows that nothing filters out.
+ 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") {
+ // `f.PART` is resolved to the schema's `part` before it ever becomes a runtime filter, so this
+ // pins the end-to-end path rather than any name matching: what the scan receives is identical
+ // to the lower-case query's. The test below pins how the scan matches the name.
+ 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, got ${rows.length}")
+ val numOutputRows = factScanOf(df).metrics("numOutputRows").value
+ assert(numOutputRows == 10,
+ s"expected the fact scan to read 10 rows, got $numOutputRows")
+ }
+ }
+ }
+
+ test("a runtime filter the scan cannot apply is rejected rather than silently dropped") {
+ // Unreachable through SQL: PushDownUtils screens a filter against the attributes the scan
+ // declared before handing it over. Calling the scan directly is what pins the guard, which has
+ // to fail rather than drop the filter, since a fully pushed predicate dropped here would be
+ // evaluated nowhere. The reference is matched the way the rest of the scan matches partition
+ // column names, so a mixed-case reference binds rather than being rejected.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT f.id, f.part FROM fact f")
+ df.collect()
+ val scan = fileScanOf(df)
+ val e = intercept[SparkException] {
+ scan.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(AttributeReference("id", LongType)(), Literal(1L))))
+ }
+ assert(e.getCondition === "INTERNAL_ERROR")
+ assert(e.getMessage.contains("can only apply a runtime filter over its read partition"),
+ s"unexpected message: ${e.getMessage}")
+
+ val mixedCase = scan.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(AttributeReference("PART", IntegerType)(), Literal(7))))
+ assert(mixedCase.flatMap(
+ _.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0))).distinct.toSeq
+ === Seq(7))
+ }
+ }
+ }
+
+ test("a fully pushed filter stays exact when one FilePartition spans partition directories") {
+ // The fully-pushed contract used to ask for the attribute's value to be fixed within every
+ // InputPartition. A file scan does not satisfy that: FilePartition packs by size and can hold
+ // files from several partition directories. What makes the declaration sound is that the
+ // predicate chose those directories, so every row in the partition satisfies it even though the
+ // value varies. Pack the surviving directories into one FilePartition to exercise that.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1g",
+ SQLConf.FILES_MIN_PARTITION_NUM.key -> "1",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "0") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ // min(dim_id) over dim_val > 7 is 8, so this keeps partitions 8 and 9.
+ val df = sql(
+ """SELECT f.id FROM fact f
+ |WHERE f.part >= (SELECT min(dim_id) FROM dim WHERE dim_val > 7)""".stripMargin)
+ checkAnswer(df, (0 until 10).flatMap(i => Seq(Row(i * 10 + 8), Row(i * 10 + 9))))
+ val factScan = factScanOf(df)
+ assert(factScan.filteredPartitions.flatten.size == 1,
+ "expected both surviving partition directories to pack into one FilePartition, got " +
+ s"${factScan.filteredPartitions.flatten.size}")
+ val numOutputRows = factScan.metrics("numOutputRows").value
+ assert(numOutputRows == 20,
+ s"expected that one FilePartition to carry both directories' 20 rows, got " +
+ s"$numOutputRows")
+ assert(collect(df.queryExecution.executedPlan) { case f: FilterExec => f }.isEmpty,
+ "expected the partition filter to be fully pushed, leaving no FilterExec")
+ }
+ }
+ }
+ }
+
+ test("runtime pruning packs the surviving files for their own size") {
+ // `partitions` sized its bins for the whole file set, so handing those bins back with files
+ // removed leaves a heavily pruned scan running the few tasks the unpruned set needed. Selecting
+ // the same partitions at compile time packs for the pruned set, so the two layouts must agree.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1g",
+ SQLConf.FILES_MIN_PARTITION_NUM.key -> "3",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "0") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val expected = (0 until 10).flatMap(i => Seq(Row(i * 10 + 8, 8), Row(i * 10 + 9, 9)))
+ // min(dim_id) over dim_val > 7 is 8, so this keeps partitions 8 and 9.
+ val runtime = sql(
+ """SELECT f.id, f.part FROM fact f
+ |WHERE f.part >= (SELECT min(dim_id) FROM dim WHERE dim_val > 7)""".stripMargin)
+ checkAnswer(runtime, expected)
+ val compileTime = sql("SELECT f.id, f.part FROM fact f WHERE f.part >= 8")
+ checkAnswer(compileTime, expected)
+ // What gives the comparison teeth: the whole file set packs two or more of the surviving
+ // files into one bin, so handing those bins back with the other files removed cannot
+ // produce as many partitions as packing the survivors does. That needs the fixture to
+ // write more than one file per partition directory, and this is the assertion that fails
+ // if it stops doing so.
+ val unpruned = sql("SELECT f.id, f.part FROM fact f")
+ unpruned.collect()
+ val unprunedPartitions = factScanOf(unpruned).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val binsHoldingSurvivors =
+ unprunedPartitions.count(_.files.exists(_.partitionValues.getInt(0) >= 8))
+ val survivorsWhenUnpruned =
+ unprunedPartitions.flatMap(_.files).count(_.partitionValues.getInt(0) >= 8)
+ assert(binsHoldingSurvivors < survivorsWhenUnpruned,
+ s"expected the whole set to pack $survivorsWhenUnpruned surviving files into fewer " +
+ s"than $survivorsWhenUnpruned bins, got $binsHoldingSurvivors")
+ val compileTimePartitions = factScanOf(compileTime).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val runtimePartitions = factScanOf(runtime).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val runtimeFiles = runtimePartitions.flatMap(_.files).size
+ // The teeth above were measured on the unpruned scan; this is what ties them to this arm.
+ assert(runtimeFiles === survivorsWhenUnpruned,
+ s"expected this arm to read the $survivorsWhenUnpruned files the unpruned packing " +
+ s"held, got $runtimeFiles")
+ // Both arms must see the same files: the runtime path cannot re-split, so a compile-time
+ // maxSplitBytes below one file's length would cut that arm's files finer and the counts
+ // would differ for a reason this test is not about.
+ assert(runtimeFiles === compileTimePartitions.flatMap(_.files).size,
+ s"expected both arms to read $runtimeFiles files, got " +
+ s"${compileTimePartitions.flatMap(_.files).size} at compile time")
+ // The two arms pack from differently ordered input: `partitions` sorts within a partition
+ // directory, this path sorts the survivors globally, and Next Fit Decreasing is sensitive
+ // to that. What makes the counts comparable is that under this conf no two surviving
+ // files fit in one bin, so both arms give one partition per file whatever the order.
+ assert(runtimePartitions.size === runtimeFiles,
+ s"expected one partition per surviving file, got ${runtimePartitions.size} " +
+ s"for $runtimeFiles files")
+ assert(runtimePartitions.size === compileTimePartitions.size,
+ "expected the runtime-pruned scan to be packed like the compile-time one " +
+ s"(${compileTimePartitions.size} partitions), got ${runtimePartitions.size}")
+ }
+ }
+ }
+ }
+
+ test("a DPP filter degraded to true leaves the v2 file scan unpruned") {
+ // With reuse-only DPP and no broadcast to reuse, the subquery is replaced by true before
+ // execution. What this pins is that such a filter still reaches the scan and the query stays
+ // correct; the row count only records that nothing was pruned, which is what a true predicate
+ // does whether or not `catalystRuntimeFilters` drops it.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ 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()
+ assert(rows.length == 10, s"expected 10 rows after join+filter, got ${rows.length}")
+ val factScan = factScanOf(df)
+ assert(factScan.runtimeFilters.exists {
+ case DynamicPruningExpression(l: Literal) => l == Literal.TrueLiteral
+ case _ => false
+ }, s"expected a degraded DPP filter, got ${factScan.runtimeFilters}")
+ val numOutputRows = factScan.metrics("numOutputRows").value
+ assert(numOutputRows == 100,
+ s"expected the fact scan to read all 100 rows, got $numOutputRows")
+ }
+ }
+ }
+ }
+
+ test("a runtime filter binds to the read partition columns when only some of them are read") {
+ // `partitions` projects each file's partition values down to `readPartitionSchema` when the
+ // query reads a strict subset of the partition columns, so a filter on `p2` has to read slot 0
+ // of those values. Binding by the table's partition schema would read slot 1, which in the
+ // projected values holds `p3`: rows that satisfy the filter are lost, and with the filter
+ // fully pushed, rows that do not are returned.
+ withDppV2Conf {
+ withTempDir { dir =>
+ val tPath = new File(dir, "t").getCanonicalPath
+ spark.range(100).selectExpr("id", "id % 2 AS p1", "id % 5 AS p2", "id % 3 AS p3")
+ .write.format("parquet").partitionBy("p1", "p2", "p3").save(tPath)
+ spark.read.format("parquet").schema("id long, p1 int, p2 int, p3 int").load(tPath)
+ .createOrReplaceTempView("t")
+ writeFactAndDim(dir)
+ val expected = (0 until 100).filter(_ % 5 == 2).map(i => Row(i.toLong, i % 3))
+ Seq(
+ """SELECT id, p3 FROM t
+ |WHERE p2 = (SELECT max(dim_id) FROM dim WHERE dim_val = 2)""".stripMargin,
+ """SELECT t.id, t.p3 FROM t JOIN dim d
+ |ON t.p2 = d.dim_id WHERE d.dim_val = 2""".stripMargin).foreach { query =>
+ val df = sql(query)
+ checkAnswer(df, expected)
+ val scan = factScanOf(df)
+ assert(scan.scan.asInstanceOf[FileScan].readPartitionSchema.fieldNames.toSeq ===
+ Seq("p2", "p3"), "the scan must read a strict subset of the partition columns")
+ val numOutputRows = scan.metrics("numOutputRows").value
+ assert(numOutputRows == expected.size,
+ s"expected the scan to read ${expected.size} rows, got $numOutputRows")
+ }
+ }
+ }
+ }
+
+ test("SPARK-38148: no DPP over a partition column a static equality already fixes") {
+ // `FileScanBuilder` pushes the static partition filter fully into the scan, so no `Filter`
+ // above the scan carries it; the cleanup has to read it from the scan relation's pushed
+ // filters, as it reads it from the `Filter` over a V1 relation. V1 is the reference.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ // An int join key, so the pruning key is `part` itself rather than a cast of it.
+ val dimIntPath = new File(dir, "dim_int").getCanonicalPath
+ spark.range(10).selectExpr("CAST(id AS INT) AS dim_id", "CAST(id AS INT) AS dim_val")
+ .write.format("parquet").save(dimIntPath)
+ // (condition, the partitions it keeps, whether DPP stays, whether V1 is checked). The first
+ // six are the SPARK-38148 cases. The last two do not fix the key, so V2 keeps the DPP; they
+ // are checked under V2 only, because V1 reads an equality anywhere in its filter and drops
+ // the DPP.
+ Seq[(String, Int => Boolean, Boolean, Boolean)](
+ ("f.part = 3", _ == 3, false, true),
+ ("3 = f.part", _ == 3, false, true),
+ ("f.part <=> 3", _ == 3, false, true),
+ ("3 <=> f.part", _ == 3, false, true),
+ ("f.part > 3", _ > 3, true, true),
+ ("5 > f.part", _ < 5, true, true),
+ ("f.part <> 3", _ != 3, true, false),
+ ("(f.part = 3 OR f.part = 4)", p => p == 3 || p == 4, true, false)
+ ).foreach { case (condition, keeps, hasDpp, checkV1) =>
+ val query =
+ s"""SELECT /*+ broadcast(d) */ f.id FROM fact f JOIN dim_int d
+ |ON f.part = d.dim_id WHERE $condition AND d.dim_val > 2""".stripMargin
+ val expected = (0 until 100).filter { i =>
+ val part = i % 10
+ part > 2 && keeps(part)
+ }.map(i => Row(i.toLong))
+ (if (checkV1) Seq("parquet", "") else Seq("")).foreach { v1Sources =>
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> v1Sources) {
+ // A temp view keeps the relation it was created with, so create both under this conf.
+ spark.read.format("parquet").schema("id long, part int")
+ .load(new File(dir, "fact").getCanonicalPath).createOrReplaceTempView("fact")
+ spark.read.format("parquet").load(dimIntPath).createOrReplaceTempView("dim_int")
+ val df = sql(query)
+ val optimized = df.queryExecution.optimizedPlan
+ assert(optimized.exists(_.isInstanceOf[DataSourceV2ScanRelation]) ===
+ v1Sources.isEmpty, s"wrong source version for '$v1Sources':\n$optimized")
+ assert(collectDppFilters(optimized).nonEmpty === hasDpp,
+ s"`$condition` with v1 sources '$v1Sources': expected DPP $hasDpp, got\n$optimized")
+ checkAnswer(df, expected)
+ }
+ }
+ }
+ }
+ }
+ }
+
+ Seq("orc", "json", "csv").foreach { format =>
Review Comment:
Done in e721ae7. "DPP and a scalar subquery prune input partitions at runtime for v2 text" uses its own single-column text fixture, and `AvroV2Suite` has "SPARK-30628: Avro source v2: runtime filters prune partitions". Both check that only partition 7 is read.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "dongjoon-hyun (via GitHub) via reviews" <re...@spark.apache.org>.
dongjoon-hyun commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-6666936891
Sure, @LuciferYang .
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4489844535
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionDirectory}
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep, so `partitions` is not used. A scan that overrides it or `planInputPartitions()`, or that
+ * reports a partitioning or an ordering, should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val listingRef = new AtomicReference[Seq[PartitionDirectory]]()
+ @transient private lazy val allPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private def listing: Seq[PartitionDirectory] =
+ once(listingRef)(fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private def allPartitions: Seq[FilePartition] = once(allPartitionsRef)(filePartitions(listing))
+
+ private def once[T >: Null <: AnyRef](ref: AtomicReference[T])(compute: => T): T =
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, compute)
+ ref.get()
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = allPartitions.toArray
Review Comment:
Done in a76b80e435f. The trait now declares `final override protected def partitions = allPartitions()`, and `planInputPartitions()` is final as well, since overriding it skips the runtime-filtered path in the same way. A subclass of a built-in scan can still override `filePartitions`, and both read paths go through it.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r3220495968
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -186,6 +188,21 @@ trait FileScan extends Scan
partitions.toArray
}
+ /**
+ * SPARK-30628: produce InputPartitions taking additional runtime filters into account.
+ * Called by `BatchScanExec.filteredPartitions` with DPP and scalar-subquery filters that
+ * must apply on top of the scan's compile-time `partitionFilters`. Returns the same
+ * partitions as `planInputPartitions()` when `extraFilters` is empty.
+ */
+ def planInputPartitionsWithRuntimeFilters(
+ extraFilters: Seq[Expression]): Array[InputPartition] = {
+ if (extraFilters.isEmpty) {
+ planInputPartitions()
+ } else {
+ partitionsImpl(partitionFilters ++ extraFilters).toArray
Review Comment:
What happens for the existing 3rd-party DSv2 data sources which overrides `FileScan.partitions` already?
Could you add an explicit comment that we need to override `partitionsImpl` for runtime-filter path?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r3223459898
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -186,6 +188,21 @@ trait FileScan extends Scan
partitions.toArray
}
+ /**
+ * SPARK-30628: produce InputPartitions taking additional runtime filters into account.
+ * Called by `BatchScanExec.filteredPartitions` with DPP and scalar-subquery filters that
+ * must apply on top of the scan's compile-time `partitionFilters`. Returns the same
+ * partitions as `planInputPartitions()` when `extraFilters` is empty.
+ */
+ def planInputPartitionsWithRuntimeFilters(
+ extraFilters: Seq[Expression]): Array[InputPartition] = {
+ if (extraFilters.isEmpty) {
+ planInputPartitions()
+ } else {
+ partitionsImpl(partitionFilters ++ extraFilters).toArray
Review Comment:
Good catch — answers below, addressed in d544f201629:
1. Existing third-party `FileScan` subclasses that override the `protected def partitions` keep working. The compile-time path (`planInputPartitions`) still calls `partitions`, so their override is honored. The new `planInputPartitionsWithRuntimeFilters` only diverges from that path when `extraFilters` is non-empty — when it's empty, the method delegates to `planInputPartitions()`. So existing sources see no behavior change; they simply won't benefit from the runtime-filter pruning unless they migrate to override `partitionsImpl`.
2. Added Scaladoc on `partitions`, `partitionsImpl`, and `planInputPartitionsWithRuntimeFilters` documenting that subclasses should override `partitionsImpl` so the customization applies to both code paths.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "pan3793 (via GitHub)" <gi...@apache.org>.
pan3793 commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-6312942854
+1 for Option A. Moving SupportsRuntimeCatalystFiltering off FileScan keeps the one-interface rule intact and leaves existing FileScan subclasses on their pre-PR behavior, while the built-in formats keep the feature. Please move plannedPartitions into the new trait with the three methods. B + precedence reopens the one-interface rule for little gain.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4454062991
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala:
##########
@@ -72,35 +80,38 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty
/**
- * Filters this scan using runtime Catalyst expressions.
+ * Plans this scan's input partitions with runtime Catalyst expressions applied.
+ *
+ * The provided expressions must be interpreted as a set of predicates that are ANDed together,
+ * and applied on top of the filters already pushed at query compilation. Implementations may use
+ * them to plan fewer [[InputPartition]]s than [[org.apache.spark.sql.connector.read.Batch]]
+ * would without them.
*
- * The provided expressions must be interpreted as a set of predicates that are ANDed together.
- * Implementations may use the expressions to prune initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s.
+ * Spark calls this method at execution time, and only when at least one runtime filter survives
+ * screening; the partitions it returns replace the ones that scan node would otherwise have
+ * read. Spark may also call `Batch.planInputPartitions()` on the same scan, so neither method may
+ * assume it is the only one used. One scan instance can back several scan nodes (e.g. the two
+ * branches of a group-based UPDATE), each with its own runtime filters, so an implementation
+ * must derive its result from the given expressions alone and must not carry state over from a
+ * previous call. A group-based scan whose write replaces the groups that were read has to record
+ * what each call returned for that write, and count a scan node that read through
+ * `planInputPartitions()` as having read every group that method returns. Spark may also call
+ * that method while planning a scan node, so such a call alone does not show that the node read
+ * through it.
*
* Spark currently tracks runtime-filter eligibility by root attribute. If [[filterAttributes]]
* returns a nested reference, an expression may access another nested field under the same root.
* The scan must match each access against its own partition layout and use only expressions it
* can apply.
*
- * Spark may call this method more than once for the same scan instance: a plan can hold several
- * scan nodes sharing one scan (e.g. the two branches of a group-based UPDATE), and each pushes
- * its own copy of the runtime filters. Implementations must treat successive calls as additive,
- * ANDing the new expressions with those already pushed rather than replacing them.
- *
* If the scan also implements
* [[org.apache.spark.sql.connector.read.SupportsReportPartitioning]], it must preserve
- * the originally reported partitioning during runtime filtering. While applying runtime
- * predicates, the scan may detect that some
- * [[org.apache.spark.sql.connector.read.InputPartition]]s have no matching data, in which
- * case it can either replace the initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s that have no matching data with
- * empty [[org.apache.spark.sql.connector.read.InputPartition]]s, or report only a subset of
- * the original partition values (omitting those with no data) via
- * [[org.apache.spark.sql.connector.read.Batch#planInputPartitions]]. The scan must not
- * report new partition values that were not present in the original partitioning.
- *
- * Note that Spark will call [[Scan.toBatch]] again after filtering the scan at runtime.
+ * the originally reported partitioning. While applying runtime predicates, the scan may detect
+ * that some [[InputPartition]]s have no matching data, in which case it can either replace
+ * them with empty [[InputPartition]]s, or report only a subset of the original partition
+ * values (omitting those with no data). The scan must not report new partition values that
+ * were not present in the original partitioning.
*/
- def filter(expressions: Array[Expression]): Unit
+ def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition]
Review Comment:
Done in e721ae7. The file scans no longer declare `fullyPushedFilterAttributes()`, so a scalar-subquery filter keeps its `FilterExec`, and a node that replaces `BatchScanExec` without pruning still returns the right rows. Since every runtime filter is now evaluated again above the scan, a filter that fails to evaluate on a directory keeps it, whether it is a DPP filter or a scalar-subquery filter. "a scalar subquery filter on a partition column is evaluated again after the scan" checks the `FilterExec`. `PushDownUtils.pushRuntimeFilters` now also returns false for these scans instead of failing, so a node that pushes through it and then plans the scan reads every partition and the join or the `FilterExec` drops the extra rows.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
Review Comment:
Done in e721ae7. `TextBasedFileScan` is back to what master has, and `CSVScan`, `JsonScan` and `TextScan` mix in the trait themselves.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/TextBasedFileScan.scala:
##########
@@ -27,7 +27,7 @@ import org.apache.spark.util.Utils
abstract class TextBasedFileScan(
sparkSession: SparkSession,
- options: CaseInsensitiveStringMap) extends FileScan {
+ options: CaseInsensitiveStringMap) extends RuntimePartitionPruningFileScan {
Review Comment:
This went away in 7158934. The runtime path now filters the listed directories and passes them to `filePartitions`, which sums the length and open cost per file, as V1 does.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
+ .flatMap(_.files.filter(file => predicate.eval(file.partitionValues)))
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+ // Pack the survivors instead of handing back the planned partitions with files removed:
+ // `partitions` sized those bins for the whole file set, so a scan that pruned most of its files
+ // would run the same few tasks over a fraction of the data. V1 packs the pruned set the same
+ // way, from `dynamicallySelectedPartitions`.
+ val openCostInBytes = conf.filesOpenCostInBytes
+ val maxSplitBytes = FilePartition.maxSplitBytes(
+ sparkSession, keptFiles.map(_.length + openCostInBytes).sum)
Review Comment:
Done in e721ae7, with one change from your snippet. Since 7158934 the index is into `fileIndex.partitionSchema`, the schema the directory values follow, so the type comes from there rather than from `readPartitionSchema`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-6679165691
@dongjoon-hyun thanks for the review. Items 9-23 are addressed in a76b80e435f, on top of a merge of master (6a5338ee6f1) that brings in SPARK-60752. For item 12 I changed `PartitionPruning` rather than only noting it. For item 13 the description now has the bloom-filter bullet, and the `RuntimeFilters` line is listed for every V2 scan, including connector scans.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-5768680813
Thanks for picking this up, @LuciferYang. I read through the diff and the surrounding paths, and the correctness core looks sound to me. Two comments plus a nit on the description.
### 1. Runtime filters prune after bin-packing, which costs parallelism
`planInputPartitionsWithRuntimeFilters` filters files out of `FilePartition`s that `partitions` already bin-packed, and `maxSplitBytes` was derived there from the *unpruned* file set. V1 does the opposite: `FileSourceScanExec.createReadRDD` recomputes `FilePartition.maxSplitBytes` from `dynamicallySelectedPartitions` and bin-packs only the survivors.
Concretely: a 12.8 GB fact table over 10 partition directories, with the default `spark.sql.files.maxPartitionBytes` (128 MB) and leaf parallelism 200, gives `maxSplitBytes = min(128MB, 12.8GB/200) = 64MB` and ~200 `FilePartition`s. A DPP filter selecting one directory (1.28 GB) leaves the ~20 bins that covered it, so the scan runs ~20 tasks of 64 MB. The V1 plan for the same query recomputes `maxSplitBytes` from 1.28 GB (6.4 MB) and produces ~200 tasks -- roughly a 10x parallelism difference. Bins that straddle two directories are additionally left under-filled by `part.copy(files = kept)`.
The same loop also evaluates the predicate once per `PartitionedFile` rather than once per partition directory, so the driver does one interpreted `InSubqueryExec` lookup per file instead of per distinct partition value.
I see the tradeoff you documented (one listing per scan, honoring whatever a subclass did in `partitions`), and I think both can be kept by re-running `FilePartition.maxSplitBytes` / `FilePartition.getFilePartitions` over the surviving `PartitionedFile`s instead of returning the filtered bins as they are. If you'd rather leave it, it seems worth listing next to the missing pruning metrics in the description.
### 2. Stale comment in the new suite
In `DataSourceV2FileSourceDPPSuite`, the "a runtime filter is ANDed with the compile-time partition filters" test explains the mechanism as "`buildPartitions` gets `partitionFilters ++ expressions`", but there is no `buildPartitions` anywhere in the tree. The actual mechanism is that `planInputPartitionsWithRuntimeFilters` filters `plannedPartitions`, which `partitions` built with `partitionFilters` applied through `fileIndex.listFiles`. Looks like a leftover from an earlier iteration.
### Nit on the description
> Such a subquery filter is no longer evaluated above the scan, so the scan is its only evaluator and a file whose partition value cannot be evaluated (an ANSI cast, an overflow) fails the query rather than being filtered out row by row
This reads as more user-visible than it is. Whenever the partition directory holds rows, the pre-change plan evaluates the same cast in the `FilterExec` above the scan (or in the join key) and fails identically -- the `fact_str` test in this PR is exactly that shape. The genuinely new exposure looks narrow (e.g. a partition directory whose files carry no rows).
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4173545701
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
Review Comment:
Done in 67e88d6. The memo is an `AtomicReference` set with `compareAndSet`, so `partitions` runs outside the scan's monitor; the remaining `lazy val` only creates the empty reference. Two threads asking first at the same time may both list, and both get the stored result. `partitions runs without holding the scan's monitor` records `Thread.holdsLock(this)` inside an overridden `partitions` for both read paths, and fails against the old `lazy val`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
Review Comment:
Done in 2acbef9. For a `DataSourceV2ScanRelation`, `CleanupDynamicPruningFilters` also reads equalities from the top-level conjuncts of `pushedFilters`, so `part <> 3` or `part = 3 OR part = 4` keeps the DPP. A new test runs the six SPARK-38148 conditions on V1 and V2 parquet and expects the same DPP decision and rows on both, then checks those two extra conditions on V2. This applies to any V2 source that fully pushes such an equality: the result is the same, and the useless DPP is gone.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -182,8 +185,80 @@ trait FileScan extends Scan
FilePartition.getFilePartitions(sparkSession, splitFiles, maxSplitBytes)
}
+ /**
+ * The partitions `partitions` planned, computed once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ */
+ @transient private lazy val plannedPartitions: Seq[FilePartition] = partitions
+
override def planInputPartitions(): Array[InputPartition] = {
- partitions.toArray
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
Review Comment:
Agreed. I'll add it to the user-facing section when I next update the description.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala:
##########
@@ -862,21 +883,17 @@ private class PartitioningBreakingScan(
afterFilter: Seq[InputPartition])
extends Scan with Batch with SupportsRuntimeCatalystFiltering {
- private var filtered = false
-
override def readSchema(): StructType = new StructType().add("part", IntegerType)
override def toBatch: Batch = this
- override def planInputPartitions(): Array[InputPartition] =
- if (filtered) afterFilter.toArray else initialPartitions.toArray
+ override def planInputPartitions(): Array[InputPartition] = initialPartitions.toArray
Review Comment:
Done in 2acbef9. It is now `PartitioningBreakingScan(afterFilter) extends SupportsRuntimeCatalystFiltering`, without the `Batch` members and their imports.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "pan3793 (via GitHub)" <gi...@apache.org>.
pan3793 commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-6289031659
@LuciferYang, haven't gone through the whole conversation, but I see there is discussion about Kyuubi's `HiveScan`. From my view, `FileScan` is Spark's private API, as a downstream project, Kyuubi takes its own risk in consuming it, no need to worry about compatibility when making changes upstream.
I will take a closer look later
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "dongjoon-hyun (via GitHub) via reviews" <re...@spark.apache.org>.
dongjoon-hyun commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4486420149
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -293,11 +346,19 @@ object PushDownUtils extends Logging {
output: Seq[AttributeReference],
keyedPartitioning: Option[KeyedPartitioning],
originalPartitions: => Seq[InputPartition]): Seq[Option[InputPartition]] = {
- val filtered = pushRuntimeFilters(scan, runtimeFilters, table, output)
+ val (filtered, newPartitions) = scan match {
+ // A scan implementing both runtime-filtering interfaces falls through to pushRuntimeFilters,
+ // which rejects it. The scan is asked even when screening leaves no filter.
+ case catalystScan: SupportsRuntimeCatalystFiltering
Review Comment:
**10. [Low] Correction to my item 1: a 4.3.0 scan now re-plans when screening removes every runtime filter**
The default I suggested in [item 1](https://github.com/apache/spark/pull/55579#discussion_r4449362227) matches 4.3.0 only when a filter survives screening, as its scaladoc now says (SupportsRuntimeCatalystFiltering.scala:113-115). This branch also calls it with an empty array, and the default then skips `filter` but still plans again through `toBatch.planInputPartitions()` (SupportsRuntimeCatalystFiltering.scala:120-121). In that case 4.3.0 returned false from `pushRuntimeFilters` and read `originalPartitions` ([v4.3.0](https://github.com/apache/spark/blob/v4.3.0/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala#L263-L268), 2026-10-27).
This happens under the defaults. With `spark.sql.optimizer.dynamicPartitionPruning.reuseBroadcastOnly=true` and a sort-merge join on the scan's partition column, `PlanDynamicPruningFilters` degrades the DPP filter to `DynamicPruningExpression(true)` and `catalystRuntimeFilters` drops it (L290). Each such scan node then pays one more planning, which can be a file or manifest listing for a lakehouse-style connector, and under a `KeyedPartitioning` the re-plan also goes through the checks below. The test at DataSourceV2CatalystRuntimeFilterSuite.scala:568-579 cannot see this, because it passes `scan.toBatch.planInputPartitions().toSeq` as `originalPartitions` and its Scala scan has a mixin forwarder. The `@param originalPartitions` at L336 also no longer says when they are read.
Could this guard add `&& !implementsOnlyDeprecatedFilter(catalystScan)`? Such a scan then takes the deprecated route of `pushRuntimeFilters` (L241-243), which already does what 4.3.0 did. A `JavaDeprecatedFilterScan` test with a sentinel `originalPartitions` would pin it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionDirectory}
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep, so `partitions` is not used. A scan that overrides it or `planInputPartitions()`, or that
+ * reports a partitioning or an ordering, should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val listingRef = new AtomicReference[Seq[PartitionDirectory]]()
+ @transient private lazy val allPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private def listing: Seq[PartitionDirectory] =
+ once(listingRef)(fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private def allPartitions: Seq[FilePartition] = once(allPartitionsRef)(filePartitions(listing))
+
+ private def once[T >: Null <: AnyRef](ref: AtomicReference[T])(compute: => T): T =
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, compute)
+ ref.get()
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = allPartitions.toArray
Review Comment:
**9. [Low] A subclass of a built-in scan that narrows `partitions` now reads every file**
Since 7158934 neither read path calls `partitions`, as the scaladoc above says, but the six built-in scans still let a subclass override it. A subclass like `new ParquetScan(...) { override protected def partitions = super.partitions.map(p => p.copy(files = p.files.filter(keep))).filter(_.files.nonEmpty) }` still compiles, yet even a plain `SELECT * FROM t` now plans through this method and `filePartitions(listing)` (L48-51), so the override never runs and the files it excluded are read without any error. On master, `planInputPartitions()` reads `partitions`. bfc057d documented this case ("A subclass still customizes `partitions`") and tested it ("a subclass that narrows the file set in partitions keeps it under a runtime filter"), and 7158934 removed both. Such a subclass cannot follow "should not mix this in", because the built-in scan already mixes it in. I did not find one on GitHub, so this is a contract gap rather than a known break.
Could the trait declare `final override protected def partitions: Seq[FilePartition] = allPartitions`, so that the silent bypass becomes a compile error? Nothing in-tree overrides `partitions` on these scans, and the suite's anonymous `ParquetScan` overrides `filePartitions`, which stays open. Otherwise, could the PR description say so?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionDirectory}
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep, so `partitions` is not used. A scan that overrides it or `planInputPartitions()`, or that
+ * reports a partitioning or an ordering, should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val listingRef = new AtomicReference[Seq[PartitionDirectory]]()
+ @transient private lazy val allPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private def listing: Seq[PartitionDirectory] =
+ once(listingRef)(fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private def allPartitions: Seq[FilePartition] = once(allPartitionsRef)(filePartitions(listing))
+
+ private def once[T >: Null <: AnyRef](ref: AtomicReference[T])(compute: => T): T =
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, compute)
+ ref.get()
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = allPartitions.toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ if (expressions.isEmpty) return allPartitions.toArray
+ val readFields = readPartitionSchema.fieldNames.map(normalizePartitionName).toSet
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => readFields.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ // Directory values follow the table's partition schema, not the read one, and bind as V1
+ // binds them: with that schema's types, nullable.
+ val partitionSchema = fileIndex.partitionSchema
+ val fieldIndex = partitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ val predicates = expressions.toSeq.map { e =>
+ Predicate.createInterpreted(e.transform {
+ case a: Attribute =>
+ val i = fieldIndex(normalizePartitionName(a.name))
+ BoundReference(i, partitionSchema(i).dataType, nullable = true)
+ })
+ }
+ val kept = listing.filter { dir =>
+ predicates.forall { predicate =>
+ try predicate.eval(dir.values) catch { case NonFatal(_) => true }
+ }
+ }
+ filePartitions(kept).toArray
Review Comment:
**11. [Low] A runtime filter that keeps every directory packs the files again**
When a non-empty runtime filter keeps every listed directory, this line splits and packs the whole listing a second time, although `allPartitions` already holds the same packing. Planning forces it for every built-in file scan: `DataSourceV2Strategy` asks `supportsColumnar`, which under the default `PARTITION_DEFINED` mode reads `inputPartitions`, i.e. L59. Some runtime filters do keep everything, such as a DPP filter whose build side still holds every partition value, or a scalar subquery filter such as `part >= (SELECT min(part) FROM ...)`. The driver then reruns `maxSplitBytes`, `splitFiles` with its block-location lookups, the per-directory sort and the packing over every file before the stage launches, and `filteredPartitions` keeps a second set of `PartitionedFile`s for the plan's lifetime. The merge-base and V1 pack once.
Could this return the memo when nothing was pruned, e.g. `if (kept.length == listing.length) return allPartitions.toArray` (`filter` keeps the order)? A keep-all variant of the `packings === 1` test at DataSourceV2FileSourceDPPSuite.scala:350-375 would pin it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionDirectory}
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep, so `partitions` is not used. A scan that overrides it or `planInputPartitions()`, or that
+ * reports a partitioning or an ordering, should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val listingRef = new AtomicReference[Seq[PartitionDirectory]]()
+ @transient private lazy val allPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private def listing: Seq[PartitionDirectory] =
+ once(listingRef)(fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private def allPartitions: Seq[FilePartition] = once(allPartitionsRef)(filePartitions(listing))
+
+ private def once[T >: Null <: AnyRef](ref: AtomicReference[T])(compute: => T): T =
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, compute)
+ ref.get()
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = allPartitions.toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
Review Comment:
**12. [Low] A partitioned V2 file dim filtered only on its partition column gets no DPP**
Declaring the partition columns here is what brings V2 file facts into DPP, but `PartitionPruning.hasSelectivePredicate` (PartitionPruning.scala:234-237) only counts a likely-selective `Filter` on the filtering side, and `FileScanBuilder` takes partition-column predicates into the scan without leaving a `Filter` above it. So with `useV1SourceList` empty, `fact` partitioned by `part` and `dim(k, dp)` partitioned by `dp`, `SELECT f.id FROM fact f JOIN dim d ON f.part = d.k WHERE d.dp = 1` keeps only `isnotnull(k)` on the dim side. No DPP filter is inserted, and the fact scan reads every partition. With `useV1SourceList=parquet`, `PruneFileSourcePartitions` keeps that `Filter` (PruneFileSourcePartitions.scala:92-93, "Keep partition-pruning predicates so that they are visible in physical planning") and DPP fires. This is performance only, and the mechanism already exists on master for any DSv2 filtering side, but every dim in the new tests is unpartitioned. PartitionPruning.scala is o
utside this diff, so I left this here.
Could the description list it next to the pruning-metrics difference, or could `hasSelectivePredicate` read the pushed filters as `CleanupDynamicPruningFilters` now does, e.g. `case r: DataSourceV2ScanRelation if r.scan.isInstanceOf[RuntimePartitionPruningFileScan] => r.pushedFilters.exists(isLikelySelective)`, with a test against V1?
##########
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:
**18. [Nit] The `= Array.empty` overrides in the error-path fixtures are dead**
Now that `planInputPartitionsWithRuntimeFilters` has a default (e721ae7), this override and the ones at L1017, L1031, L1047 and L1059 are never reached: those tests fail in `runtimeFilterAttrs` before anything is planned, and this dual-interface scan is rejected in `pushRuntimeFilters` (PushDownUtils.scala:197). They also make each fixture's diff read as a rewrite rather than a plain removal of the no-op `filter`. Could they go, as `NoRuntimeFilterMethodScan` (L983-994) shows they can?
##########
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
Review Comment:
**15. [Nit] The union test would pass without the strategy-level drop**
The title and L265-267 say that letting this filter through means an INTERNAL_ERROR. But with `usableDynamicFilters` (DataSourceV2Strategy.scala:239-243) reverted, `catalystRuntimeFilters` still screens the filter over `other` against `fact_other`'s empty `filterAttributes()` (PushDownUtils.scala:281-290), and RuntimePartitionPruningFileScan.scala:82 returns `allPartitions` for the empty array, so only the `runtimeFilters.isEmpty` assertion at L296 would fail. The claim held for the execution-time screen this test was added for in fd8cd2d, and went stale when 2acbef9 retargeted it at the strategy. Could the title say e.g. "... is dropped at planning", with the strategy comment's reason: keeping it would plan its subquery for nothing and leave a dead filter in EXPLAIN?
##########
connector/avro/src/main/scala/org/apache/spark/sql/v2/avro/AvroScan.scala:
##########
@@ -41,7 +41,7 @@ case class AvroScan(
options: CaseInsensitiveStringMap,
pushedFilters: Array[Filter],
partitionFilters: Seq[Expression] = Seq.empty,
- dataFilters: Seq[Expression] = Seq.empty) extends FileScan {
+ dataFilters: Seq[Expression] = Seq.empty) extends RuntimePartitionPruningFileScan {
Review Comment:
**21. [Nit] Four comments still point at `FileScan.partitions`**
The comment at L53, like the ones at CSVScan.scala:86, JsonScan.scala:79 and OrcScan.scala:74, says the partition values are truncated in `FileScan.partitions`. These scans now plan through `RuntimePartitionPruningFileScan`, which never calls `partitions`, and the truncation happens in `FileScan.filePartitions`. L53 is outside this diff, so I left this on the line that mixes the trait in. Could the four comments say `FileScan.filePartitions`?
##########
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:
**17. [Nit] The packing count needs a full `ParquetScan` subclass**
This subclasses `ParquetScan` through all 12 constructor fields only to count `filePartitions` calls, so the next rename, retype or new required field (as SPARK-54656 did with the last one) breaks the compilation of this suite. The memo hands back the same `FilePartition` objects and every packing allocates new ones, so an identity check pins the same property, and the title literally:
```scala
val scan = fileScanOf(df)
val planned = scan.planInputPartitions()
assert(planned.nonEmpty && scan.planInputPartitionsWithRuntimeFilters(Array.empty).toSeq
.corresponds(planned.toSeq)(_ eq _))
```
Could it use that instead?
##########
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.
Review Comment:
**23. [Nit] The read schema is explicit for another reason**
The fixture's own schemas, `part int` (L79) and `dim_id long` (L84), put a cast on every `f.part = d.dim_id` join key, so every DPP test in this suite prunes with the key `cast(part as bigint)`, and DPP still fires; `dim_int` at L698 exists precisely to make the key `part` itself. What csv needs the schema for is that the files are written without a header. A reader who trusts this comment and changes the schema to `part long` would silently move every DPP test from a cast key to a bare attribute key. Could it say that instead, e.g. "The read schema is given explicitly because csv is written without a header; it also keeps the join key `cast(part as bigint)` in every format"?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala:
##########
@@ -785,22 +793,37 @@ abstract class InMemoryBaseTable(
private val catalystPredicates = ArrayBuffer.empty[CatalystExpression]
private var filterCalls = 0
+ private val plannedPartitions = ArrayBuffer.empty[Seq[InputPartition]]
- override def filter(expressions: Array[CatalystExpression]): Unit = {
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[CatalystExpression]): Array[InputPartition] = {
catalystPredicates ++= expressions
filterCalls += 1
+ // The result comes from `expressions` and the unfiltered `data` alone. Spark calls this once
+ // per scan node, and several nodes can share one scan, so carrying one call's pruning into
+ // the next would drop partitions the second node still needs. `catalystPredicates` and
+ // `filterCalls` only record what happened, for test assertions; `plannedPartitions` is read
+ // back through `readPartitions`, which is what the group-based write replaces.
+ val filtered = prunePartitions(self.data, expressions)
Review Comment:
**14. [Low] After merging master, no fixture reproduces SPARK-60752**
[SPARK-60752](https://github.com/apache/spark/commit/71316c632b0) landed on master today, after this PR's merge-base, and merges cleanly. Its scalar-subquery regression test reproduces `Can't zip RDDs with unequal numbers of partitions: List(7, 1)` only in `KeyGroupedPartitioningCatalystRuntimeFilterSuite`, because master's Catalyst fixture prunes `self.data` in place (InMemoryBaseTable.scala:803 on master), while the V2 fixture prunes only an `In` filter. With the non-mutating fixture here, the DPP build-side node returns 1 partition but `data` keeps 7, so the join's copy re-plans 7, and that test, like the Catalyst run of its nested-DPP test, would pass even with `PlannedScan` reverted. The bug class stays live for mutating scans, i.e. `SupportsRuntimeV2Filtering.filter` and the deprecated Catalyst default. The comment at KeyGroupedPartitioningSuite.scala:556 ("which prunes the shared `Scan`") and the `PlannedScan` doc ("A pushed runtime filter replans them from a new `Batch`")
would also go stale.
When you merge master, could `KeyGroupedPartitioningRuntimeFilterTests` also run against a catalog whose Catalyst scan overrides only the deprecated `filter`, or one built on `InMemoryTableWithV2Filter`, which prunes `=` in place?
##########
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
Review Comment:
**22. [Nit] Aggregate pushdown does not need the GROUP BY to cover the partition columns**
`AggregatePushDownUtils` applies the partition-column check only to a non-empty GROUP BY (AggregatePushDownUtils.scala:102-103), so a global aggregate such as the `count(part)` in the test above is pushed with no partition column at all, and that test relies on it. With a GROUP BY the condition is also "exactly the partition columns", not "covers". Could this read "With a GROUP BY, aggregate pushdown only happens when it groups by exactly the partition columns, ..."?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionDirectory}
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep, so `partitions` is not used. A scan that overrides it or `planInputPartitions()`, or that
+ * reports a partitioning or an ordering, should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val listingRef = new AtomicReference[Seq[PartitionDirectory]]()
+ @transient private lazy val allPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private def listing: Seq[PartitionDirectory] =
+ once(listingRef)(fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private def allPartitions: Seq[FilePartition] = once(allPartitionsRef)(filePartitions(listing))
+
+ private def once[T >: Null <: AnyRef](ref: AtomicReference[T])(compute: => T): T =
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, compute)
+ ref.get()
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = allPartitions.toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ if (expressions.isEmpty) return allPartitions.toArray
+ val readFields = readPartitionSchema.fieldNames.map(normalizePartitionName).toSet
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => readFields.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
Review Comment:
**19. [Nit] This check guards input Spark never passes**
`catalystRuntimeFilters` keeps only filters over `filterAttributes()` (PushDownUtils.scala:286-287), a subset of `readPartitionSchema`, so L83-90 cannot fire from SQL, as the direct-call test at DataSourceV2FileSourceDPPSuite.scala:566-591 says ("Unreachable through SQL"). Its reason, that a fully pushed filter would otherwise be applied nowhere, is gone: 7158934 removed the comment saying so, and e721ae7 stopped declaring the columns fully pushed. What remains is a third definition of the applicable columns within L66-100 (`filterAttributes`, `readPartitionSchema` here, `fileIndex.partitionSchema` at L93). Could the check fold into L99 as `fieldIndex.getOrElse(normalizePartitionName(a.name), throw SparkException.internalError(...))`, as `filePartitions` does at FileScan.scala:149-153, or keep a comment pointing at PushDownUtils.scala:286-287?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionDirectory}
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep, so `partitions` is not used. A scan that overrides it or `planInputPartitions()`, or that
+ * reports a partitioning or an ordering, should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val listingRef = new AtomicReference[Seq[PartitionDirectory]]()
+ @transient private lazy val allPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private def listing: Seq[PartitionDirectory] =
+ once(listingRef)(fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private def allPartitions: Seq[FilePartition] = once(allPartitionsRef)(filePartitions(listing))
+
+ private def once[T >: Null <: AnyRef](ref: AtomicReference[T])(compute: => T): T =
Review Comment:
**20. [Nit] `once` and its holders re-implement `TransientBestEffortLazyVal`**
Following up on my lock-free memo ask in [round 4](https://github.com/apache/spark/pull/55579#discussion_r4169230750): core already has one, `org.apache.spark.util.TransientBestEffortLazyVal`. It is lock-free, stores a single result, does not cache a failure and resets on deserialization, and `QueryPlan` (`_outputSet` at QueryPlan.scala:98) and `KeyedPartitioning` (partitioning.scala:887-892) use it for the same Scala 2 lazy-val monitor reason. The two `@transient lazy val` holders here also still take the scan's monitor on first access. Could L41-57 become
```scala
private val listing = new TransientBestEffortLazyVal(() =>
fileIndex.listFiles(partitionFilters, dataFilters))
private val allPartitions = new TransientBestEffortLazyVal(() => filePartitions(listing()))
```
with `listing()` and `allPartitions()` at the call sites? That drops `once`, both holders and the `AtomicReference` import, and the memo tests should keep passing.
##########
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:
**16. [Nit] Three tests repeat others**
As L548-550 say, the scan receives the same input here as in the lower-case test at L425, and L584-588 pin the scan-side name matching, so this test runs the same plan and assertions as L425. "filterAttributes reports the read partition columns" (L446) repeats the assertion at L160 on a scan of the same shape, and "DPP prunes input partitions at runtime for v2 parquet" (L209) is the body of the format loop at L775 for parquet. The first two alone cost 4 query executions and 8 fixture writes across AEOff and AEOn. Could this test and L446 go, and could the loop at L775 cover `"parquet"` too in place of L209, keeping its plan dump in the message?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-6719169474
@dongjoon-hyun thanks for the review. Items 24-38 are addressed in c1ba7eacfd7 and fd6f5adfc85, on top of a merge of master (b286b8a06f4) that brings in SPARK-60702 and resolves the conflict in `StorageFilterTestSource.scala`. Items 24, 26, 28, 30, 34 and 38 change code, items 25 and 35-37 change docs or comments, and items 27, 29, 31 and 33 are in the description, whose bloom-filter bullet now follows SPARK-60702. Items 29 and 32 hold on master too, so they go to SPARK-61173 and SPARK-61174.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r3220519836
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala:
##########
@@ -63,10 +63,19 @@ case class BatchScanExec(
@transient private[sql] lazy val filteredPartitions: Seq[Option[InputPartition]] = {
val originalPartitioning = outputPartitioning
- val filtered = PushDownUtils.pushRuntimeFilters(scan, runtimeFilters, table, output)
+ // SPARK-30628: V2 file scans don't implement SupportsRuntimeV2Filtering, so route
+ // runtime filters through FileScan.planInputPartitionsWithRuntimeFilters as Catalyst
+ // expressions. This preserves DynamicPruningExpression and scalar-subquery semantics
+ // that V2-Predicate translation would drop.
+ val (filtered, newPartitions) = scan match {
+ case fs: FileScan if runtimeFilters.nonEmpty =>
Review Comment:
For the following, do you think we can add a guard explicitly?
> SPARK-30628: V2 file scans don't implement SupportsRuntimeV2Filtering
```scala
- case fs: FileScan if runtimeFilters.nonEmpty =>
+ case fs: FileScan if runtimeFilters.nonEmpty && !fs.isInstanceOf[SupportsRuntimeV2Filtering] =>
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-6570101604
@pan3793 Done in bfc057d, after merging master. `SupportsRuntimeCatalystFiltering`, its three methods and the `plannedPartitions` memo now live in `RuntimePartitionPruningFileScan`, which `ParquetScan`, `OrcScan`, `AvroScan` and `TextBasedFileScan` extend, and `FileScan` is the same as on master. `DataSourceV2FileScanSubclassSuite` covers the shapes from the probe: a Kyuubi-like `FileScan with SupportsRuntimeFiltering`, the same with a SPYT-like cache in `partitions`, and a Sedona-like scan with its own partition type. @dongjoon-hyun this covers items 1, 2 and 4; the description update comes next.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-6624297987
@pan3793 a follow-up on item 5. A review of #60053 pointed out that a DSv2 source can accept `part = 3` fully, apply it row by row, and leave the pruning of its planned partitions to DPP. For such a source the general rule would turn a pruned scan into a full one, so I closed #60053. In 33a96ff, `CleanupDynamicPruningFilters` reads `pushedFilters` only for the built-in file scans, and every other DSv2 source keeps the DPP filter, as on master. A new test checks this for a Kyuubi-like `FileScan` that implements `SupportsRuntimeFiltering`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4505534960
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
Review Comment:
Done in c1ba7eacfd7. The user-facing section says that such a subclass fails every query, next to the finals.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
+ fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private val allPartitions = new TransientBestEffortLazyVal(() => filePartitions(listing()))
+
+ // Final so that a subclass of a built-in scan cannot change one read path without the other.
+ final override protected def partitions: Seq[FilePartition] = allPartitions()
+
+ final override def planInputPartitions(): Array[InputPartition] = allPartitions().toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
Review Comment:
Done in c1ba7eacfd7. `planInputPartitionsWithRuntimeFilters` is final too, and the user-facing section says that a subclass overriding any of the three no longer compiles, and one built against 4.3.0 fails to load.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala:
##########
@@ -233,6 +233,10 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join
private def hasSelectivePredicate(plan: LogicalPlan): Boolean = plan.exists {
case f: Filter => isLikelySelective(f.condition)
+ // A built-in V2 file scan takes its partition filters in without leaving a `Filter` above it,
Review Comment:
It is in the description now, next to the bloom-filter bullet, and I opened SPARK-61173 for it. This PR does not change that path, and the gap holds on master for any DSv2 dimension whose filter the scan takes in fully, so the fix changes plans for every DSv2 source. It also needs its own care: the creation side must not read a scan with a pushed limit or top-N, and `FileScan.estimateStatistics` reports the unpruned table size, so the size check would still differ from V1.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "dongjoon-hyun (via GitHub) via reviews" <re...@spark.apache.org>.
dongjoon-hyun commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4503842447
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -230,48 +238,92 @@ object PushDownUtils extends Logging {
translatedFiltersPushed || partPredicatesPushed
- case catalystScan: SupportsRuntimeCatalystFiltering if pushableFilters.nonEmpty =>
- // A runtime filter is normally evaluated twice: the source prunes with it, and the
- // FilterExec above the scan applies it again. The two have to agree, so this screen
- // pushes only predicates the source can be trusted to evaluate in Spark's place.
- //
- // But pushing a non-deterministic one would let the source prune on its own coin flip and
- // Spark flip again for the rows that survive, so we handle that specially.
- //
- // Not pushing a filter is safe only while its FilterExec still evaluates it. A fully
- // pushed filter has none, so the source is its only evaluator. That is why
- // DataSourceV2Strategy deletes a FilterExec at planning time only for a filter we push
- // below, and we use the same method (isPushablePartitionFilter) to determine if we
- // should push it.
- //
- // A DPP filter degrades to TrueLiteral once its subquery is pruned away, so it matches
- // every row. The V2 path above drops these implicitly, since translateRuntimeFilterV2
- // returns None; here we push Catalyst expressions directly, so we remove them explicitly.
- val catalystFilters = pushableFilters
- .flatMap(unwrapRuntimeFilterExpression)
- .filterNot(_ == Literal.TrueLiteral)
- if (catalystFilters.nonEmpty) {
- catalystScan.filter(catalystFilters.toArray)
- true
- } else {
- false
+ case catalystScan: SupportsRuntimeCatalystFiltering
+ if runtimeFilters.nonEmpty && implementsOnlyDeprecatedFilter(catalystScan) =>
+ filterWithDeprecatedMethod(catalystScan, runtimeFilters, output)
+
+ case catalystScan: SupportsRuntimeCatalystFiltering =>
+ // Resolved only when a filter could be fully pushed, so a node with no such filter does
+ // not ask the scan.
+ lazy val fullyPushedAttrs = V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs(
+ catalystScan.fullyPushedFilterAttributes(),
+ output,
+ "fullyPushedFilterAttributes()",
+ catalystScan.getClass.getName)
+ val evaluatedOnlyByScan =
+ runtimeFilters.filter(isFullyPushedRuntimeFilter(_, fullyPushedAttrs))
+ if (evaluatedOnlyByScan.nonEmpty) {
+ throw SparkException.internalError(
+ s"${scan.getClass.getName} fully evaluates the runtime filters " +
+ s"${evaluatedOnlyByScan.mkString(", ")}, so they must go through " +
+ "PushDownUtils.replanWithRuntimeFilters, not pushRuntimeFilters.")
}
+ false
case _ =>
false
}
}
+ /**
+ * Screens runtime filters for a [[SupportsRuntimeCatalystFiltering]] scan and unwraps them to the
+ * Catalyst predicates the scan plans with, untranslated.
+ *
+ * A filter is kept only if it references attributes the scan declared filterable and passes
+ * `isPushablePartitionFilter`, the check `DataSourceV2Strategy` uses to drop a pushed filter's
+ * `FilterExec`. A DPP filter degraded to true is removed, as the V2 path does by not translating
+ * it.
+ */
+ private def catalystRuntimeFilters(
+ scan: SupportsRuntimeCatalystFiltering,
+ runtimeFilters: Seq[Expression],
+ output: Seq[AttributeReference]): Seq[Expression] = {
+ val filterAttrs = V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs(
+ scan.filterAttributes(),
+ output,
+ "filterAttributes()",
+ scan.getClass.getName)
+ runtimeFilters
+ .filter(f => f.references.subsetOf(filterAttrs))
+ .filter(isPushablePartitionFilter(_, includeSubquery = true))
+ .flatMap(unwrapRuntimeFilterExpression)
+ .filterNot(_ == Literal.TrueLiteral)
+ }
+
+ /**
+ * Whether `scan` gets `planInputPartitionsWithRuntimeFilters` from the default in
+ * [[SupportsRuntimeCatalystFiltering]], such as a class compiled against Spark 4.3 or a Java
+ * class implementing the trait directly that overrides only `filter`. A Scala class that mixes in
+ * the trait gets a forwarder, so it and its subclasses take the general path.
+ */
+ private def implementsOnlyDeprecatedFilter(scan: SupportsRuntimeCatalystFiltering): Boolean =
+ scan.getClass.getMethod("planInputPartitionsWithRuntimeFilters", classOf[Array[Expression]])
Review Comment:
**26. [Low] `implementsOnlyDeprecatedFilter` can fail a query with `NoClassDefFoundError`**
`getMethod` builds the public methods of each class it searches, starting with the scan's own class, and that resolves every type in their signatures. So a `SupportsRuntimeCatalystFiltering` scan with a public method over a type that is missing at runtime, e.g. one from an optional dependency, now fails every query that has a runtime filter on it with `NoClassDefFoundError` at L242 or L356, while 4.3.0 called `filter` without reflection and ran the query. The lookup also runs again on every execution of every such scan node. `CodeGenerator.leavesGenCodeToExpression` (CodeGenerator.scala:2929-2936) makes the same inherited-default check behind a `ClassValue` cache and `catch { case _: LinkageError => false }`. Could this do the same? With `false` such a scan takes the general path, whose default still calls `filter`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -226,22 +227,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
}
// Screen with the same test pushdown applies, or a filter dropped here and rejected there
// would be evaluated nowhere.
- val fullyPushedRuntimeFilters = scalarSubqueryFilters.filter { f =>
- f.references.subsetOf(relation.fullyPushedRuntimeFilterAttrs) &&
- PushDownUtils.isPushablePartitionFilter(f, includeSubquery = true)
- }
+ val fullyPushedRuntimeFilters = scalarSubqueryFilters.filter(
+ PushDownUtils.isFullyPushedRuntimeFilter(_, relation.fullyPushedRuntimeFilterAttrs))
// dynamicFilters need no such check: a DynamicPruningSubquery over a non-deterministic
// filtering plan is itself non-deterministic, so CleanupDynamicPruningFilters has already
// rewritten it to TrueLiteral by the time we get here.
+ // A Catalyst runtime-filtering scan only applies a filter over attributes it declared, so a
+ // DPP filter over any other column is dropped here rather than planned for nothing. One
+ // PartitionPruning inserted over a union lands in every branch, including one that cannot
+ // use it; the join it was derived from still applies it.
+ val usableDynamicFilters = relation.scan match {
+ case _: SupportsRuntimeCatalystFiltering =>
+ dynamicFilters.filter(_.references.subsetOf(relation.runtimeFilterAttrs))
Review Comment:
**31. [Nit] An unpartitioned branch of a union gets no runtime filter when the DPP filter is kept**
When `PartitionPruning` inserts a DPP filter over a union of a partitioned V2 file table and an unpartitioned one, `InjectRuntimeFilter` finds it on the key and adds no bloom filter (InjectRuntimeFilter.scala:423), the filter then lands in both branches (L236-238), and this line drops the copy in the unpartitioned branch. When that DPP filter is kept, e.g. as its own subquery with `reuseBroadcastOnly=false`, the unpartitioned branch gets no runtime filter, while on master, where a V2 file scan gets no DPP filter, the bloom filter goes into both branches where its checks pass. SPARK-60702 does not change this, since such a DPP filter counts as kept. The results are the same. Could the description list it with the bloom-filter bullet?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala:
##########
@@ -233,6 +233,10 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join
private def hasSelectivePredicate(plan: LogicalPlan): Boolean = plan.exists {
case f: Filter => isLikelySelective(f.condition)
+ // A built-in V2 file scan takes its partition filters in without leaving a `Filter` above it,
Review Comment:
**29. [Low] A V2 file dim filtered only on its partition column cannot be a bloom filter's creation side**
Item 12 fixed this for DPP, but `InjectRuntimeFilter.extractSelectiveFilterOverScan` makes the same check for a bloom filter's creation side: it counts a selective predicate only at a `Filter` (InjectRuntimeFilter.scala:140) and accepts a leaf only after one (L237). A V2 file dim filtered only on its partition column keeps no `Filter` above its scan, so it cannot be the creation side, while V1 can, since `PruneFileSourcePartitions` keeps the `Filter`. This already holds on master for any DSv2 dim, but it matters more after merging master. With SPARK-60702, a key whose `onlyInBroadcast` DPP filter physical planning drops now gets a bloom filter when this path finds a creation side. So with the default `reuseBroadcastOnly` on a shuffle join, the query of "DPP for a dimension filtered only on its partition column, as in V1" would get a bloom filter in V1, where its size checks pass, and no runtime filter at all in V2. `InjectRuntimeFilter` is in sql/catalyst and cannot see this trait
, but it could read `DataSourceV2ScanRelation.pushedFilters`, as the check below does. InjectRuntimeFilter.scala is outside this diff, so I left this here. Could it be a follow-up JIRA, or listed in the description next to the bloom-filter bullet?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
+ fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private val allPartitions = new TransientBestEffortLazyVal(() => filePartitions(listing()))
+
+ // Final so that a subclass of a built-in scan cannot change one read path without the other.
+ final override protected def partitions: Seq[FilePartition] = allPartitions()
+
+ final override def planInputPartitions(): Array[InputPartition] = allPartitions().toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
Review Comment:
**28. [Nit] `planInputPartitionsWithRuntimeFilters` is still open to a subclass**
L50 makes `partitions` and `planInputPartitions()` final "so that a subclass of a built-in scan cannot change one read path without the other", and the description says such a subclass "can change the packing only by overriding `filePartitions`". This method is not final, though, so a subclass can still change the runtime-filtered path alone. Could it be `final override def` as well? Nothing in-tree overrides it on a built-in scan.
The two finals are also a binary change for such a subclass. scalac emits them as `final` forwarders in each built-in scan class, so a subclass built against 4.3.0 that overrides `planInputPartitions()` now fails at class loading ("overrides final method"), and a recompiled one fails to compile. MiMa does not report it, since `org.apache.spark.sql.execution.*` is excluded (MimaExcludes.scala:191). Could the user-facing section mention both finals?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
+ fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private val allPartitions = new TransientBestEffortLazyVal(() => filePartitions(listing()))
+
+ // Final so that a subclass of a built-in scan cannot change one read path without the other.
+ final override protected def partitions: Seq[FilePartition] = allPartitions()
+
+ final override def planInputPartitions(): Array[InputPartition] = allPartitions().toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
Review Comment:
**24. [Low] A group-based row-level operation on a built-in scan now gets a group filter that the scan does not record**
With the partition columns declared here, `RowLevelOperationRuntimeGroupFiltering` injects a group filter into a group-based (`ReplaceData`) operation whose scan is a built-in file scan, since it checks only `runtimeFilterAttrs.nonEmpty` (RowLevelOperationRuntimeGroupFiltering.scala:67-72). The read then gets only the partition directories that have a matching row. The trait asks a group-based scan whose write replaces the groups that were read to record what each call returned (SupportsRuntimeCatalystFiltering.scala:96-98), but this one records nothing, and `planInputPartitions()` still returns every directory. So a connector whose `SupportsRowLevelOperations` builds its scan from a built-in one, and whose write replaces the groups that scan plans, e.g. from `planInputPartitions()` or the file index, would now rewrite only the matching partitions of `DELETE FROM t WHERE x = 5` but replace all of them, which loses the rows of the others. On master the scan has no runtime filter at
tributes, so the rule does nothing and the read covers every group. No in-tree `FileTable` supports row-level operations, and I did not find such a connector on GitHub, so this is a contract gap rather than a known break.
Could `canInjectGroupFilters` skip a `RuntimePartitionPruningFileScan`, as on master, or could the scaladoc at L67-73 say that the scan does not record what each call returned, so it must not back such a write?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -293,11 +347,21 @@ object PushDownUtils extends Logging {
output: Seq[AttributeReference],
keyedPartitioning: Option[KeyedPartitioning],
originalPartitions: => Seq[InputPartition]): Seq[Option[InputPartition]] = {
- val filtered = pushRuntimeFilters(scan, runtimeFilters, table, output)
+ val (filtered, newPartitions) = scan match {
+ // A Catalyst scan is asked even when screening leaves no filter, unless it also implements
+ // SupportsRuntimeV2Filtering, which pushRuntimeFilters rejects, or declares only the
+ // deprecated `filter`, which pushRuntimeFilters calls as Spark 4.3 did.
+ case catalystScan: SupportsRuntimeCatalystFiltering
+ if runtimeFilters.nonEmpty && !catalystScan.isInstanceOf[SupportsRuntimeV2Filtering] &&
+ !implementsOnlyDeprecatedFilter(catalystScan) =>
Review Comment:
**25. [Low] Follow-up to my item 10: a Scala scan recompiled against this trait still plans again when screening removes every runtime filter**
This guard sends a scan to `pushRuntimeFilters` only when its classes do not declare `planInputPartitionsWithRuntimeFilters`, i.e. one built against 4.3.0 or a Java one (L293-298). A Scala scan recompiled against this trait gets a mixin forwarder for it, so one that still overrides only `filter`, which the description says needs only `override` on it, takes this branch, as PushDownUtils.scala:182-185 says. When screening removes every runtime filter, the default then skips `filter` but plans again through `toBatch.planInputPartitions()` (SupportsRuntimeCatalystFiltering.scala:120-121) instead of reading `originalPartitions`. That is item 10 again for every such scan recompiled from now on: one more planning per scan node on a sort-merge join under the default DPP settings, and, for a scan instance that another node already narrowed through `filter`, a read of the narrowed set, the SPARK-60752 class from item 14.
The test at DataSourceV2CatalystRuntimeFilterSuite.scala:568-579 cannot see this. `LegacyFilterScan` is such a recompiled Scala scan, although its doc says "A scan written against Spark 4.3" (L965), and the test passes `scan.toBatch.planInputPartitions()` as `originalPartitions`, so a new plan and the planned partitions agree, and L576 reads as if the scan took the 4.3.0 route.
Reflection sees a forwarder as a declaration, so I do not see a check here that tells the two apart. Could the scaladoc of `filter` (SupportsRuntimeCatalystFiltering.scala:124-128) and the description say that such a scan should implement `planInputPartitionsWithRuntimeFilters`, since it otherwise plans again when screening removes every runtime filter, and could L576 and L965 say that this Scala fixture takes the general path?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
Review Comment:
**27. [Low] A subclass of a built-in scan that adds `SupportsRuntimeV2Filtering` now fails every query**
The six built-in scans are case classes that are not final, and through this trait they now implement `SupportsRuntimeCatalystFiltering`. A subclass that adds `SupportsRuntimeFiltering` or `SupportsRuntimeV2Filtering`, e.g. `class MyScan(...) extends ParquetScan(...) with SupportsRuntimeV2Filtering`, then implements both, and even a plain `SELECT` fails with INTERNAL_ERROR, since DataSourceV2Strategy.scala:218 reads `runtimeFilterAttrs`, which calls `checkRuntimeFilteringInterfaces()` (DataSourceV2Relation.scala:225-229 and 392-400). On master such a subclass works, and its runtime filters reach its own `filter`. The per-format opt-in keeps a `FileScan` subclass as on master, but not a subclass of a built-in scan. As with item 9, I did not find one on GitHub, so this is a contract gap rather than a known break. Could the user-facing section of the description say so, next to the finals of item 28?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
Review Comment:
**30. [Low] Every built-in file scan keeps its file listing for the plan's lifetime**
`listing` holds the `Seq[PartitionDirectory]` from `fileIndex.listFiles` for as long as the scan lives, on every built-in file scan, including an unpartitioned table, a scan that reads no partition column and a plan without runtime filters, none of which reads it again after `allPartitions`. `listFiles` wraps each file in a new `FileStatusWithMetadata` per call (FileIndex.scala:103-104), so this keeps an object per file on top of `allPartitions`, tens of MB for a table with a million files, for as long as a `QueryExecution` or a cached plan holds the scan. After `InMemoryFileIndex.refresh()`, a retained scan also keeps the previous listing's `FileStatus`es alive. The merge-base drops the listing after packing. Could `listing` be memoized only when `filterAttributes()` is not empty, and otherwise stay local to `allPartitions`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
+ fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private val allPartitions = new TransientBestEffortLazyVal(() => filePartitions(listing()))
+
+ // Final so that a subclass of a built-in scan cannot change one read path without the other.
+ final override protected def partitions: Seq[FilePartition] = allPartitions()
+
+ final override def planInputPartitions(): Array[InputPartition] = allPartitions().toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ if (expressions.isEmpty) return allPartitions().toArray
+ // Directory values follow the table's partition schema, not the read one, and bind as V1
+ // binds them: with that schema's types, nullable.
+ val partitionSchema = fileIndex.partitionSchema
+ val fieldIndex = partitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ val predicates = expressions.toSeq.map { e =>
+ Predicate.createInterpreted(e.transform {
+ case a: Attribute =>
+ // PushDownUtils passes only filters over `filterAttributes()`.
+ val i = fieldIndex.getOrElse(normalizePartitionName(a.name),
+ throw SparkException.internalError("A file scan can only apply a runtime filter " +
+ s"over its partition columns ${partitionSchema.fieldNames.mkString("[", ", ", "]")}" +
+ s", got $e"))
+ BoundReference(i, partitionSchema(i).dataType, nullable = true)
+ })
+ }
+ val all = listing()
+ val kept = all.filter { dir =>
+ predicates.forall { predicate =>
+ try predicate.eval(dir.values) catch { case NonFatal(_) => true }
+ }
+ }
+ if (kept.length == all.length) allPartitions().toArray else filePartitions(kept).toArray
Review Comment:
**32. [Nit] A scan that DPP prunes still packs every file at planning**
When a runtime filter keeps few directories, the node packs twice: planning packs every listed file through `allPartitions`, and this line packs the survivors. Planning has to, because `FileScan` keeps the default `columnarSupportMode()`, `PARTITION_DEFINED`, so `supportsColumnar` (DataSourceV2ScanExecBase.scala:200-210), which DataSourceV2Strategy.scala:275 asks, reads `inputPartitions`. With DPP keeping 10 of 10,000 partitions of a table with a million files, the driver splits, sorts and packs a million files at planning only to learn one boolean. V1 packs only `dynamicallySelectedPartitions`. The merge-base pays the same at planning, so this is not a regression. Could a follow-up have the built-in scans report `columnarSupportMode()` themselves? Their reader factories decide columnar reads from the schema and the conf, not from the partition.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -230,48 +238,92 @@ object PushDownUtils extends Logging {
translatedFiltersPushed || partPredicatesPushed
- case catalystScan: SupportsRuntimeCatalystFiltering if pushableFilters.nonEmpty =>
- // A runtime filter is normally evaluated twice: the source prunes with it, and the
- // FilterExec above the scan applies it again. The two have to agree, so this screen
- // pushes only predicates the source can be trusted to evaluate in Spark's place.
- //
- // But pushing a non-deterministic one would let the source prune on its own coin flip and
- // Spark flip again for the rows that survive, so we handle that specially.
- //
- // Not pushing a filter is safe only while its FilterExec still evaluates it. A fully
- // pushed filter has none, so the source is its only evaluator. That is why
- // DataSourceV2Strategy deletes a FilterExec at planning time only for a filter we push
- // below, and we use the same method (isPushablePartitionFilter) to determine if we
- // should push it.
- //
- // A DPP filter degrades to TrueLiteral once its subquery is pruned away, so it matches
- // every row. The V2 path above drops these implicitly, since translateRuntimeFilterV2
- // returns None; here we push Catalyst expressions directly, so we remove them explicitly.
- val catalystFilters = pushableFilters
- .flatMap(unwrapRuntimeFilterExpression)
- .filterNot(_ == Literal.TrueLiteral)
- if (catalystFilters.nonEmpty) {
- catalystScan.filter(catalystFilters.toArray)
- true
- } else {
- false
+ case catalystScan: SupportsRuntimeCatalystFiltering
+ if runtimeFilters.nonEmpty && implementsOnlyDeprecatedFilter(catalystScan) =>
+ filterWithDeprecatedMethod(catalystScan, runtimeFilters, output)
+
+ case catalystScan: SupportsRuntimeCatalystFiltering =>
+ // Resolved only when a filter could be fully pushed, so a node with no such filter does
+ // not ask the scan.
+ lazy val fullyPushedAttrs = V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs(
+ catalystScan.fullyPushedFilterAttributes(),
+ output,
+ "fullyPushedFilterAttributes()",
+ catalystScan.getClass.getName)
+ val evaluatedOnlyByScan =
+ runtimeFilters.filter(isFullyPushedRuntimeFilter(_, fullyPushedAttrs))
+ if (evaluatedOnlyByScan.nonEmpty) {
+ throw SparkException.internalError(
+ s"${scan.getClass.getName} fully evaluates the runtime filters " +
+ s"${evaluatedOnlyByScan.mkString(", ")}, so they must go through " +
+ "PushDownUtils.replanWithRuntimeFilters, not pushRuntimeFilters.")
}
+ false
case _ =>
false
}
}
+ /**
+ * Screens runtime filters for a [[SupportsRuntimeCatalystFiltering]] scan and unwraps them to the
+ * Catalyst predicates the scan plans with, untranslated.
+ *
+ * A filter is kept only if it references attributes the scan declared filterable and passes
+ * `isPushablePartitionFilter`, the check `DataSourceV2Strategy` uses to drop a pushed filter's
+ * `FilterExec`. A DPP filter degraded to true is removed, as the V2 path does by not translating
+ * it.
+ */
+ private def catalystRuntimeFilters(
+ scan: SupportsRuntimeCatalystFiltering,
+ runtimeFilters: Seq[Expression],
+ output: Seq[AttributeReference]): Seq[Expression] = {
+ val filterAttrs = V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs(
+ scan.filterAttributes(),
+ output,
+ "filterAttributes()",
+ scan.getClass.getName)
+ runtimeFilters
+ .filter(f => f.references.subsetOf(filterAttrs))
Review Comment:
**36. [Nit] The `filterAttributes()` screen at execution no longer drops anything Spark planned**
Since 2acbef9 the strategy keeps only the DPP filters over `relation.runtimeFilterAttrs` for a Catalyst scan (DataSourceV2Strategy.scala:239-243) and routes only scalar subquery filters over them (L218-224). `BatchScanExec` is built only at DataSourceV2Strategy.scala:269, and the copies that DPP planning makes keep the same columns, so this `subsetOf` cannot drop a filter of a node Spark planned. It only guards a node built elsewhere, at the cost of resolving `filterAttributes()` again on each execution. The `true` removal at L290 is still needed. Could the scaladoc at L268-276 say that the strategy already drops these? Alternatively, `CleanupDynamicPruningFilters` could drop a DPP filter whose `pruningKey.references` are not in `r.runtimeFilterAttrs` for every `DataSourceV2ScanRelation`, which would also keep it from the `filter` of a `SupportsRuntimeV2Filtering` branch, which gets it on master.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/CleanupDynamicPruningFilters.scala:
##########
@@ -79,9 +90,14 @@ object CleanupDynamicPruningFilters extends Rule[LogicalPlan] with PredicateHelp
case p @ NodeWithOnlyDeterministicProjectAndFilter(
HiveTableRelation(_, _, _, _, _)) =>
removeUnnecessaryDynamicPruningSubquery(p)
- case p @ NodeWithOnlyDeterministicProjectAndFilter(
- _: DataSourceV2ScanRelation) =>
- removeUnnecessaryDynamicPruningSubquery(p)
+ // Only a built-in file scan is known to prune what it plans with a partition filter it fully
+ // took. Another DSv2 source may apply such a filter row by row and leave the pruning to DPP.
Review Comment:
**35. [Nit] The reason at L93-94 does not hold for the `Filter` path of the same rule**
L93-94 keep the pushed path to built-in file scans because another DSv2 source "may apply such a filter row by row and leave the pruning to DPP". But the `Filter` path of this rule (L62-65), which master already applies to every `DataSourceV2ScanRelation`, drops the DPP filter for an equality left in a `Filter` above the scan, which the source did not take and so cannot prune by. The rule thus trusts a DSv2 source less when it took the equality than when it did not. Keeping other sources as on master after #60053 is a fine reason by itself; could the comment say that instead?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala:
##########
@@ -233,6 +233,10 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join
private def hasSelectivePredicate(plan: LogicalPlan): Boolean = plan.exists {
case f: Filter => isLikelySelective(f.condition)
+ // A built-in V2 file scan takes its partition filters in without leaving a `Filter` above it,
+ // where V1 keeps one.
+ case r: DataSourceV2ScanRelation if r.scan.isInstanceOf[RuntimePartitionPruningFileScan] =>
Review Comment:
**34. [Nit] Follow-up to my item 12: the filtering side need not be a built-in file scan**
The `isInstanceOf[RuntimePartitionPruningFileScan]` here came from my suggestion in item 12, but the filtering side does not need it. This only asks whether a likely selective predicate restricts the dim, and a predicate that any DSv2 source took in fully, leaving no `Filter`, does so just as one in a `Filter` does. The scope in `CleanupDynamicPruningFilters` is about what the pruned scan guarantees, which does not apply here. So a dim on a connector scan with a fully pushed `region = 'EU'` still gets no DPP, as on master. Could this be `case r: DataSourceV2ScanRelation => r.pushedFilters.exists(isLikelySelective)`? It is a benefit heuristic, so the results stay the same, though with `reuseBroadcastOnly=false` it can add DPP subqueries for such dims.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -226,22 +227,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
}
// Screen with the same test pushdown applies, or a filter dropped here and rejected there
// would be evaluated nowhere.
- val fullyPushedRuntimeFilters = scalarSubqueryFilters.filter { f =>
- f.references.subsetOf(relation.fullyPushedRuntimeFilterAttrs) &&
- PushDownUtils.isPushablePartitionFilter(f, includeSubquery = true)
- }
+ val fullyPushedRuntimeFilters = scalarSubqueryFilters.filter(
+ PushDownUtils.isFullyPushedRuntimeFilter(_, relation.fullyPushedRuntimeFilterAttrs))
// dynamicFilters need no such check: a DynamicPruningSubquery over a non-deterministic
// filtering plan is itself non-deterministic, so CleanupDynamicPruningFilters has already
// rewritten it to TrueLiteral by the time we get here.
+ // A Catalyst runtime-filtering scan only applies a filter over attributes it declared, so a
+ // DPP filter over any other column is dropped here rather than planned for nothing. One
+ // PartitionPruning inserted over a union lands in every branch, including one that cannot
+ // use it; the join it was derived from still applies it.
+ val usableDynamicFilters = relation.scan match {
+ case _: SupportsRuntimeCatalystFiltering =>
+ dynamicFilters.filter(_.references.subsetOf(relation.runtimeFilterAttrs))
+ case _ => dynamicFilters
+ }
// A group-based read keeps only the group filter, for the reason above. PartitionPruning
// plans a DynamicPruningSubquery from the join of one read, so it is a filter of one read
// too. With an IN subquery it lands on the read of the updated rows of an UPDATE. With a
// NOT EXISTS subquery the negated condition becomes a semi join, so it lands on the read
// of the rows to copy over.
val runtimeFilters = if (groupBasedRead) {
- dynamicFilters.filterNot(_.isInstanceOf[DynamicPruningSubquery])
+ usableDynamicFilters.filterNot(_.isInstanceOf[DynamicPruningSubquery])
} else {
- dynamicFilters ++ scalarSubqueryFilters
+ usableDynamicFilters ++ scalarSubqueryFilters
Review Comment:
**33. [Nit] A scalar subquery filter over a partition column is planned twice**
Since e721ae7 the file scans declare no partition column fully pushed (item 2), so `part = (SELECT ...)` goes into `runtimeFilters` here and also stays in the `Filter` above (L265), as for a `SupportsRuntimeV2Filtering` scan on master. Under the defaults, subquery reuse runs it once, but with `spark.sql.execution.reuseSubquery=false` it runs twice, and if its source changes in between, the scan prunes with one value while the `Filter` keeps another. V1 and the merge-base evaluate it once. That is a fair price for item 2, but could the user-facing section say that such a filter now shows up both in `RuntimeFilters` and in the `Filter` above the scan?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
+ fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private val allPartitions = new TransientBestEffortLazyVal(() => filePartitions(listing()))
+
+ // Final so that a subclass of a built-in scan cannot change one read path without the other.
+ final override protected def partitions: Seq[FilePartition] = allPartitions()
+
+ final override def planInputPartitions(): Array[InputPartition] = allPartitions().toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ if (expressions.isEmpty) return allPartitions().toArray
+ // Directory values follow the table's partition schema, not the read one, and bind as V1
+ // binds them: with that schema's types, nullable.
+ val partitionSchema = fileIndex.partitionSchema
+ val fieldIndex = partitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ val predicates = expressions.toSeq.map { e =>
+ Predicate.createInterpreted(e.transform {
+ case a: Attribute =>
+ // PushDownUtils passes only filters over `filterAttributes()`.
+ val i = fieldIndex.getOrElse(normalizePartitionName(a.name),
+ throw SparkException.internalError("A file scan can only apply a runtime filter " +
+ s"over its partition columns ${partitionSchema.fieldNames.mkString("[", ", ", "]")}" +
+ s", got $e"))
+ BoundReference(i, partitionSchema(i).dataType, nullable = true)
+ })
+ }
+ val all = listing()
+ val kept = all.filter { dir =>
+ predicates.forall { predicate =>
+ try predicate.eval(dir.values) catch { case NonFatal(_) => true }
Review Comment:
**38. [Nit] A runtime filter that fails on a directory leaves no log**
Keeping the directory is right, but nothing records that the filter failed, so a user who sees no pruning has nothing to go on. `PartitionPredicateImpl.eval` keeps a partition for the same reason and logs a warning with the expression and the exception (PartitionPredicateImpl.scala:86-95). Could this log one too, e.g. once per call rather than per directory? This is also another binder of directory values next to `PartitioningAwareFileIndex.prunePartitions`, `FileSourceScanExec.dynamicallySelectedPartitions` and `PartitionPredicateImpl`. None of them fits as is, so a shared one could wait for a follow-up.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
+ fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private val allPartitions = new TransientBestEffortLazyVal(() => filePartitions(listing()))
+
+ // Final so that a subclass of a built-in scan cannot change one read path without the other.
+ final override protected def partitions: Seq[FilePartition] = allPartitions()
+
+ final override def planInputPartitions(): Array[InputPartition] = allPartitions().toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ if (expressions.isEmpty) return allPartitions().toArray
+ // Directory values follow the table's partition schema, not the read one, and bind as V1
+ // binds them: with that schema's types, nullable.
+ val partitionSchema = fileIndex.partitionSchema
+ val fieldIndex = partitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ val predicates = expressions.toSeq.map { e =>
+ Predicate.createInterpreted(e.transform {
+ case a: Attribute =>
+ // PushDownUtils passes only filters over `filterAttributes()`.
+ val i = fieldIndex.getOrElse(normalizePartitionName(a.name),
+ throw SparkException.internalError("A file scan can only apply a runtime filter " +
+ s"over its partition columns ${partitionSchema.fieldNames.mkString("[", ", ", "]")}" +
+ s", got $e"))
+ BoundReference(i, partitionSchema(i).dataType, nullable = true)
+ })
+ }
+ val all = listing()
+ val kept = all.filter { dir =>
+ predicates.forall { predicate =>
+ try predicate.eval(dir.values) catch { case NonFatal(_) => true }
+ }
+ }
+ if (kept.length == all.length) allPartitions().toArray else filePartitions(kept).toArray
+ }
+
+ // Read when the scan is built, as `FileScan` does.
+ private val isCaseSensitive = conf.caseSensitiveAnalysis
+
+ private def normalizePartitionName(name: String): String = {
Review Comment:
**37. [Nit] Following up on item 6: a helper under a new name avoids the clash with Kyuubi**
Your reply in item 6 is right that a `protected` or package-private `normalizeName` would clash with Kyuubi's `private def normalizeName`, but that comes from reusing the name. A helper under a new name has no such clash, e.g. in `FileScan`:
```scala
private[v2] def normalizePartitionColumnName(name: String): String = normalizeName(name)
```
A subclass outside the package neither sees nor overrides it, and as a `def` it needs no field in a subclass built against 4.3.0. `filePartitions` and this trait would then share one rule again, instead of two copies of the build-time `conf.caseSensitiveAnalysis` (L102-107 and FileScan.scala:217-225).
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4454062991
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala:
##########
@@ -72,35 +80,38 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty
/**
- * Filters this scan using runtime Catalyst expressions.
+ * Plans this scan's input partitions with runtime Catalyst expressions applied.
+ *
+ * The provided expressions must be interpreted as a set of predicates that are ANDed together,
+ * and applied on top of the filters already pushed at query compilation. Implementations may use
+ * them to plan fewer [[InputPartition]]s than [[org.apache.spark.sql.connector.read.Batch]]
+ * would without them.
*
- * The provided expressions must be interpreted as a set of predicates that are ANDed together.
- * Implementations may use the expressions to prune initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s.
+ * Spark calls this method at execution time, and only when at least one runtime filter survives
+ * screening; the partitions it returns replace the ones that scan node would otherwise have
+ * read. Spark may also call `Batch.planInputPartitions()` on the same scan, so neither method may
+ * assume it is the only one used. One scan instance can back several scan nodes (e.g. the two
+ * branches of a group-based UPDATE), each with its own runtime filters, so an implementation
+ * must derive its result from the given expressions alone and must not carry state over from a
+ * previous call. A group-based scan whose write replaces the groups that were read has to record
+ * what each call returned for that write, and count a scan node that read through
+ * `planInputPartitions()` as having read every group that method returns. Spark may also call
+ * that method while planning a scan node, so such a call alone does not show that the node read
+ * through it.
*
* Spark currently tracks runtime-filter eligibility by root attribute. If [[filterAttributes]]
* returns a nested reference, an expression may access another nested field under the same root.
* The scan must match each access against its own partition layout and use only expressions it
* can apply.
*
- * Spark may call this method more than once for the same scan instance: a plan can hold several
- * scan nodes sharing one scan (e.g. the two branches of a group-based UPDATE), and each pushes
- * its own copy of the runtime filters. Implementations must treat successive calls as additive,
- * ANDing the new expressions with those already pushed rather than replacing them.
- *
* If the scan also implements
* [[org.apache.spark.sql.connector.read.SupportsReportPartitioning]], it must preserve
- * the originally reported partitioning during runtime filtering. While applying runtime
- * predicates, the scan may detect that some
- * [[org.apache.spark.sql.connector.read.InputPartition]]s have no matching data, in which
- * case it can either replace the initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s that have no matching data with
- * empty [[org.apache.spark.sql.connector.read.InputPartition]]s, or report only a subset of
- * the original partition values (omitting those with no data) via
- * [[org.apache.spark.sql.connector.read.Batch#planInputPartitions]]. The scan must not
- * report new partition values that were not present in the original partitioning.
- *
- * Note that Spark will call [[Scan.toBatch]] again after filtering the scan at runtime.
+ * the originally reported partitioning. While applying runtime predicates, the scan may detect
+ * that some [[InputPartition]]s have no matching data, in which case it can either replace
+ * them with empty [[InputPartition]]s, or report only a subset of the original partition
+ * values (omitting those with no data). The scan must not report new partition values that
+ * were not present in the original partitioning.
*/
- def filter(expressions: Array[Expression]): Unit
+ def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition]
Review Comment:
Done in e721ae7. "DPP and a scalar subquery prune input partitions at runtime for v2 text" uses its own single-column text fixture, and `AvroV2Suite` has "SPARK-30628: Avro source v2: runtime filters prune partitions". Both check that only partition 7 is read.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,787 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+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
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, FileScan, RuntimePartitionPruningFileScan}
+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.{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 {
+
+ // 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")
Review Comment:
Done in e721ae7. `DataSourceV2FileSourceDPPSuite` clears the temp views after each test, and `withFactAndDim` in `DataSourceV2FileScanSubclassSuite` uses `withTempView`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-6602662925
@dongjoon-hyun thanks for the review. Addressed in e721ae7, except item 4, which went away in 7158934, and item 6, which I answered inline. Item 1 keeps a scan built against 4.3.0 working, and for item 2 the file scans no longer declare their partition columns fully pushed.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4454064947
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType, a.nullable)
Review Comment:
I'd rather keep the copy. Kyuubi's `HiveScan` extends `FileScan` and defines its own `private def normalizeName`. If the one in `FileScan` becomes `protected` or package-private, scalac 2.13 rejects that private method with "weaker access privileges in overriding", so Kyuubi would stop compiling.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
+ .flatMap(_.files.filter(file => predicate.eval(file.partitionValues)))
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+ // Pack the survivors instead of handing back the planned partitions with files removed:
+ // `partitions` sized those bins for the whole file set, so a scan that pruned most of its files
+ // would run the same few tasks over a fraction of the data. V1 packs the pruned set the same
+ // way, from `dynamicallySelectedPartitions`.
+ val openCostInBytes = conf.filesOpenCostInBytes
+ val maxSplitBytes = FilePartition.maxSplitBytes(
+ sparkSession, keptFiles.map(_.length + openCostInBytes).sum)
+ FilePartition.getFilePartitions(sparkSession, keptFiles, maxSplitBytes).toArray
+ }
+
+ // Read when the scan is built, as `FileScan` reads it for the names `partitions` matches, so the
+ // runtime filters bind the same way even if the session's setting changes before they arrive.
+ private val isCaseSensitive = conf.caseSensitiveAnalysis
+
+ private def normalizePartitionName(name: String): String = {
+ if (isCaseSensitive) name else name.toLowerCase(Locale.ROOT)
+ }
Review Comment:
Done in e721ae7. "DPP and a scalar subquery prune input partitions at runtime for v2 text" uses its own single-column text fixture, and `AvroV2Suite` has "SPARK-30628: Avro source v2: runtime filters prune partitions". Both check that only partition 7 is read.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,787 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+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
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, FileScan, RuntimePartitionPruningFileScan}
+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.{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 {
+
+ // 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]
+
+ 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 partition value the runtime filter cannot evaluate fails the query") {
+ // A partition value the filter cannot be evaluated against has to fail the query: keeping the
+ // file would return rows that nothing filters, since the filter is declared fully pushed. What
+ // this pins is that the failure is not swallowed. The exposure is the same one a compile-time
+ // partition filter over the same column has, and it is why this path must not gain the
+ // eval-failure tolerance the iterative PartitionPredicate path has. That no FilterExec is left
+ // above the scan is pinned separately.
+ withDppV2Conf {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ 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")
+ val e = intercept[SparkThrowable] {
+ sql(
+ """SELECT f.id FROM fact_str f
+ |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)""".stripMargin)
+ .collect()
+ }
+ assert(e.getCondition === "CAST_INVALID_INPUT",
+ s"expected the failing cast to surface, got ${e.getCondition}")
+ }
+ }
+ }
+ }
+
+ test("a subclass that narrows the file set in partitions keeps it under a runtime filter") {
+ // The runtime-filter path filters the files `partitions` returned, so a file an override of it
+ // drops stays dropped. Listing the files again instead would hand it back as soon as a filter
+ // fired, and with the filter declared fully pushed nothing above the scan would remove 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]
+ // Keeps only partition directory 3's files and counts how often the scan asks for its
+ // partitions.
+ var partitionsCalls = 0
+ val narrowed = new ParquetScan(scan.sparkSession, scan.hadoopConf, scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema, scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters, scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ partitionsCalls += 1
+ super.partitions.map { part =>
+ part.copy(files = part.files.filter(_.partitionValues.getInt(0) == 3))
+ }.filter(_.files.nonEmpty)
+ }
+ }
+ def partValues(parts: Array[InputPartition]): Seq[Int] =
+ parts.flatMap(_.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0)))
+ .distinct.toSeq
+ val partAttr = AttributeReference("part", IntegerType)()
+ // The plain read path first, as planning does through `supportsColumnar`.
+ assert(partValues(narrowed.planInputPartitions()) === Seq(3))
+ // A filter selecting a directory the override dropped comes back with nothing, and the
+ // calls after it, plain and filtered, still find the directory it kept: no call narrows
+ // what the next one starts from.
+ assert(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(7)))).isEmpty)
+ assert(partValues(narrowed.planInputPartitions()) === Seq(3))
+ assert(partValues(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(3))))) === Seq(3))
+ // All four calls above shared one listing, which is what keeps the plan and the execution
+ // on one snapshot of the file index.
+ assert(partitionsCalls === 1, s"expected one listing, got $partitionsCalls")
+ }
+ }
+ }
+
+ test("partitions runs without holding the scan's monitor") {
+ // An override of `partitions` may hand work to another thread and wait for it. If the memo
+ // held the scan's monitor meanwhile, that thread would block when it tried to initialize one
+ // of the scan's lazy vals.
+ 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 =>
+ val heldLock = mutable.ArrayBuffer.empty[Boolean]
+ val probe = new ParquetScan(scan.sparkSession, scan.hadoopConf, scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema, scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters, scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ heldLock += Thread.holdsLock(this)
+ super.partitions
+ }
+ }
+ firstCall(probe)
+ assert(heldLock.toSeq === Seq(false))
+ }
+ }
+ }
+ }
+
+ 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("scalar subquery on a partition column is not re-evaluated after the scan") {
+ // The file scan declares its partition columns in fullyPushedFilterAttributes: applying such a
+ // filter selects partition directories, so every row the scan returns satisfies it and Spark
+ // drops the post-scan FilterExec, the same treatment a compile-time partition filter gets.
+ 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, got ${rows.length}")
+ val executedPlan = df.queryExecution.executedPlan
+ assert(collect(executedPlan) { case f: FilterExec => f }.isEmpty,
+ "a fully pushed partition filter should leave no FilterExec above the scan, got:\n" +
+ executedPlan.treeString)
+ }
+ }
+ }
+
+ test("a runtime filter is ANDed with the compile-time partition filters, not substituted") {
+ // The two halves are applied in different places: `partitions` puts `partitionFilters` through
+ // `fileIndex.listFiles`, and `planInputPartitionsWithRuntimeFilters` filters what that planned.
+ // Both are the scan's only evaluator: FileScanBuilder.pushFilters returns just the data filters
+ // as post-scan filters, and a partition-column runtime filter is declared fully pushed, so
+ // dropping either half returns rows that nothing filters out.
+ 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") {
+ // `f.PART` is resolved to the schema's `part` before it ever becomes a runtime filter, so this
+ // pins the end-to-end path rather than any name matching: what the scan receives is identical
+ // to the lower-case query's. The test below pins how the scan matches the name.
+ 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, got ${rows.length}")
+ val numOutputRows = factScanOf(df).metrics("numOutputRows").value
+ assert(numOutputRows == 10,
+ s"expected the fact scan to read 10 rows, got $numOutputRows")
+ }
+ }
+ }
+
+ test("a runtime filter the scan cannot apply is rejected rather than silently dropped") {
+ // Unreachable through SQL: PushDownUtils screens a filter against the attributes the scan
+ // declared before handing it over. Calling the scan directly is what pins the guard, which has
+ // to fail rather than drop the filter, since a fully pushed predicate dropped here would be
+ // evaluated nowhere. The reference is matched the way the rest of the scan matches partition
+ // column names, so a mixed-case reference binds rather than being rejected.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT f.id, f.part FROM fact f")
+ df.collect()
+ val scan = fileScanOf(df)
+ val e = intercept[SparkException] {
+ scan.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(AttributeReference("id", LongType)(), Literal(1L))))
+ }
+ assert(e.getCondition === "INTERNAL_ERROR")
+ assert(e.getMessage.contains("can only apply a runtime filter over its read partition"),
+ s"unexpected message: ${e.getMessage}")
+
+ val mixedCase = scan.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(AttributeReference("PART", IntegerType)(), Literal(7))))
+ assert(mixedCase.flatMap(
+ _.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0))).distinct.toSeq
+ === Seq(7))
+ }
+ }
+ }
+
+ test("a fully pushed filter stays exact when one FilePartition spans partition directories") {
+ // The fully-pushed contract used to ask for the attribute's value to be fixed within every
+ // InputPartition. A file scan does not satisfy that: FilePartition packs by size and can hold
+ // files from several partition directories. What makes the declaration sound is that the
+ // predicate chose those directories, so every row in the partition satisfies it even though the
+ // value varies. Pack the surviving directories into one FilePartition to exercise that.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1g",
+ SQLConf.FILES_MIN_PARTITION_NUM.key -> "1",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "0") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ // min(dim_id) over dim_val > 7 is 8, so this keeps partitions 8 and 9.
+ val df = sql(
+ """SELECT f.id FROM fact f
+ |WHERE f.part >= (SELECT min(dim_id) FROM dim WHERE dim_val > 7)""".stripMargin)
+ checkAnswer(df, (0 until 10).flatMap(i => Seq(Row(i * 10 + 8), Row(i * 10 + 9))))
+ val factScan = factScanOf(df)
+ assert(factScan.filteredPartitions.flatten.size == 1,
+ "expected both surviving partition directories to pack into one FilePartition, got " +
+ s"${factScan.filteredPartitions.flatten.size}")
+ val numOutputRows = factScan.metrics("numOutputRows").value
+ assert(numOutputRows == 20,
+ s"expected that one FilePartition to carry both directories' 20 rows, got " +
+ s"$numOutputRows")
+ assert(collect(df.queryExecution.executedPlan) { case f: FilterExec => f }.isEmpty,
+ "expected the partition filter to be fully pushed, leaving no FilterExec")
+ }
+ }
+ }
+ }
+
+ test("runtime pruning packs the surviving files for their own size") {
+ // `partitions` sized its bins for the whole file set, so handing those bins back with files
+ // removed leaves a heavily pruned scan running the few tasks the unpruned set needed. Selecting
+ // the same partitions at compile time packs for the pruned set, so the two layouts must agree.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1g",
+ SQLConf.FILES_MIN_PARTITION_NUM.key -> "3",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "0") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val expected = (0 until 10).flatMap(i => Seq(Row(i * 10 + 8, 8), Row(i * 10 + 9, 9)))
+ // min(dim_id) over dim_val > 7 is 8, so this keeps partitions 8 and 9.
+ val runtime = sql(
+ """SELECT f.id, f.part FROM fact f
+ |WHERE f.part >= (SELECT min(dim_id) FROM dim WHERE dim_val > 7)""".stripMargin)
+ checkAnswer(runtime, expected)
+ val compileTime = sql("SELECT f.id, f.part FROM fact f WHERE f.part >= 8")
+ checkAnswer(compileTime, expected)
+ // What gives the comparison teeth: the whole file set packs two or more of the surviving
+ // files into one bin, so handing those bins back with the other files removed cannot
+ // produce as many partitions as packing the survivors does. That needs the fixture to
+ // write more than one file per partition directory, and this is the assertion that fails
+ // if it stops doing so.
+ val unpruned = sql("SELECT f.id, f.part FROM fact f")
+ unpruned.collect()
+ val unprunedPartitions = factScanOf(unpruned).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val binsHoldingSurvivors =
+ unprunedPartitions.count(_.files.exists(_.partitionValues.getInt(0) >= 8))
+ val survivorsWhenUnpruned =
+ unprunedPartitions.flatMap(_.files).count(_.partitionValues.getInt(0) >= 8)
+ assert(binsHoldingSurvivors < survivorsWhenUnpruned,
+ s"expected the whole set to pack $survivorsWhenUnpruned surviving files into fewer " +
+ s"than $survivorsWhenUnpruned bins, got $binsHoldingSurvivors")
+ val compileTimePartitions = factScanOf(compileTime).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val runtimePartitions = factScanOf(runtime).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val runtimeFiles = runtimePartitions.flatMap(_.files).size
+ // The teeth above were measured on the unpruned scan; this is what ties them to this arm.
+ assert(runtimeFiles === survivorsWhenUnpruned,
+ s"expected this arm to read the $survivorsWhenUnpruned files the unpruned packing " +
+ s"held, got $runtimeFiles")
+ // Both arms must see the same files: the runtime path cannot re-split, so a compile-time
+ // maxSplitBytes below one file's length would cut that arm's files finer and the counts
+ // would differ for a reason this test is not about.
+ assert(runtimeFiles === compileTimePartitions.flatMap(_.files).size,
+ s"expected both arms to read $runtimeFiles files, got " +
+ s"${compileTimePartitions.flatMap(_.files).size} at compile time")
+ // The two arms pack from differently ordered input: `partitions` sorts within a partition
+ // directory, this path sorts the survivors globally, and Next Fit Decreasing is sensitive
+ // to that. What makes the counts comparable is that under this conf no two surviving
+ // files fit in one bin, so both arms give one partition per file whatever the order.
+ assert(runtimePartitions.size === runtimeFiles,
+ s"expected one partition per surviving file, got ${runtimePartitions.size} " +
+ s"for $runtimeFiles files")
+ assert(runtimePartitions.size === compileTimePartitions.size,
+ "expected the runtime-pruned scan to be packed like the compile-time one " +
+ s"(${compileTimePartitions.size} partitions), got ${runtimePartitions.size}")
+ }
+ }
+ }
+ }
+
+ test("a DPP filter degraded to true leaves the v2 file scan unpruned") {
+ // With reuse-only DPP and no broadcast to reuse, the subquery is replaced by true before
+ // execution. What this pins is that such a filter still reaches the scan and the query stays
+ // correct; the row count only records that nothing was pruned, which is what a true predicate
+ // does whether or not `catalystRuntimeFilters` drops it.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ 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()
+ assert(rows.length == 10, s"expected 10 rows after join+filter, got ${rows.length}")
+ val factScan = factScanOf(df)
+ assert(factScan.runtimeFilters.exists {
+ case DynamicPruningExpression(l: Literal) => l == Literal.TrueLiteral
+ case _ => false
+ }, s"expected a degraded DPP filter, got ${factScan.runtimeFilters}")
+ val numOutputRows = factScan.metrics("numOutputRows").value
+ assert(numOutputRows == 100,
+ s"expected the fact scan to read all 100 rows, got $numOutputRows")
+ }
+ }
+ }
+ }
+
+ test("a runtime filter binds to the read partition columns when only some of them are read") {
+ // `partitions` projects each file's partition values down to `readPartitionSchema` when the
+ // query reads a strict subset of the partition columns, so a filter on `p2` has to read slot 0
+ // of those values. Binding by the table's partition schema would read slot 1, which in the
+ // projected values holds `p3`: rows that satisfy the filter are lost, and with the filter
+ // fully pushed, rows that do not are returned.
+ withDppV2Conf {
+ withTempDir { dir =>
+ val tPath = new File(dir, "t").getCanonicalPath
+ spark.range(100).selectExpr("id", "id % 2 AS p1", "id % 5 AS p2", "id % 3 AS p3")
+ .write.format("parquet").partitionBy("p1", "p2", "p3").save(tPath)
+ spark.read.format("parquet").schema("id long, p1 int, p2 int, p3 int").load(tPath)
+ .createOrReplaceTempView("t")
+ writeFactAndDim(dir)
+ val expected = (0 until 100).filter(_ % 5 == 2).map(i => Row(i.toLong, i % 3))
+ Seq(
+ """SELECT id, p3 FROM t
+ |WHERE p2 = (SELECT max(dim_id) FROM dim WHERE dim_val = 2)""".stripMargin,
+ """SELECT t.id, t.p3 FROM t JOIN dim d
+ |ON t.p2 = d.dim_id WHERE d.dim_val = 2""".stripMargin).foreach { query =>
+ val df = sql(query)
+ checkAnswer(df, expected)
+ val scan = factScanOf(df)
+ assert(scan.scan.asInstanceOf[FileScan].readPartitionSchema.fieldNames.toSeq ===
+ Seq("p2", "p3"), "the scan must read a strict subset of the partition columns")
+ val numOutputRows = scan.metrics("numOutputRows").value
+ assert(numOutputRows == expected.size,
+ s"expected the scan to read ${expected.size} rows, got $numOutputRows")
+ }
+ }
+ }
+ }
+
+ test("SPARK-38148: no DPP over a partition column a static equality already fixes") {
+ // `FileScanBuilder` pushes the static partition filter fully into the scan, so no `Filter`
+ // above the scan carries it; the cleanup has to read it from the scan relation's pushed
+ // filters, as it reads it from the `Filter` over a V1 relation. V1 is the reference.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ // An int join key, so the pruning key is `part` itself rather than a cast of it.
+ val dimIntPath = new File(dir, "dim_int").getCanonicalPath
+ spark.range(10).selectExpr("CAST(id AS INT) AS dim_id", "CAST(id AS INT) AS dim_val")
+ .write.format("parquet").save(dimIntPath)
+ // (condition, the partitions it keeps, whether DPP stays, whether V1 is checked). The first
+ // six are the SPARK-38148 cases. The last two do not fix the key, so V2 keeps the DPP; they
+ // are checked under V2 only, because V1 reads an equality anywhere in its filter and drops
+ // the DPP.
+ Seq[(String, Int => Boolean, Boolean, Boolean)](
+ ("f.part = 3", _ == 3, false, true),
+ ("3 = f.part", _ == 3, false, true),
+ ("f.part <=> 3", _ == 3, false, true),
+ ("3 <=> f.part", _ == 3, false, true),
+ ("f.part > 3", _ > 3, true, true),
+ ("5 > f.part", _ < 5, true, true),
+ ("f.part <> 3", _ != 3, true, false),
+ ("(f.part = 3 OR f.part = 4)", p => p == 3 || p == 4, true, false)
+ ).foreach { case (condition, keeps, hasDpp, checkV1) =>
+ val query =
+ s"""SELECT /*+ broadcast(d) */ f.id FROM fact f JOIN dim_int d
+ |ON f.part = d.dim_id WHERE $condition AND d.dim_val > 2""".stripMargin
+ val expected = (0 until 100).filter { i =>
+ val part = i % 10
+ part > 2 && keeps(part)
+ }.map(i => Row(i.toLong))
+ (if (checkV1) Seq("parquet", "") else Seq("")).foreach { v1Sources =>
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> v1Sources) {
+ // A temp view keeps the relation it was created with, so create both under this conf.
+ spark.read.format("parquet").schema("id long, part int")
+ .load(new File(dir, "fact").getCanonicalPath).createOrReplaceTempView("fact")
+ spark.read.format("parquet").load(dimIntPath).createOrReplaceTempView("dim_int")
+ val df = sql(query)
+ val optimized = df.queryExecution.optimizedPlan
+ assert(optimized.exists(_.isInstanceOf[DataSourceV2ScanRelation]) ===
+ v1Sources.isEmpty, s"wrong source version for '$v1Sources':\n$optimized")
+ assert(collectDppFilters(optimized).nonEmpty === hasDpp,
+ s"`$condition` with v1 sources '$v1Sources': expected DPP $hasDpp, got\n$optimized")
+ checkAnswer(df, expected)
+ }
+ }
+ }
+ }
+ }
+ }
+
+ Seq("orc", "json", "csv").foreach { format =>
Review Comment:
Done in e721ae7. `DataSourceV2FileSourceDPPSuite` clears the temp views after each test, and `withFactAndDim` in `DataSourceV2FileScanSubclassSuite` uses `withTempView`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-6650290830
@pan3793 has given approval. Would you mind reviewing it again when you get a chance? @dongjoon-hyun Thanks
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4489847020
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala:
##########
@@ -785,22 +793,37 @@ abstract class InMemoryBaseTable(
private val catalystPredicates = ArrayBuffer.empty[CatalystExpression]
private var filterCalls = 0
+ private val plannedPartitions = ArrayBuffer.empty[Seq[InputPartition]]
- override def filter(expressions: Array[CatalystExpression]): Unit = {
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[CatalystExpression]): Array[InputPartition] = {
catalystPredicates ++= expressions
filterCalls += 1
+ // The result comes from `expressions` and the unfiltered `data` alone. Spark calls this once
+ // per scan node, and several nodes can share one scan, so carrying one call's pruning into
+ // the next would drop partitions the second node still needs. `catalystPredicates` and
+ // `filterCalls` only record what happened, for test assertions; `plannedPartitions` is read
+ // back through `readPartitions`, which is what the group-based write replaces.
+ val filtered = prunePartitions(self.data, expressions)
Review Comment:
Done in a76b80e435f, after merging master. `KeyGroupedPartitioningV2FilterRuntimeFilterSuite` runs `KeyGroupedPartitioningRuntimeFilterTests` against an `InMemoryCatalog` whose tables are `InMemoryTableWithV2Filter`, which prunes in place on `=` and `IN` over a single identity partition column. With `BatchScanExec` planning per node again, both SPARK-60752 tests in that suite fail with `Can't zip RDDs with unequal numbers of partitions: List(7, 1)`. I also updated the test comments that say which fixture prunes the shared `Scan`, and the `PlannedScan` doc.
##########
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
Review Comment:
Done in a76b80e435f. It is now "a DPP filter a union branch cannot apply is dropped at planning", and the comment gives the strategy's reason.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4505536866
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
+ fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private val allPartitions = new TransientBestEffortLazyVal(() => filePartitions(listing()))
+
+ // Final so that a subclass of a built-in scan cannot change one read path without the other.
+ final override protected def partitions: Seq[FilePartition] = allPartitions()
+
+ final override def planInputPartitions(): Array[InputPartition] = allPartitions().toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ if (expressions.isEmpty) return allPartitions().toArray
+ // Directory values follow the table's partition schema, not the read one, and bind as V1
+ // binds them: with that schema's types, nullable.
+ val partitionSchema = fileIndex.partitionSchema
+ val fieldIndex = partitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ val predicates = expressions.toSeq.map { e =>
+ Predicate.createInterpreted(e.transform {
+ case a: Attribute =>
+ // PushDownUtils passes only filters over `filterAttributes()`.
+ val i = fieldIndex.getOrElse(normalizePartitionName(a.name),
+ throw SparkException.internalError("A file scan can only apply a runtime filter " +
+ s"over its partition columns ${partitionSchema.fieldNames.mkString("[", ", ", "]")}" +
+ s", got $e"))
+ BoundReference(i, partitionSchema(i).dataType, nullable = true)
+ })
+ }
+ val all = listing()
+ val kept = all.filter { dir =>
+ predicates.forall { predicate =>
+ try predicate.eval(dir.values) catch { case NonFatal(_) => true }
Review Comment:
Done in c1ba7eacfd7. Each call logs at most one warning, with the number of directories it kept because a filter failed on them, and the first failing expression and its exception. "a runtime filter keeps a partition directory it cannot evaluate" checks it. A shared binder can wait for a follow-up, as you said.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-5966973332
@dongjoon-hyun Items 1 and 2 set the direction for several of the others, so before picking one I ran both alternatives against `FileScan` subclasses modeled on the ones you linked.
The probe is a throwaway suite with three subclasses: one shaped like Kyuubi's `HiveScan` (`FileScan with SupportsRuntimeFiltering`, re-listing with the runtime filters in `partitions`), one like SPYT's `YtScan` (the same, except it caches `partitions` and clears the cache in `filter`), and one like Sedona's `ShapefileScan` (its own partition type from `planInputPartitions()`, read by a factory that rejects any other type; Sedona's `RasterScan`, `NetCdfMetadataScan` and `GeoTiffMetadataScan` have the same shape). Each runs a plain `SELECT`, a DPP join and a scalar subquery on the partition column, next to built-in parquet.
| | built-in parquet | Kyuubi-like | SPYT-like | Sedona-like |
|---|---|---|---|---|
| master | no pruning | DPP prunes | DPP prunes | works, no pruning |
| this PR | DPP and subquery prune | every query fails at planning | every query fails at planning | DPP and subquery fail: `Unexpected partition type: FilePartition` |
| A: per-format opt-in | DPP and subquery prune | same as master | same as master | same as master |
| B: overridable hook | DPP and subquery prune | every query fails at planning | every query fails at planning | same as master |
| B + precedence rule, memo behind the flag | DPP and subquery prune | same as master | same as master | same as master |
In the table, every query that ran returned the right rows. The planning failure is the `INTERNAL_ERROR` from the dual-interface check.
Option A moves `SupportsRuntimeCatalystFiltering` and its three methods off `FileScan` into a trait that `ParquetScan`, `OrcScan`, `AvroScan` and `TextBasedFileScan` (CSV, JSON, text) mix in, so any other `FileScan` subclass sees what it sees on master. The memo has to move into that trait as well. When I left it in `FileScan.planInputPartitions()`, the SPYT-like scan's DPP stopped pruning without any error: its fact scan read 100 rows instead of 10, and only the join kept the answer right. `supportsColumnar` fills the memo at planning, and the `filter()` that later clears YtScan's own cache never reaches it. With A, items 4 and 12 also apply only to the scans that mix the trait in.
Option B keeps the interface on `FileScan` and returns no filter attributes unless a subclass turns on an overridable flag, which the built-in scans do. On its own that fixes Sedona but not Kyuubi or SPYT: both dual-interface checks (DataSourceV2Relation.scala:309, PushDownUtils.scala:190) match on the type, and a subclass cannot drop a trait it inherits. Dropping the checks returns wrong rows: the Kyuubi-like and SPYT-like scans then inherit a `fullyPushedFilterAttributes()` that returns their own partition columns, Spark removes the `FilterExec`, and their `filter()` ignores the scalar subquery, so that query returned 100 rows instead of 10. What worked in the probe is the last row: a precedence rule, where a dual-interface scan takes the V2 path and Spark ignores its `fullyPushedFilterAttributes()`, plus the memo behind the same flag. Without the memo change the SPYT-like fact scan read 100 rows under DPP, the same silent loss as leaving the memo in `FileScan` under A. The prec
edence rule replaces the one-interface rule SPARK-58523 set, and the probe covers only the shapes above. The other route is that each such connector moves to the Catalyst interface.
@pan3793, for Kyuubi the second route means `HiveScan`, `KyuubiParquetScan` and `KyuubiOrcScan` drop `SupportsRuntimeFiltering` and each override `planInputPartitionsWithRuntimeFilters`. All three re-list with the runtime filters today, `HiveScan` through `listHiveFiles(partitionFilters ++ runtimeFilters)` and the Parquet and ORC wrappers through a sibling scan built with `partitionFilters ++ runtimeFilters`, while the inherited default filters a listing made without them. The filters would arrive as Catalyst expressions rather than `In` filters, and since `FileScan` declares the partition columns fully pushed, each override has to apply every expression it gets, or return nothing from `fullyPushedFilterAttributes()`. The interface is already in 4.3, with a mutating `filter(Array[Expression])` that this PR replaces, and the connector is one binary for Spark 3.5 through 4.2 (compiled against 3.5, tested on 4.0 to 4.2), so these scans would need separately compiled classes for 4.4,
as this PR is meant for branch-4.x too. Does that match how you see it?
My preference is A: among `FileScan` subclasses it changes behavior only for the built-in scans, and it leaves the one-interface rule alone. A third-party `FileScan` that wants runtime pruning can mix the trait in itself. I'd like to hear what you both think before I change it. I'll answer the other comments inline.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4173544622
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,658 @@
+/*
+ * 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, SparkThrowable}
+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.execution.FilterExec
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, FileScan}
+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.{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 {
+
+ // 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)
Review Comment:
Done in 2acbef9. The new test partitions by `p1, p2, p3` and runs `SELECT id, p3 FROM t WHERE p2 = (SELECT ...)` plus a DPP join on `p2`. It asserts `readPartitionSchema.fieldNames === Seq("p2", "p3")`, the rows, and how many rows the scan read. Binding by `fileIndex.partitionSchema` ordinals fails it.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala:
##########
@@ -40,9 +39,15 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
/**
* Returns attributes this scan can be filtered by at runtime.
*
- * Spark will call [[filter]] if it can derive a runtime filter for any of these attributes.
- * Each reference must resolve against the scan relation output when Spark builds it. Attributes
- * pruned out of [[Scan.readSchema]] fail to resolve.
+ * Spark will call [[planInputPartitionsWithRuntimeFilters]] if it can derive a runtime filter
+ * for any of these attributes. Each reference must resolve against the scan relation output
+ * when Spark builds it. Attributes pruned out of [[Scan.readSchema]] fail to resolve.
+ *
+ * A dynamic partition pruning filter has no post-scan evaluator at all: Spark removes it from
+ * the post-scan filters whether or not the attribute is also in [[fullyPushedFilterAttributes]],
+ * because the join it was derived from already implies it. So an attribute exposed here must be
+ * one this scan can filter on exactly, not approximately. [[fullyPushedFilterAttributes]]
Review Comment:
Done in 2acbef9. The doc now says the scan may return rows that fail the filter, names what drops them for each kind of filter, and says it must never drop a row that satisfies the filter.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,658 @@
+/*
+ * 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, SparkThrowable}
+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.execution.FilterExec
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, FileScan}
+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.{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 {
+
+ // 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): FileScan = factScanOf(df).scan.asInstanceOf[FileScan]
+
+ 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 handed a filter over a column it never declared filterable. Dropping it costs pruning
+ // on that branch and nothing else, since 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: that is what makes this a regression
+ // test rather than a query that happens to work. Both the unpartitioned branch and `dim`
+ // are scans with no partition column, so ask for the one that was handed a live DPP filter.
+ val unusable = collect(df.queryExecution.executedPlan) {
+ case b: BatchScanExec if b.scan.isInstanceOf[FileScan] &&
+ b.scan.asInstanceOf[FileScan].readPartitionSchema.isEmpty &&
+ b.runtimeFilters.exists {
+ case DynamicPruningExpression(child) => !child.isInstanceOf[Literal]
+ case _ => false
+ } => b
+ }
+ assert(unusable.size == 1,
+ "expected the DPP filter to reach exactly the branch that cannot apply it, got " +
+ unusable.map(_.scan.description()).mkString("\n"))
+ // 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 partition value the runtime filter cannot evaluate fails the query") {
+ // A partition value the filter cannot be evaluated against has to fail the query: keeping the
+ // file would return rows that nothing filters, since the filter is declared fully pushed. What
+ // this pins is that the failure is not swallowed. The exposure is the same one a compile-time
+ // partition filter over the same column has, and it is why this path must not gain the
+ // eval-failure tolerance the iterative PartitionPredicate path has. That no FilterExec is left
+ // above the scan is pinned separately.
+ withDppV2Conf {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ 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")
+ val e = intercept[SparkThrowable] {
+ sql(
+ """SELECT f.id FROM fact_str f
+ |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)""".stripMargin)
+ .collect()
+ }
+ assert(e.getCondition === "CAST_INVALID_INPUT",
+ s"expected the failing cast to surface, got ${e.getCondition}")
+ }
+ }
+ }
+ }
+
+ test("a subclass that narrows the file set in partitions keeps it under a runtime filter") {
+ // The runtime-filter path filters what `partitions` returned, so an override of it is honored.
+ // Listing the files again instead would hand the excluded ones back as soon as a filter fired,
+ // and with the filter declared fully pushed nothing above the scan would remove them.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT f.id, f.part FROM fact f")
+ df.collect()
+ val scan = fileScanOf(df).asInstanceOf[ParquetScan]
+ // Keeps the files of one partition directory, which the second filter below does not
+ // select, and counts how often the scan asks for its partitions.
+ var partitionsCalls = 0
+ val narrowed = new ParquetScan(scan.sparkSession, scan.hadoopConf, scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema, scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters, scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ partitionsCalls += 1
+ super.partitions.map { part =>
+ part.copy(files = part.files.filter(_.partitionValues.getInt(0) == 3))
+ }.filter(_.files.nonEmpty)
+ }
+ }
+ val partAttr = AttributeReference("part", IntegerType)()
+ val kept = narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(3))))
+ assert(kept.flatMap(_.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0)))
+ .distinct.toSeq === Seq(3))
+ // And a filter selecting a directory the override dropped comes back with nothing.
+ assert(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(7)))).isEmpty)
+ // Both calls filtered one listing, which is what keeps the plan and the execution on one
+ // snapshot of the file index.
+ assert(partitionsCalls === 1, s"expected one listing, got $partitionsCalls")
+ }
+ }
+ }
+
+ 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("scalar subquery on a partition column is not re-evaluated after the scan") {
+ // FileScan declares its partition columns in fullyPushedFilterAttributes: applying such a
+ // filter selects partition directories, so every row the scan returns satisfies it and Spark
+ // drops the post-scan FilterExec, the same treatment a compile-time partition filter gets.
+ 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, got ${rows.length}")
+ val executedPlan = df.queryExecution.executedPlan
+ assert(collect(executedPlan) { case f: FilterExec => f }.isEmpty,
+ "a fully pushed partition filter should leave no FilterExec above the scan, got:\n" +
+ executedPlan.treeString)
+ }
+ }
+ }
+
+ test("a runtime filter is ANDed with the compile-time partition filters, not substituted") {
+ // The two halves are applied in different places: `partitions` puts `partitionFilters` through
+ // `fileIndex.listFiles`, and `planInputPartitionsWithRuntimeFilters` filters what that planned.
+ // Both are the scan's only evaluator: FileScanBuilder.pushFilters returns just the data filters
+ // as post-scan filters, and a partition-column runtime filter is declared fully pushed, so
+ // dropping either half returns rows that nothing filters out.
+ 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") {
+ // `f.PART` is resolved to the schema's `part` before it ever becomes a runtime filter, so this
+ // pins the end-to-end path rather than any name matching: what the scan receives is identical
+ // to the lower-case query's. The test below is the one that pins matching by exact name.
Review Comment:
Fixed the test comment in 2acbef9 ("pins how the scan matches the name"). I'll fix the description line when I next update it.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala:
##########
@@ -72,35 +80,34 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty
/**
- * Filters this scan using runtime Catalyst expressions.
+ * Plans this scan's input partitions with runtime Catalyst expressions applied.
+ *
+ * The provided expressions must be interpreted as a set of predicates that are ANDed together,
+ * and applied on top of the filters already pushed at query compilation. Implementations may use
+ * them to plan fewer [[InputPartition]]s than [[org.apache.spark.sql.connector.read.Batch]]
+ * would without them.
*
- * The provided expressions must be interpreted as a set of predicates that are ANDed together.
- * Implementations may use the expressions to prune initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s.
+ * Spark calls this method at execution time, and only when at least one runtime filter survives
+ * screening; the partitions it returns replace the ones that scan node would otherwise have
+ * read. Spark may also call `Batch.planInputPartitions()` on the same scan, so neither method may
+ * assume it is the only one used. One scan instance can back several scan nodes (e.g. the two
+ * branches of a group-based UPDATE), each with its own runtime filters, so an implementation
Review Comment:
Done in 2acbef9. The paragraph now says a group-based scan whose write replaces the groups it read has to record what each call returned, and count a node that read through `planInputPartitions()` as having read every group that method returns. Because Spark may also call that method while planning a scan node, it adds that such a call alone does not show a read.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -230,48 +231,64 @@ object PushDownUtils extends Logging {
translatedFiltersPushed || partPredicatesPushed
- case catalystScan: SupportsRuntimeCatalystFiltering if pushableFilters.nonEmpty =>
- // A runtime filter is normally evaluated twice: the source prunes with it, and the
- // FilterExec above the scan applies it again. The two have to agree, so this screen
- // pushes only predicates the source can be trusted to evaluate in Spark's place.
- //
- // But pushing a non-deterministic one would let the source prune on its own coin flip and
- // Spark flip again for the rows that survive, so we handle that specially.
- //
- // Not pushing a filter is safe only while its FilterExec still evaluates it. A fully
- // pushed filter has none, so the source is its only evaluator. That is why
- // DataSourceV2Strategy deletes a FilterExec at planning time only for a filter we push
- // below, and we use the same method (isPushablePartitionFilter) to determine if we
- // should push it.
- //
- // A DPP filter degrades to TrueLiteral once its subquery is pruned away, so it matches
- // every row. The V2 path above drops these implicitly, since translateRuntimeFilterV2
- // returns None; here we push Catalyst expressions directly, so we remove them explicitly.
- val catalystFilters = pushableFilters
- .flatMap(unwrapRuntimeFilterExpression)
- .filterNot(_ == Literal.TrueLiteral)
- if (catalystFilters.nonEmpty) {
- catalystScan.filter(catalystFilters.toArray)
- true
- } else {
- false
- }
-
case _ =>
false
}
}
+ /**
+ * Screens runtime filters for a [[SupportsRuntimeCatalystFiltering]] scan and unwraps them to the
+ * Catalyst predicates the scan should plan with. No translation to connector predicates happens
+ * here: what a scan can apply is the scan's own business.
+ *
+ * Only filters over attributes the scan declared filterable are kept. `DataSourceV2Strategy`
+ * routes a dynamic pruning filter without checking that declaration, and the key of a filter
+ * inserted over one union branch is pushed positionally into the others, so a scan can be handed
+ * a filter over a column it never declared. Dropping it costs pruning and not correctness: its
+ * rows are filtered anyway, by the join a DPP filter was derived from, and by the rewrite
+ * re-applying its own condition for a row-level operation's group filter. For a filter derived
+ * from a scalar subquery this drops nothing: the strategy already required its references to be a
+ * subset of the same declaration.
+ *
+ * The determinism screen is the other half. A scalar subquery filter is evaluated twice, by the
+ * source and by the FilterExec above the scan, so pushing a non-deterministic one would let the
+ * two disagree. `DataSourceV2Strategy` deletes that FilterExec only for a filter this method
+ * keeps, which is why both sites decide with `isPushablePartitionFilter`.
+ *
+ * A DPP filter degrades to TrueLiteral once its subquery is pruned away, so it matches every row.
+ * The V2 path drops these implicitly, since translateRuntimeFilterV2 returns None; here the
+ * expressions reach the scan directly, so we remove them explicitly.
+ */
+ private def catalystRuntimeFilters(
+ scan: SupportsRuntimeCatalystFiltering,
+ runtimeFilters: Seq[Expression],
+ output: Seq[AttributeReference]): Seq[Expression] = {
+ val filterAttrs = V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs(
+ scan.filterAttributes(),
+ output,
+ "filterAttributes()",
+ scan.getClass.getName)
+ runtimeFilters
+ .filter(f => f.references.subsetOf(filterAttrs))
Review Comment:
Done in 2acbef9. For a `SupportsRuntimeCatalystFiltering` scan the strategy keeps only the DPP filters over `relation.runtimeFilterAttrs`, and `catalystRuntimeFilters` keeps its check for a `BatchScanExec` built elsewhere. The union test now asserts that the optimized plan still has the DPP filter over `fact_other`, and that the `fact_other` scan has no runtime filter.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "dongjoon-hyun (via GitHub) via reviews" <re...@spark.apache.org>.
dongjoon-hyun commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4449362227
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala:
##########
@@ -72,35 +80,38 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty
/**
- * Filters this scan using runtime Catalyst expressions.
+ * Plans this scan's input partitions with runtime Catalyst expressions applied.
+ *
+ * The provided expressions must be interpreted as a set of predicates that are ANDed together,
+ * and applied on top of the filters already pushed at query compilation. Implementations may use
+ * them to plan fewer [[InputPartition]]s than [[org.apache.spark.sql.connector.read.Batch]]
+ * would without them.
*
- * The provided expressions must be interpreted as a set of predicates that are ANDed together.
- * Implementations may use the expressions to prune initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s.
+ * Spark calls this method at execution time, and only when at least one runtime filter survives
+ * screening; the partitions it returns replace the ones that scan node would otherwise have
+ * read. Spark may also call `Batch.planInputPartitions()` on the same scan, so neither method may
+ * assume it is the only one used. One scan instance can back several scan nodes (e.g. the two
+ * branches of a group-based UPDATE), each with its own runtime filters, so an implementation
+ * must derive its result from the given expressions alone and must not carry state over from a
+ * previous call. A group-based scan whose write replaces the groups that were read has to record
+ * what each call returned for that write, and count a scan node that read through
+ * `planInputPartitions()` as having read every group that method returns. Spark may also call
+ * that method while planning a scan node, so such a call alone does not show that the node read
+ * through it.
*
* Spark currently tracks runtime-filter eligibility by root attribute. If [[filterAttributes]]
* returns a nested reference, an expression may access another nested field under the same root.
* The scan must match each access against its own partition layout and use only expressions it
* can apply.
*
- * Spark may call this method more than once for the same scan instance: a plan can hold several
- * scan nodes sharing one scan (e.g. the two branches of a group-based UPDATE), and each pushes
- * its own copy of the runtime filters. Implementations must treat successive calls as additive,
- * ANDing the new expressions with those already pushed rather than replacing them.
- *
* If the scan also implements
* [[org.apache.spark.sql.connector.read.SupportsReportPartitioning]], it must preserve
- * the originally reported partitioning during runtime filtering. While applying runtime
- * predicates, the scan may detect that some
- * [[org.apache.spark.sql.connector.read.InputPartition]]s have no matching data, in which
- * case it can either replace the initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s that have no matching data with
- * empty [[org.apache.spark.sql.connector.read.InputPartition]]s, or report only a subset of
- * the original partition values (omitting those with no data) via
- * [[org.apache.spark.sql.connector.read.Batch#planInputPartitions]]. The scan must not
- * report new partition values that were not present in the original partitioning.
- *
- * Note that Spark will call [[Scan.toBatch]] again after filtering the scan at runtime.
+ * the originally reported partitioning. While applying runtime predicates, the scan may detect
+ * that some [[InputPartition]]s have no matching data, in which case it can either replace
+ * them with empty [[InputPartition]]s, or report only a subset of the original partition
+ * values (omitting those with no data). The scan must not report new partition values that
+ * were not present in the original partitioning.
*/
- def filter(expressions: Array[Expression]): Unit
+ def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition]
Review Comment:
**1. [Medium] A scan built against 4.3.0 fails with `AbstractMethodError`**
`filter(expressions: Array[Expression]): Unit` shipped in [v4.3.0](https://github.com/apache/spark/blob/v4.3.0/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala#L105) (2026-10-27), and this PR replaces it with this abstract method. A scan compiled against 4.3.0 that implements the trait has no `planInputPartitionsWithRuntimeFilters`. The first query that gives it a runtime filter fails at execution with `AbstractMethodError` from `PushDownUtils.replanWithRuntimeFilters`. MiMa does not catch it, because `org.apache.spark.sql.internal.*` is excluded.
You mentioned the replacement in [your comparison](https://github.com/apache/spark/pull/55579#issuecomment-5966973332), but the PR description does not list it, and the PR is meant for `branch-4.x` too. Could we keep such a scan working? One way is to give this method a default that does what Spark did for that scan in 4.3.0, and to deprecate `filter` with a default that fails. A new scan then implements only this method:
```scala
def filter(expressions: Array[Expression]): Unit =
throw SparkException.internalError(
s"${getClass.getName} must implement planInputPartitionsWithRuntimeFilters")
def planInputPartitionsWithRuntimeFilters(
expressions: Array[Expression]): Array[InputPartition] = {
filter(expressions)
toBatch.planInputPartitions()
}
```
If you would rather keep the break, could the user-facing section of the description list it?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/TextBasedFileScan.scala:
##########
@@ -27,7 +27,7 @@ import org.apache.spark.util.Utils
abstract class TextBasedFileScan(
sparkSession: SparkSession,
- options: CaseInsensitiveStringMap) extends FileScan {
+ options: CaseInsensitiveStringMap) extends RuntimePartitionPruningFileScan {
Review Comment:
**3. [Low] `TextBasedFileScan` opts in scans outside Spark**
[Your comparison](https://github.com/apache/spark/pull/55579#issuecomment-5966973332) chose the per-format opt-in so that any other `FileScan` subclass "sees what it sees on master", and the scaladoc of `RuntimePartitionPruningFileScan` says the same. `TextBasedFileScan` is also a base class outside Spark, though, so its subclasses get the trait without opting in. A few I found on GitHub are `ExcelScan` in `nightscape/spark-excel`, `WarcScan` in `rubigdata/warc-for-spark` and `SSTScan` in `tikv/migration`.
None of these three overrides `planInputPartitions()` or implements a runtime-filtering interface, so I expect them to start pruning rather than fail. It is still a change they did not ask for, including item 2. Mixing the trait into `CSVScan`, `JsonScan` and `TextScan` instead of `TextBasedFileScan` would keep it to the built-in scans. Otherwise the scaladoc could say that `TextBasedFileScan` subclasses get it too.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
Review Comment:
**2. [Medium] A plugin that replaces `BatchScanExec` can drop the fully pushed filter**
With the partition columns declared here, `DataSourceV2Strategy` removes the `FilterExec` of a scalar-subquery filter on them at planning. From then on only `PushDownUtils.replanWithRuntimeFilters` applies that filter, and only when it is handed this scan. A plugin node that replaces `BatchScanExec` and its scan loses it. For example, the [Spark 5.0 shim of `GpuBatchScanExec`](https://github.com/NVIDIA/cudf-spark/blob/33dcf0491730a8118e9bf064d2b36379d14f63ce/sql-plugin/src/main/spark500/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala#L97-L106) in the RAPIDS Accelerator calls `replanWithRuntimeFilters` with its `GpuParquetScan`, which extends `FileScan` but not this trait:
- `pushRuntimeFilters` returns `false` for that scan, so the node reads every planned partition.
- No `FilterExec` is left above the node, so `SELECT id FROM fact WHERE part = (SELECT max(dim_id) FROM dim)` returns the rows of every partition.
- Its [`TagScanForRuntimeFiltering`](https://github.com/NVIDIA/cudf-spark/blob/33dcf0491730a8118e9bf064d2b36379d14f63ce/sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/TagScanForRuntimeFiltering.scala) keeps a scan on the CPU only when it implements `SupportsRuntimeV2Filtering`, so it does not catch these scans.
I read this code but did not run it. A DPP filter is safe there, because the join still applies it. The silent case comes from the fully pushed declaration, and the `FilterExec` it removes only compares one column with a literal per row, which is cheap next to reading the files. Would it be worth keeping that `FilterExec` for the file scans, by returning nothing from `fullyPushedFilterAttributes()`? If you prefer to keep the declaration, the description could say that a node replacing `BatchScanExec` now has to apply `runtimeFilters` itself.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
+ .flatMap(_.files.filter(file => predicate.eval(file.partitionValues)))
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+ // Pack the survivors instead of handing back the planned partitions with files removed:
+ // `partitions` sized those bins for the whole file set, so a scan that pruned most of its files
+ // would run the same few tasks over a fraction of the data. V1 packs the pruned set the same
+ // way, from `dynamicallySelectedPartitions`.
+ val openCostInBytes = conf.filesOpenCostInBytes
+ val maxSplitBytes = FilePartition.maxSplitBytes(
+ sparkSession, keptFiles.map(_.length + openCostInBytes).sum)
Review Comment:
**4. [Nit] The repack counts one open cost per split, not per file**
This sums `length + openCostInBytes` over the surviving splits. V1 sums `getLen + openCostInBytes` over the surviving files (`FilePartition.maxSplitBytes` over `PartitionDirectory`s), and so does `FileScan.partitions`. A file that `partitions` cut into k splits adds k open costs here instead of one, so `maxSplitBytes` comes out higher by up to `(splits - files) * openCostInBytes / minPartitionNum`.
The effect on the task count is small, but the comment above says this packs the way V1 does. `PartitionedFile.fileSize` holds the whole file length, so counting each file once matches it:
```scala
val totalBytes = keptFiles.map(f => f.filePath -> f.fileSize).distinct
.map(_._2 + openCostInBytes).sum
val maxSplitBytes = FilePartition.maxSplitBytes(sparkSession, totalBytes)
```
If the runtime path moves to pruning before splitting, as suggested in [today's review](https://github.com/apache/spark/pull/55579#pullrequestreview-5733830527), this goes away with it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType, a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
+ .flatMap(_.files.filter(file => predicate.eval(file.partitionValues)))
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+ // Pack the survivors instead of handing back the planned partitions with files removed:
+ // `partitions` sized those bins for the whole file set, so a scan that pruned most of its files
+ // would run the same few tasks over a fraction of the data. V1 packs the pruned set the same
+ // way, from `dynamicallySelectedPartitions`.
+ val openCostInBytes = conf.filesOpenCostInBytes
+ val maxSplitBytes = FilePartition.maxSplitBytes(
+ sparkSession, keptFiles.map(_.length + openCostInBytes).sum)
+ FilePartition.getFilePartitions(sparkSession, keptFiles, maxSplitBytes).toArray
+ }
+
+ // Read when the scan is built, as `FileScan` reads it for the names `partitions` matches, so the
+ // runtime filters bind the same way even if the session's setting changes before they arrive.
+ private val isCaseSensitive = conf.caseSensitiveAnalysis
+
+ private def normalizePartitionName(name: String): String = {
+ if (isCaseSensitive) name else name.toLowerCase(Locale.ROOT)
+ }
Review Comment:
**6. [Nit] `normalizePartitionName` copies `FileScan.normalizeName`**
`isCaseSensitive` and `normalizePartitionName` repeat the private `isCaseSensitive` and `normalizeName` of `FileScan`. The binding is only right while the two copies agree, since `partitions` matches the names with the other one. Making `FileScan.normalizeName` `protected` and calling it here would keep a single rule.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters, the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns (SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what `partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan --
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the survivors group into
+ * partitions is decided here again, for their own size. Where each file was cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition returned satisfies
+ * the expressions even though one partition can still hold files from several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside the read partition
+ // schema has no value to evaluate against, so fail loudly rather than return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType, a.nullable)
Review Comment:
**5. [Nit] The binding takes the type and nullability from the attribute**
V1 binds the same kind of predicate with the partition schema's type and `nullable = true` (`FileSourceScanExec.dynamicallySelectedPartitions` and `PartitioningAwareFileIndex.prunePartitions`), since that is what the partition values hold. Here both come from the filter's attribute. Today they agree, because the partition columns of a V2 file table are always nullable and have the index's types, so I could not build a failing query.
If an attribute ever arrived non-nullable, though, the interpreted `BoundReference` would read a null partition value of an int column as 0, and with the filter fully pushed nothing would catch it. Taking both from the schema removes the dependency:
```scala
case a: Attribute =>
val i = fieldIndex(normalizePartitionName(a.name))
BoundReference(i, readPartitionSchema(i).dataType, nullable = true)
```
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,787 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+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
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, FileScan, RuntimePartitionPruningFileScan}
+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.{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 {
+
+ // 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")
Review Comment:
**8. [Nit] The temp views are never dropped**
`writeFactAndDim` registers `fact` and `dim` with `createOrReplaceTempView`, and so do the tests for `fact5`, `fact_other`, `fact_str`, `t` and `dim_int`, as well as `withFactAndDim` in `DataSourceV2FileScanSubclassSuite`. None of them is dropped, so after a test the views still point at its deleted temp directory. A later test that forgot to write its fixture would then fail on missing files instead of a missing view. The new `ExplainSuite` test already uses `withTempView`, which would work here too.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,787 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+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
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, FileScan, RuntimePartitionPruningFileScan}
+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.{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 {
+
+ // 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]
+
+ 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 partition value the runtime filter cannot evaluate fails the query") {
+ // A partition value the filter cannot be evaluated against has to fail the query: keeping the
+ // file would return rows that nothing filters, since the filter is declared fully pushed. What
+ // this pins is that the failure is not swallowed. The exposure is the same one a compile-time
+ // partition filter over the same column has, and it is why this path must not gain the
+ // eval-failure tolerance the iterative PartitionPredicate path has. That no FilterExec is left
+ // above the scan is pinned separately.
+ withDppV2Conf {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ 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")
+ val e = intercept[SparkThrowable] {
+ sql(
+ """SELECT f.id FROM fact_str f
+ |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val = 7)""".stripMargin)
+ .collect()
+ }
+ assert(e.getCondition === "CAST_INVALID_INPUT",
+ s"expected the failing cast to surface, got ${e.getCondition}")
+ }
+ }
+ }
+ }
+
+ test("a subclass that narrows the file set in partitions keeps it under a runtime filter") {
+ // The runtime-filter path filters the files `partitions` returned, so a file an override of it
+ // drops stays dropped. Listing the files again instead would hand it back as soon as a filter
+ // fired, and with the filter declared fully pushed nothing above the scan would remove 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]
+ // Keeps only partition directory 3's files and counts how often the scan asks for its
+ // partitions.
+ var partitionsCalls = 0
+ val narrowed = new ParquetScan(scan.sparkSession, scan.hadoopConf, scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema, scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters, scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ partitionsCalls += 1
+ super.partitions.map { part =>
+ part.copy(files = part.files.filter(_.partitionValues.getInt(0) == 3))
+ }.filter(_.files.nonEmpty)
+ }
+ }
+ def partValues(parts: Array[InputPartition]): Seq[Int] =
+ parts.flatMap(_.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0)))
+ .distinct.toSeq
+ val partAttr = AttributeReference("part", IntegerType)()
+ // The plain read path first, as planning does through `supportsColumnar`.
+ assert(partValues(narrowed.planInputPartitions()) === Seq(3))
+ // A filter selecting a directory the override dropped comes back with nothing, and the
+ // calls after it, plain and filtered, still find the directory it kept: no call narrows
+ // what the next one starts from.
+ assert(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(7)))).isEmpty)
+ assert(partValues(narrowed.planInputPartitions()) === Seq(3))
+ assert(partValues(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(3))))) === Seq(3))
+ // All four calls above shared one listing, which is what keeps the plan and the execution
+ // on one snapshot of the file index.
+ assert(partitionsCalls === 1, s"expected one listing, got $partitionsCalls")
+ }
+ }
+ }
+
+ test("partitions runs without holding the scan's monitor") {
+ // An override of `partitions` may hand work to another thread and wait for it. If the memo
+ // held the scan's monitor meanwhile, that thread would block when it tried to initialize one
+ // of the scan's lazy vals.
+ 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 =>
+ val heldLock = mutable.ArrayBuffer.empty[Boolean]
+ val probe = new ParquetScan(scan.sparkSession, scan.hadoopConf, scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema, scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters, scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ heldLock += Thread.holdsLock(this)
+ super.partitions
+ }
+ }
+ firstCall(probe)
+ assert(heldLock.toSeq === Seq(false))
+ }
+ }
+ }
+ }
+
+ 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("scalar subquery on a partition column is not re-evaluated after the scan") {
+ // The file scan declares its partition columns in fullyPushedFilterAttributes: applying such a
+ // filter selects partition directories, so every row the scan returns satisfies it and Spark
+ // drops the post-scan FilterExec, the same treatment a compile-time partition filter gets.
+ 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, got ${rows.length}")
+ val executedPlan = df.queryExecution.executedPlan
+ assert(collect(executedPlan) { case f: FilterExec => f }.isEmpty,
+ "a fully pushed partition filter should leave no FilterExec above the scan, got:\n" +
+ executedPlan.treeString)
+ }
+ }
+ }
+
+ test("a runtime filter is ANDed with the compile-time partition filters, not substituted") {
+ // The two halves are applied in different places: `partitions` puts `partitionFilters` through
+ // `fileIndex.listFiles`, and `planInputPartitionsWithRuntimeFilters` filters what that planned.
+ // Both are the scan's only evaluator: FileScanBuilder.pushFilters returns just the data filters
+ // as post-scan filters, and a partition-column runtime filter is declared fully pushed, so
+ // dropping either half returns rows that nothing filters out.
+ 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") {
+ // `f.PART` is resolved to the schema's `part` before it ever becomes a runtime filter, so this
+ // pins the end-to-end path rather than any name matching: what the scan receives is identical
+ // to the lower-case query's. The test below pins how the scan matches the name.
+ 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, got ${rows.length}")
+ val numOutputRows = factScanOf(df).metrics("numOutputRows").value
+ assert(numOutputRows == 10,
+ s"expected the fact scan to read 10 rows, got $numOutputRows")
+ }
+ }
+ }
+
+ test("a runtime filter the scan cannot apply is rejected rather than silently dropped") {
+ // Unreachable through SQL: PushDownUtils screens a filter against the attributes the scan
+ // declared before handing it over. Calling the scan directly is what pins the guard, which has
+ // to fail rather than drop the filter, since a fully pushed predicate dropped here would be
+ // evaluated nowhere. The reference is matched the way the rest of the scan matches partition
+ // column names, so a mixed-case reference binds rather than being rejected.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT f.id, f.part FROM fact f")
+ df.collect()
+ val scan = fileScanOf(df)
+ val e = intercept[SparkException] {
+ scan.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(AttributeReference("id", LongType)(), Literal(1L))))
+ }
+ assert(e.getCondition === "INTERNAL_ERROR")
+ assert(e.getMessage.contains("can only apply a runtime filter over its read partition"),
+ s"unexpected message: ${e.getMessage}")
+
+ val mixedCase = scan.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(AttributeReference("PART", IntegerType)(), Literal(7))))
+ assert(mixedCase.flatMap(
+ _.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0))).distinct.toSeq
+ === Seq(7))
+ }
+ }
+ }
+
+ test("a fully pushed filter stays exact when one FilePartition spans partition directories") {
+ // The fully-pushed contract used to ask for the attribute's value to be fixed within every
+ // InputPartition. A file scan does not satisfy that: FilePartition packs by size and can hold
+ // files from several partition directories. What makes the declaration sound is that the
+ // predicate chose those directories, so every row in the partition satisfies it even though the
+ // value varies. Pack the surviving directories into one FilePartition to exercise that.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1g",
+ SQLConf.FILES_MIN_PARTITION_NUM.key -> "1",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "0") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ // min(dim_id) over dim_val > 7 is 8, so this keeps partitions 8 and 9.
+ val df = sql(
+ """SELECT f.id FROM fact f
+ |WHERE f.part >= (SELECT min(dim_id) FROM dim WHERE dim_val > 7)""".stripMargin)
+ checkAnswer(df, (0 until 10).flatMap(i => Seq(Row(i * 10 + 8), Row(i * 10 + 9))))
+ val factScan = factScanOf(df)
+ assert(factScan.filteredPartitions.flatten.size == 1,
+ "expected both surviving partition directories to pack into one FilePartition, got " +
+ s"${factScan.filteredPartitions.flatten.size}")
+ val numOutputRows = factScan.metrics("numOutputRows").value
+ assert(numOutputRows == 20,
+ s"expected that one FilePartition to carry both directories' 20 rows, got " +
+ s"$numOutputRows")
+ assert(collect(df.queryExecution.executedPlan) { case f: FilterExec => f }.isEmpty,
+ "expected the partition filter to be fully pushed, leaving no FilterExec")
+ }
+ }
+ }
+ }
+
+ test("runtime pruning packs the surviving files for their own size") {
+ // `partitions` sized its bins for the whole file set, so handing those bins back with files
+ // removed leaves a heavily pruned scan running the few tasks the unpruned set needed. Selecting
+ // the same partitions at compile time packs for the pruned set, so the two layouts must agree.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1g",
+ SQLConf.FILES_MIN_PARTITION_NUM.key -> "3",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "0") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val expected = (0 until 10).flatMap(i => Seq(Row(i * 10 + 8, 8), Row(i * 10 + 9, 9)))
+ // min(dim_id) over dim_val > 7 is 8, so this keeps partitions 8 and 9.
+ val runtime = sql(
+ """SELECT f.id, f.part FROM fact f
+ |WHERE f.part >= (SELECT min(dim_id) FROM dim WHERE dim_val > 7)""".stripMargin)
+ checkAnswer(runtime, expected)
+ val compileTime = sql("SELECT f.id, f.part FROM fact f WHERE f.part >= 8")
+ checkAnswer(compileTime, expected)
+ // What gives the comparison teeth: the whole file set packs two or more of the surviving
+ // files into one bin, so handing those bins back with the other files removed cannot
+ // produce as many partitions as packing the survivors does. That needs the fixture to
+ // write more than one file per partition directory, and this is the assertion that fails
+ // if it stops doing so.
+ val unpruned = sql("SELECT f.id, f.part FROM fact f")
+ unpruned.collect()
+ val unprunedPartitions = factScanOf(unpruned).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val binsHoldingSurvivors =
+ unprunedPartitions.count(_.files.exists(_.partitionValues.getInt(0) >= 8))
+ val survivorsWhenUnpruned =
+ unprunedPartitions.flatMap(_.files).count(_.partitionValues.getInt(0) >= 8)
+ assert(binsHoldingSurvivors < survivorsWhenUnpruned,
+ s"expected the whole set to pack $survivorsWhenUnpruned surviving files into fewer " +
+ s"than $survivorsWhenUnpruned bins, got $binsHoldingSurvivors")
+ val compileTimePartitions = factScanOf(compileTime).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val runtimePartitions = factScanOf(runtime).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val runtimeFiles = runtimePartitions.flatMap(_.files).size
+ // The teeth above were measured on the unpruned scan; this is what ties them to this arm.
+ assert(runtimeFiles === survivorsWhenUnpruned,
+ s"expected this arm to read the $survivorsWhenUnpruned files the unpruned packing " +
+ s"held, got $runtimeFiles")
+ // Both arms must see the same files: the runtime path cannot re-split, so a compile-time
+ // maxSplitBytes below one file's length would cut that arm's files finer and the counts
+ // would differ for a reason this test is not about.
+ assert(runtimeFiles === compileTimePartitions.flatMap(_.files).size,
+ s"expected both arms to read $runtimeFiles files, got " +
+ s"${compileTimePartitions.flatMap(_.files).size} at compile time")
+ // The two arms pack from differently ordered input: `partitions` sorts within a partition
+ // directory, this path sorts the survivors globally, and Next Fit Decreasing is sensitive
+ // to that. What makes the counts comparable is that under this conf no two surviving
+ // files fit in one bin, so both arms give one partition per file whatever the order.
+ assert(runtimePartitions.size === runtimeFiles,
+ s"expected one partition per surviving file, got ${runtimePartitions.size} " +
+ s"for $runtimeFiles files")
+ assert(runtimePartitions.size === compileTimePartitions.size,
+ "expected the runtime-pruned scan to be packed like the compile-time one " +
+ s"(${compileTimePartitions.size} partitions), got ${runtimePartitions.size}")
+ }
+ }
+ }
+ }
+
+ test("a DPP filter degraded to true leaves the v2 file scan unpruned") {
+ // With reuse-only DPP and no broadcast to reuse, the subquery is replaced by true before
+ // execution. What this pins is that such a filter still reaches the scan and the query stays
+ // correct; the row count only records that nothing was pruned, which is what a true predicate
+ // does whether or not `catalystRuntimeFilters` drops it.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ 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()
+ assert(rows.length == 10, s"expected 10 rows after join+filter, got ${rows.length}")
+ val factScan = factScanOf(df)
+ assert(factScan.runtimeFilters.exists {
+ case DynamicPruningExpression(l: Literal) => l == Literal.TrueLiteral
+ case _ => false
+ }, s"expected a degraded DPP filter, got ${factScan.runtimeFilters}")
+ val numOutputRows = factScan.metrics("numOutputRows").value
+ assert(numOutputRows == 100,
+ s"expected the fact scan to read all 100 rows, got $numOutputRows")
+ }
+ }
+ }
+ }
+
+ test("a runtime filter binds to the read partition columns when only some of them are read") {
+ // `partitions` projects each file's partition values down to `readPartitionSchema` when the
+ // query reads a strict subset of the partition columns, so a filter on `p2` has to read slot 0
+ // of those values. Binding by the table's partition schema would read slot 1, which in the
+ // projected values holds `p3`: rows that satisfy the filter are lost, and with the filter
+ // fully pushed, rows that do not are returned.
+ withDppV2Conf {
+ withTempDir { dir =>
+ val tPath = new File(dir, "t").getCanonicalPath
+ spark.range(100).selectExpr("id", "id % 2 AS p1", "id % 5 AS p2", "id % 3 AS p3")
+ .write.format("parquet").partitionBy("p1", "p2", "p3").save(tPath)
+ spark.read.format("parquet").schema("id long, p1 int, p2 int, p3 int").load(tPath)
+ .createOrReplaceTempView("t")
+ writeFactAndDim(dir)
+ val expected = (0 until 100).filter(_ % 5 == 2).map(i => Row(i.toLong, i % 3))
+ Seq(
+ """SELECT id, p3 FROM t
+ |WHERE p2 = (SELECT max(dim_id) FROM dim WHERE dim_val = 2)""".stripMargin,
+ """SELECT t.id, t.p3 FROM t JOIN dim d
+ |ON t.p2 = d.dim_id WHERE d.dim_val = 2""".stripMargin).foreach { query =>
+ val df = sql(query)
+ checkAnswer(df, expected)
+ val scan = factScanOf(df)
+ assert(scan.scan.asInstanceOf[FileScan].readPartitionSchema.fieldNames.toSeq ===
+ Seq("p2", "p3"), "the scan must read a strict subset of the partition columns")
+ val numOutputRows = scan.metrics("numOutputRows").value
+ assert(numOutputRows == expected.size,
+ s"expected the scan to read ${expected.size} rows, got $numOutputRows")
+ }
+ }
+ }
+ }
+
+ test("SPARK-38148: no DPP over a partition column a static equality already fixes") {
+ // `FileScanBuilder` pushes the static partition filter fully into the scan, so no `Filter`
+ // above the scan carries it; the cleanup has to read it from the scan relation's pushed
+ // filters, as it reads it from the `Filter` over a V1 relation. V1 is the reference.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ // An int join key, so the pruning key is `part` itself rather than a cast of it.
+ val dimIntPath = new File(dir, "dim_int").getCanonicalPath
+ spark.range(10).selectExpr("CAST(id AS INT) AS dim_id", "CAST(id AS INT) AS dim_val")
+ .write.format("parquet").save(dimIntPath)
+ // (condition, the partitions it keeps, whether DPP stays, whether V1 is checked). The first
+ // six are the SPARK-38148 cases. The last two do not fix the key, so V2 keeps the DPP; they
+ // are checked under V2 only, because V1 reads an equality anywhere in its filter and drops
+ // the DPP.
+ Seq[(String, Int => Boolean, Boolean, Boolean)](
+ ("f.part = 3", _ == 3, false, true),
+ ("3 = f.part", _ == 3, false, true),
+ ("f.part <=> 3", _ == 3, false, true),
+ ("3 <=> f.part", _ == 3, false, true),
+ ("f.part > 3", _ > 3, true, true),
+ ("5 > f.part", _ < 5, true, true),
+ ("f.part <> 3", _ != 3, true, false),
+ ("(f.part = 3 OR f.part = 4)", p => p == 3 || p == 4, true, false)
+ ).foreach { case (condition, keeps, hasDpp, checkV1) =>
+ val query =
+ s"""SELECT /*+ broadcast(d) */ f.id FROM fact f JOIN dim_int d
+ |ON f.part = d.dim_id WHERE $condition AND d.dim_val > 2""".stripMargin
+ val expected = (0 until 100).filter { i =>
+ val part = i % 10
+ part > 2 && keeps(part)
+ }.map(i => Row(i.toLong))
+ (if (checkV1) Seq("parquet", "") else Seq("")).foreach { v1Sources =>
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> v1Sources) {
+ // A temp view keeps the relation it was created with, so create both under this conf.
+ spark.read.format("parquet").schema("id long, part int")
+ .load(new File(dir, "fact").getCanonicalPath).createOrReplaceTempView("fact")
+ spark.read.format("parquet").load(dimIntPath).createOrReplaceTempView("dim_int")
+ val df = sql(query)
+ val optimized = df.queryExecution.optimizedPlan
+ assert(optimized.exists(_.isInstanceOf[DataSourceV2ScanRelation]) ===
+ v1Sources.isEmpty, s"wrong source version for '$v1Sources':\n$optimized")
+ assert(collectDppFilters(optimized).nonEmpty === hasDpp,
+ s"`$condition` with v1 sources '$v1Sources': expected DPP $hasDpp, got\n$optimized")
+ checkAnswer(df, expected)
+ }
+ }
+ }
+ }
+ }
+ }
+
+ Seq("orc", "json", "csv").foreach { format =>
Review Comment:
**7. [Nit] Nothing runs a runtime filter on Avro or text**
This loop covers parquet, orc, json and csv, but `AvroScan` and `TextScan` extend `RuntimePartitionPruningFileScan` too, and no test gives either of them a DPP or scalar-subquery filter. Text needs its own fixture, since the text writer takes a single string column, and Avro would go under `connector/avro`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55579:
URL: https://github.com/apache/spark/pull/55579#issuecomment-5772975648
Thanks @dongjoon-hyun. All three hold up; the first one is fixed, with a caveat about your example.
### 1. Packing
Fixed. The runtime path now collects the surviving `PartitionedFile`s, recomputes `FilePartition.maxSplitBytes` from their own size, and runs `getFilePartitions` over them, instead of returning `part.copy(files = kept)`. A new test compares the layout against the same partitions selected at compile time (`f.part >= 8`); reverting only that hunk fails it with `1 did not equal 4 expected one partition per surviving file, got 1 for 4 files`.
The caveat: this does not change your example, and the reason is worth stating. `getFilePartitions` packs but never splits, so the floor on the task count is the number of `PartitionedFile` pieces, which `partitions` already fixed at the unpruned `maxSplitBytes`. With 64 MB pieces the repack gives each piece its own bin and you still get ~20 tasks, not 200. Where it does match `createReadRDD` exactly is the many-small-files shape: a file below `maxSplitBytes` is never split, so repacking the survivors at the pruned set's `maxSplitBytes` reproduces V1's layout. The new test also measures the boundary between the two: raising `spark.sql.files.minPartitionNum` until the pruned set's `maxSplitBytes` drops below one file's length makes the compile-time arm cut its files finer, and the arms separate (6 partitions against the runtime arm's 4). Closing that means re-splitting, which needs the pre-split file list, i.e. the second read of the index that the earlier revision did and that r40
48745686 asked me to drop. I changed that mechanism without re-checking the packing, which is how this got in. The residual is now in the description next to the missing metrics.
On the per-file evaluation: `InSubqueryExec.eval` delegates to an `InSet` built once from `result`, so each file costs a hash lookup. I left it per file rather than grouping by partition values.
### 2. Stale comment
Fixed, thanks. It described the previous mechanism. It now says that `partitions` applies `partitionFilters` through `fileIndex.listFiles` and `planInputPartitionsWithRuntimeFilters` filters what that planned.
### 3. Description
You are right, and it is narrower still: `PartitioningAwareFileIndex.prunePartitions` evaluates the same expression once per partition directory with no tolerance either, so V1 fails even for a directory holding no files, while the scan here evaluates per file. Reworded to say that the new exposure is a directory contributing no rows, and that whenever a directory does contribute, the cast already fails above the scan or in the join key.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4048745686
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScan.scala:
##########
@@ -186,6 +209,47 @@ trait FileScan extends Scan
partitions.toArray
}
+ /**
+ * The partition columns Spark can derive a runtime filter on (SPARK-30628), restricted to the
+ * ones `readSchema()` still exposes: a reference missing from the scan relation output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it groups by.
+ *
+ * A filter over one of them is applied by selecting partition directories in `buildPartitions`,
+ * the same treatment a compile-time `partitionFilters` entry gets, so the scan evaluates it in
+ * full and Spark does not evaluate it again after the scan -- `FileScanBuilder.pushFilters`
+ * already keeps compile-time partition filters out of the post-scan filters for that reason.
+ * Directory selection matches a predicate's references against `fileIndex.partitionSchema` by
+ * name, so the names reported here have to be the ones that schema uses; both sides derive from
+ * it today, and a divergence would silently drop the predicate instead of failing.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizeName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizeName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ // Directory selection is the only thing that applies these, and it silently ignores a predicate
+ // whose references are not all partition columns -- which, for attributes declared fully
+ // pushed, would leave the predicate evaluated nowhere. Spark screens for that before it gets
+ // here; fail loudly rather than return wrong rows if that ever stops being true. Compare the
+ // way the file index does, by exact name against its own partition schema, so this guard cannot
+ // pass something the index will then drop.
+ val partitionNames = fileIndex.partitionSchema.fieldNames.toSet
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => partitionNames.contains(a.name)))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"partition columns ${fileIndex.partitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ buildPartitions(partitionFilters ++ expressions).toArray
Review Comment:
This path bypasses `partitions` / `planInputPartitions()` entirely. Since every `FileScan` now opts in by default (`filterAttributes()` and `fullyPushedFilterAttributes()` return the read partition columns), an existing 3rd-party `FileScan` subclass that narrows the file set by overriding `partitions` will silently get the excluded files back as soon as DPP or a scalar-subquery filter fires. For the scalar-subquery case, the `FilterExec` above the scan is also removed, so nothing downstream masks it.
This is different from the earlier answer in this thread ("existing sources see no behavior change"), and the PR description's user-facing section doesn't mention it either.
Could we honor the overrides instead? For example, like V1, prune the `PartitionedFile`s already produced by `partitions` by evaluating the runtime filters on their `partitionValues` (bound to `readPartitionSchema`, which `filterAttributes()` already guarantees covers the referenced columns), then repack. That would also remove the duplicate driver-side listing/packing mentioned in the PR description. Alternatively, make the opt-in explicit per built-in format rather than the default for every `FileScan`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4489850200
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionDirectory}
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep, so `partitions` is not used. A scan that overrides it or `planInputPartitions()`, or that
+ * reports a partitioning or an ordering, should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val listingRef = new AtomicReference[Seq[PartitionDirectory]]()
+ @transient private lazy val allPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private def listing: Seq[PartitionDirectory] =
+ once(listingRef)(fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private def allPartitions: Seq[FilePartition] = once(allPartitionsRef)(filePartitions(listing))
+
+ private def once[T >: Null <: AnyRef](ref: AtomicReference[T])(compute: => T): T =
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, compute)
+ ref.get()
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = allPartitions.toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ if (expressions.isEmpty) return allPartitions.toArray
+ val readFields = readPartitionSchema.fieldNames.map(normalizePartitionName).toSet
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => readFields.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
Review Comment:
Done in a76b80e435f. The check is in the `fieldIndex.getOrElse`, against `fileIndex.partitionSchema`, with a comment that `PushDownUtils` passes only filters over `filterAttributes()`. A direct call with a filter over a partition column the scan does not read now evaluates it on the directory values instead of failing.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionDirectory}
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep, so `partitions` is not used. A scan that overrides it or `planInputPartitions()`, or that
+ * reports a partitioning or an ordering, should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val listingRef = new AtomicReference[Seq[PartitionDirectory]]()
+ @transient private lazy val allPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private def listing: Seq[PartitionDirectory] =
+ once(listingRef)(fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private def allPartitions: Seq[FilePartition] = once(allPartitionsRef)(filePartitions(listing))
+
+ private def once[T >: Null <: AnyRef](ref: AtomicReference[T])(compute: => T): T =
Review Comment:
Done in a76b80e435f. Both memos are `TransientBestEffortLazyVal`, and `once`, the holders and the `AtomicReference` import are gone.
##########
connector/avro/src/main/scala/org/apache/spark/sql/v2/avro/AvroScan.scala:
##########
@@ -41,7 +41,7 @@ case class AvroScan(
options: CaseInsensitiveStringMap,
pushedFilters: Array[Filter],
partitionFilters: Seq[Expression] = Seq.empty,
- dataFilters: Seq[Expression] = Seq.empty) extends FileScan {
+ dataFilters: Seq[Expression] = Seq.empty) extends RuntimePartitionPruningFileScan {
Review Comment:
Done in a76b80e435f, in all four scans.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4489845570
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -293,11 +346,19 @@ object PushDownUtils extends Logging {
output: Seq[AttributeReference],
keyedPartitioning: Option[KeyedPartitioning],
originalPartitions: => Seq[InputPartition]): Seq[Option[InputPartition]] = {
- val filtered = pushRuntimeFilters(scan, runtimeFilters, table, output)
+ val (filtered, newPartitions) = scan match {
+ // A scan implementing both runtime-filtering interfaces falls through to pushRuntimeFilters,
+ // which rejects it. The scan is asked even when screening leaves no filter.
+ case catalystScan: SupportsRuntimeCatalystFiltering
Review Comment:
Done in a76b80e435f. `replanWithRuntimeFilters` now leaves a scan that declares only `filter` to `pushRuntimeFilters`, which calls `filter` as 4.3.0 did, so when every runtime filter is screened out the node reads `originalPartitions`. "scan implementing only the deprecated filter -> still pruned, as in Spark 4.3" now replans `JavaDeprecatedFilterScan` with a sentinel `originalPartitions`. It gets the sentinel back when every filter is screened out, without a call to `filter`. The `@param originalPartitions` doc says when they are read.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionDirectory}
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep, so `partitions` is not used. A scan that overrides it or `planInputPartitions()`, or that
+ * reports a partitioning or an ordering, should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val listingRef = new AtomicReference[Seq[PartitionDirectory]]()
+ @transient private lazy val allPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private def listing: Seq[PartitionDirectory] =
+ once(listingRef)(fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private def allPartitions: Seq[FilePartition] = once(allPartitionsRef)(filePartitions(listing))
+
+ private def once[T >: Null <: AnyRef](ref: AtomicReference[T])(compute: => T): T =
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, compute)
+ ref.get()
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = allPartitions.toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ if (expressions.isEmpty) return allPartitions.toArray
+ val readFields = readPartitionSchema.fieldNames.map(normalizePartitionName).toSet
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a => readFields.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime filter over its " +
+ s"read partition columns ${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ // Directory values follow the table's partition schema, not the read one, and bind as V1
+ // binds them: with that schema's types, nullable.
+ val partitionSchema = fileIndex.partitionSchema
+ val fieldIndex = partitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ val predicates = expressions.toSeq.map { e =>
+ Predicate.createInterpreted(e.transform {
+ case a: Attribute =>
+ val i = fieldIndex(normalizePartitionName(a.name))
+ BoundReference(i, partitionSchema(i).dataType, nullable = true)
+ })
+ }
+ val kept = listing.filter { dir =>
+ predicates.forall { predicate =>
+ try predicate.eval(dir.values) catch { case NonFatal(_) => true }
+ }
+ }
+ filePartitions(kept).toArray
Review Comment:
Done in a76b80e435f. When the filter keeps every listed directory, the method returns the memo. The test from item 17 covers this case.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionDirectory}
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep, so `partitions` is not used. A scan that overrides it or `planInputPartitions()`, or that
+ * reports a partitioning or an ordering, should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val listingRef = new AtomicReference[Seq[PartitionDirectory]]()
+ @transient private lazy val allPartitionsRef = new AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private def listing: Seq[PartitionDirectory] =
+ once(listingRef)(fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private def allPartitions: Seq[FilePartition] = once(allPartitionsRef)(filePartitions(listing))
+
+ private def once[T >: Null <: AnyRef](ref: AtomicReference[T])(compute: => T): T =
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, compute)
+ ref.get()
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = allPartitions.toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
Review Comment:
Done in a76b80e435f, with your second option. `hasSelectivePredicate` also counts a likely selective filter in the `pushedFilters` of a `RuntimePartitionPruningFileScan` relation, the same scope as the `CleanupDynamicPruningFilters` change. "DPP for a dimension filtered only on its partition column, as in V1" runs your query shape under V1 and V2 parquet and expects DPP and the same rows from both, and under V2 that the fact scan reads only the five matching partitions.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4505535485
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
Review Comment:
Done in c1ba7eacfd7. Only a scan with runtime-filterable columns keeps the listing. Any other scan packs a listing it does not keep. "only a scan with runtime-filterable columns keeps its file listing" checks that such a scan lists again for a direct call with a runtime filter.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -226,22 +227,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
}
// Screen with the same test pushdown applies, or a filter dropped here and rejected there
// would be evaluated nowhere.
- val fullyPushedRuntimeFilters = scalarSubqueryFilters.filter { f =>
- f.references.subsetOf(relation.fullyPushedRuntimeFilterAttrs) &&
- PushDownUtils.isPushablePartitionFilter(f, includeSubquery = true)
- }
+ val fullyPushedRuntimeFilters = scalarSubqueryFilters.filter(
+ PushDownUtils.isFullyPushedRuntimeFilter(_, relation.fullyPushedRuntimeFilterAttrs))
// dynamicFilters need no such check: a DynamicPruningSubquery over a non-deterministic
// filtering plan is itself non-deterministic, so CleanupDynamicPruningFilters has already
// rewritten it to TrueLiteral by the time we get here.
+ // A Catalyst runtime-filtering scan only applies a filter over attributes it declared, so a
+ // DPP filter over any other column is dropped here rather than planned for nothing. One
+ // PartitionPruning inserted over a union lands in every branch, including one that cannot
+ // use it; the join it was derived from still applies it.
+ val usableDynamicFilters = relation.scan match {
+ case _: SupportsRuntimeCatalystFiltering =>
+ dynamicFilters.filter(_.references.subsetOf(relation.runtimeFilterAttrs))
Review Comment:
It is in the description now, with the bloom-filter bullet.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
+ fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private val allPartitions = new TransientBestEffortLazyVal(() => filePartitions(listing()))
+
+ // Final so that a subclass of a built-in scan cannot change one read path without the other.
+ final override protected def partitions: Seq[FilePartition] = allPartitions()
+
+ final override def planInputPartitions(): Array[InputPartition] = allPartitions().toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ /**
+ * Keeps the listed partition directories whose values satisfy `expressions`, then splits and
+ * packs their files for their own size, as V1 does for `dynamicallySelectedPartitions`.
+ *
+ * A filter that fails to evaluate on a directory keeps it, since the join or the `FilterExec`
+ * above the scan applies the filter again.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ if (expressions.isEmpty) return allPartitions().toArray
+ // Directory values follow the table's partition schema, not the read one, and bind as V1
+ // binds them: with that schema's types, nullable.
+ val partitionSchema = fileIndex.partitionSchema
+ val fieldIndex = partitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ val predicates = expressions.toSeq.map { e =>
+ Predicate.createInterpreted(e.transform {
+ case a: Attribute =>
+ // PushDownUtils passes only filters over `filterAttributes()`.
+ val i = fieldIndex.getOrElse(normalizePartitionName(a.name),
+ throw SparkException.internalError("A file scan can only apply a runtime filter " +
+ s"over its partition columns ${partitionSchema.fieldNames.mkString("[", ", ", "]")}" +
+ s", got $e"))
+ BoundReference(i, partitionSchema(i).dataType, nullable = true)
+ })
+ }
+ val all = listing()
+ val kept = all.filter { dir =>
+ predicates.forall { predicate =>
+ try predicate.eval(dir.values) catch { case NonFatal(_) => true }
+ }
+ }
+ if (kept.length == all.length) allPartitions().toArray else filePartitions(kept).toArray
Review Comment:
Agreed, and I opened SPARK-61174 for it. The reader factories decide from the schema and the conf only (Parquet and ORC, the other four return false), and nothing else reads `inputPartitions` at planning for a file scan, so reporting the mode would skip the packing. Since the merge-base pays the same, I'd keep it out of this PR, so it does not also change how the built-in scans decide columnar reads.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -226,22 +227,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
}
// Screen with the same test pushdown applies, or a filter dropped here and rejected there
// would be evaluated nowhere.
- val fullyPushedRuntimeFilters = scalarSubqueryFilters.filter { f =>
- f.references.subsetOf(relation.fullyPushedRuntimeFilterAttrs) &&
- PushDownUtils.isPushablePartitionFilter(f, includeSubquery = true)
- }
+ val fullyPushedRuntimeFilters = scalarSubqueryFilters.filter(
+ PushDownUtils.isFullyPushedRuntimeFilter(_, relation.fullyPushedRuntimeFilterAttrs))
// dynamicFilters need no such check: a DynamicPruningSubquery over a non-deterministic
// filtering plan is itself non-deterministic, so CleanupDynamicPruningFilters has already
// rewritten it to TrueLiteral by the time we get here.
+ // A Catalyst runtime-filtering scan only applies a filter over attributes it declared, so a
+ // DPP filter over any other column is dropped here rather than planned for nothing. One
+ // PartitionPruning inserted over a union lands in every branch, including one that cannot
+ // use it; the join it was derived from still applies it.
+ val usableDynamicFilters = relation.scan match {
+ case _: SupportsRuntimeCatalystFiltering =>
+ dynamicFilters.filter(_.references.subsetOf(relation.runtimeFilterAttrs))
+ case _ => dynamicFilters
+ }
// A group-based read keeps only the group filter, for the reason above. PartitionPruning
// plans a DynamicPruningSubquery from the join of one read, so it is a filter of one read
// too. With an IN subquery it lands on the read of the updated rows of an UPDATE. With a
// NOT EXISTS subquery the negated condition becomes a semi join, so it lands on the read
// of the rows to copy over.
val runtimeFilters = if (groupBasedRead) {
- dynamicFilters.filterNot(_.isInstanceOf[DynamicPruningSubquery])
+ usableDynamicFilters.filterNot(_.isInstanceOf[DynamicPruningSubquery])
} else {
- dynamicFilters ++ scalarSubqueryFilters
+ usableDynamicFilters ++ scalarSubqueryFilters
Review Comment:
It is in the user-facing section now.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-30628][SQL] Support DPP and subquery partition pruning for V2 file sources [spark]
Posted by "LuciferYang (via GitHub) via reviews" <re...@spark.apache.org>.
LuciferYang commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4505534495
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,108 @@
+/*
+ * 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.v2
+
+import java.util.Locale
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.util.TransientBestEffortLazyVal
+
+/**
+ * A [[FileScan]] that prunes its partition directories with the runtime filters Spark derives over
+ * its partition columns (SPARK-30628). The built-in file formats mix this in.
+ *
+ * Both read paths start from one listing of the file index and split and pack the files they
+ * keep through `filePartitions`, so `partitions` and `planInputPartitions()` are final here. A
+ * scan that reports a partitioning or an ordering should not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with SupportsRuntimeCatalystFiltering {
+
+ /**
+ * The file index listing and its packing, each taken once per scan so planning and execution
+ * read one snapshot. A `lazy val` would compute them while holding the scan's monitor.
+ */
+ private val listing = new TransientBestEffortLazyVal(() =>
+ fileIndex.listFiles(partitionFilters, dataFilters))
+
+ private val allPartitions = new TransientBestEffortLazyVal(() => filePartitions(listing()))
+
+ // Final so that a subclass of a built-in scan cannot change one read path without the other.
+ final override protected def partitions: Seq[FilePartition] = allPartitions()
+
+ final override def planInputPartitions(): Array[InputPartition] = allPartitions().toArray
+
+ /**
+ * The read partition columns that `readSchema()` still exposes; a pushed-down aggregate keeps
+ * only the ones it groups by. None is declared fully pushed, so a scalar subquery filter over
+ * them is evaluated again above the scan.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
Review Comment:
Done in c1ba7eacfd7 and fd6f5adfc85. `RowLevelOperationRuntimeGroupFiltering` skips a group-based operation whose scan is a `RuntimePartitionPruningFileScan`, as on master, and the trait's scaladoc says why. A delta-based operation still gets the group filter, since its write does not depend on which groups were read. "no group filter for a group-based row-level operation over a built-in file scan" runs the rule on a `ReplaceData` over such a scan and expects it unchanged.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -293,11 +347,21 @@ object PushDownUtils extends Logging {
output: Seq[AttributeReference],
keyedPartitioning: Option[KeyedPartitioning],
originalPartitions: => Seq[InputPartition]): Seq[Option[InputPartition]] = {
- val filtered = pushRuntimeFilters(scan, runtimeFilters, table, output)
+ val (filtered, newPartitions) = scan match {
+ // A Catalyst scan is asked even when screening leaves no filter, unless it also implements
+ // SupportsRuntimeV2Filtering, which pushRuntimeFilters rejects, or declares only the
+ // deprecated `filter`, which pushRuntimeFilters calls as Spark 4.3 did.
+ case catalystScan: SupportsRuntimeCatalystFiltering
+ if runtimeFilters.nonEmpty && !catalystScan.isInstanceOf[SupportsRuntimeV2Filtering] &&
+ !implementsOnlyDeprecatedFilter(catalystScan) =>
Review Comment:
Done in c1ba7eacfd7, as you suggested. The scaladoc of `filter` and the description say that a Scala scan recompiled against this version should implement `planInputPartitionsWithRuntimeFilters`, since it otherwise plans again when screening removes every runtime filter. The comments on `LegacyFilterScan` and at L576 say that this Scala fixture takes the general path.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -230,48 +238,92 @@ object PushDownUtils extends Logging {
translatedFiltersPushed || partPredicatesPushed
- case catalystScan: SupportsRuntimeCatalystFiltering if pushableFilters.nonEmpty =>
- // A runtime filter is normally evaluated twice: the source prunes with it, and the
- // FilterExec above the scan applies it again. The two have to agree, so this screen
- // pushes only predicates the source can be trusted to evaluate in Spark's place.
- //
- // But pushing a non-deterministic one would let the source prune on its own coin flip and
- // Spark flip again for the rows that survive, so we handle that specially.
- //
- // Not pushing a filter is safe only while its FilterExec still evaluates it. A fully
- // pushed filter has none, so the source is its only evaluator. That is why
- // DataSourceV2Strategy deletes a FilterExec at planning time only for a filter we push
- // below, and we use the same method (isPushablePartitionFilter) to determine if we
- // should push it.
- //
- // A DPP filter degrades to TrueLiteral once its subquery is pruned away, so it matches
- // every row. The V2 path above drops these implicitly, since translateRuntimeFilterV2
- // returns None; here we push Catalyst expressions directly, so we remove them explicitly.
- val catalystFilters = pushableFilters
- .flatMap(unwrapRuntimeFilterExpression)
- .filterNot(_ == Literal.TrueLiteral)
- if (catalystFilters.nonEmpty) {
- catalystScan.filter(catalystFilters.toArray)
- true
- } else {
- false
+ case catalystScan: SupportsRuntimeCatalystFiltering
+ if runtimeFilters.nonEmpty && implementsOnlyDeprecatedFilter(catalystScan) =>
+ filterWithDeprecatedMethod(catalystScan, runtimeFilters, output)
+
+ case catalystScan: SupportsRuntimeCatalystFiltering =>
+ // Resolved only when a filter could be fully pushed, so a node with no such filter does
+ // not ask the scan.
+ lazy val fullyPushedAttrs = V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs(
+ catalystScan.fullyPushedFilterAttributes(),
+ output,
+ "fullyPushedFilterAttributes()",
+ catalystScan.getClass.getName)
+ val evaluatedOnlyByScan =
+ runtimeFilters.filter(isFullyPushedRuntimeFilter(_, fullyPushedAttrs))
+ if (evaluatedOnlyByScan.nonEmpty) {
+ throw SparkException.internalError(
+ s"${scan.getClass.getName} fully evaluates the runtime filters " +
+ s"${evaluatedOnlyByScan.mkString(", ")}, so they must go through " +
+ "PushDownUtils.replanWithRuntimeFilters, not pushRuntimeFilters.")
}
+ false
case _ =>
false
}
}
+ /**
+ * Screens runtime filters for a [[SupportsRuntimeCatalystFiltering]] scan and unwraps them to the
+ * Catalyst predicates the scan plans with, untranslated.
+ *
+ * A filter is kept only if it references attributes the scan declared filterable and passes
+ * `isPushablePartitionFilter`, the check `DataSourceV2Strategy` uses to drop a pushed filter's
+ * `FilterExec`. A DPP filter degraded to true is removed, as the V2 path does by not translating
+ * it.
+ */
+ private def catalystRuntimeFilters(
+ scan: SupportsRuntimeCatalystFiltering,
+ runtimeFilters: Seq[Expression],
+ output: Seq[AttributeReference]): Seq[Expression] = {
+ val filterAttrs = V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs(
+ scan.filterAttributes(),
+ output,
+ "filterAttributes()",
+ scan.getClass.getName)
+ runtimeFilters
+ .filter(f => f.references.subsetOf(filterAttrs))
+ .filter(isPushablePartitionFilter(_, includeSubquery = true))
+ .flatMap(unwrapRuntimeFilterExpression)
+ .filterNot(_ == Literal.TrueLiteral)
+ }
+
+ /**
+ * Whether `scan` gets `planInputPartitionsWithRuntimeFilters` from the default in
+ * [[SupportsRuntimeCatalystFiltering]], such as a class compiled against Spark 4.3 or a Java
+ * class implementing the trait directly that overrides only `filter`. A Scala class that mixes in
+ * the trait gets a forwarder, so it and its subclasses take the general path.
+ */
+ private def implementsOnlyDeprecatedFilter(scan: SupportsRuntimeCatalystFiltering): Boolean =
+ scan.getClass.getMethod("planInputPartitionsWithRuntimeFilters", classOf[Array[Expression]])
Review Comment:
Done in c1ba7eacfd7. The lookup is behind a `ClassValue`, as in `CodeGenerator.leavesGenCodeToExpression`. A `LinkageError` answers false, and so does any non-fatal exception, since a class loader may reject a class with a `RuntimeException` that `getMethods` passes through unwrapped. Such a scan then takes the general path, whose default still calls `filter` when a filter survives screening. "a scan whose public methods cannot all be resolved still gets its runtime filters" loads a Java scan through a loader that refuses a class in one of its signatures, once with `ClassNotFoundException` and once with `IllegalStateException`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org