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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
