dongjoon-hyun commented on code in PR #55579:
URL: https://github.com/apache/spark/pull/55579#discussion_r4449362227
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala:
##########
@@ -72,35 +80,38 @@ trait SupportsRuntimeCatalystFiltering extends Scan {
def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty
/**
- * Filters this scan using runtime Catalyst expressions.
+ * Plans this scan's input partitions with runtime Catalyst expressions
applied.
+ *
+ * The provided expressions must be interpreted as a set of predicates that
are ANDed together,
+ * and applied on top of the filters already pushed at query compilation.
Implementations may use
+ * them to plan fewer [[InputPartition]]s than
[[org.apache.spark.sql.connector.read.Batch]]
+ * would without them.
*
- * The provided expressions must be interpreted as a set of predicates that
are ANDed together.
- * Implementations may use the expressions to prune initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s.
+ * Spark calls this method at execution time, and only when at least one
runtime filter survives
+ * screening; the partitions it returns replace the ones that scan node
would otherwise have
+ * read. Spark may also call `Batch.planInputPartitions()` on the same scan,
so neither method may
+ * assume it is the only one used. One scan instance can back several scan
nodes (e.g. the two
+ * branches of a group-based UPDATE), each with its own runtime filters, so
an implementation
+ * must derive its result from the given expressions alone and must not
carry state over from a
+ * previous call. A group-based scan whose write replaces the groups that
were read has to record
+ * what each call returned for that write, and count a scan node that read
through
+ * `planInputPartitions()` as having read every group that method returns.
Spark may also call
+ * that method while planning a scan node, so such a call alone does not
show that the node read
+ * through it.
*
* Spark currently tracks runtime-filter eligibility by root attribute. If
[[filterAttributes]]
* returns a nested reference, an expression may access another nested field
under the same root.
* The scan must match each access against its own partition layout and use
only expressions it
* can apply.
*
- * Spark may call this method more than once for the same scan instance: a
plan can hold several
- * scan nodes sharing one scan (e.g. the two branches of a group-based
UPDATE), and each pushes
- * its own copy of the runtime filters. Implementations must treat
successive calls as additive,
- * ANDing the new expressions with those already pushed rather than
replacing them.
- *
* If the scan also implements
* [[org.apache.spark.sql.connector.read.SupportsReportPartitioning]], it
must preserve
- * the originally reported partitioning during runtime filtering. While
applying runtime
- * predicates, the scan may detect that some
- * [[org.apache.spark.sql.connector.read.InputPartition]]s have no matching
data, in which
- * case it can either replace the initially planned
- * [[org.apache.spark.sql.connector.read.InputPartition]]s that have no
matching data with
- * empty [[org.apache.spark.sql.connector.read.InputPartition]]s, or report
only a subset of
- * the original partition values (omitting those with no data) via
- * [[org.apache.spark.sql.connector.read.Batch#planInputPartitions]]. The
scan must not
- * report new partition values that were not present in the original
partitioning.
- *
- * Note that Spark will call [[Scan.toBatch]] again after filtering the scan
at runtime.
+ * the originally reported partitioning. While applying runtime predicates,
the scan may detect
+ * that some [[InputPartition]]s have no matching data, in which case it can
either replace
+ * them with empty [[InputPartition]]s, or report only a subset of the
original partition
+ * values (omitting those with no data). The scan must not report new
partition values that
+ * were not present in the original partitioning.
*/
- def filter(expressions: Array[Expression]): Unit
+ def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition]
Review Comment:
**1. [Medium] A scan built against 4.3.0 fails with `AbstractMethodError`**
`filter(expressions: Array[Expression]): Unit` shipped in
[v4.3.0](https://github.com/apache/spark/blob/v4.3.0/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala#L105)
(2026-10-27), and this PR replaces it with this abstract method. A scan
compiled against 4.3.0 that implements the trait has no
`planInputPartitionsWithRuntimeFilters`. The first query that gives it a
runtime filter fails at execution with `AbstractMethodError` from
`PushDownUtils.replanWithRuntimeFilters`. MiMa does not catch it, because
`org.apache.spark.sql.internal.*` is excluded.
You mentioned the replacement in [your
comparison](https://github.com/apache/spark/pull/55579#issuecomment-5966973332),
but the PR description does not list it, and the PR is meant for `branch-4.x`
too. Could we keep such a scan working? One way is to give this method a
default that does what Spark did for that scan in 4.3.0, and to deprecate
`filter` with a default that fails. A new scan then implements only this method:
```scala
def filter(expressions: Array[Expression]): Unit =
throw SparkException.internalError(
s"${getClass.getName} must implement
planInputPartitionsWithRuntimeFilters")
def planInputPartitionsWithRuntimeFilters(
expressions: Array[Expression]): Array[InputPartition] = {
filter(expressions)
toBatch.planInputPartitions()
}
```
If you would rather keep the break, could the user-facing section of the
description list it?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/TextBasedFileScan.scala:
##########
@@ -27,7 +27,7 @@ import org.apache.spark.util.Utils
abstract class TextBasedFileScan(
sparkSession: SparkSession,
- options: CaseInsensitiveStringMap) extends FileScan {
+ options: CaseInsensitiveStringMap) extends RuntimePartitionPruningFileScan
{
Review Comment:
**3. [Low] `TextBasedFileScan` opts in scans outside Spark**
[Your
comparison](https://github.com/apache/spark/pull/55579#issuecomment-5966973332)
chose the per-format opt-in so that any other `FileScan` subclass "sees what it
sees on master", and the scaladoc of `RuntimePartitionPruningFileScan` says the
same. `TextBasedFileScan` is also a base class outside Spark, though, so its
subclasses get the trait without opting in. A few I found on GitHub are
`ExcelScan` in `nightscape/spark-excel`, `WarcScan` in
`rubigdata/warc-for-spark` and `SSTScan` in `tikv/migration`.
None of these three overrides `planInputPartitions()` or implements a
runtime-filtering interface, so I expect them to start pruning rather than
fail. It is still a change they did not ask for, including item 2. Mixing the
trait into `CSVScan`, `JsonScan` and `TextScan` instead of `TextBasedFileScan`
would keep it to the built-in scans. Otherwise the scaladoc could say that
`TextBasedFileScan` subclasses get it too.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.datasources.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute,
BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference,
NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters,
the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns
(SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their
scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what
`partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning
or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with
SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new
AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both
the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets
runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice
and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would
hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was
stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to
the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation
output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it
groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition
values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan
--
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition
filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] =
filterAttributes()
Review Comment:
**2. [Medium] A plugin that replaces `BatchScanExec` can drop the fully
pushed filter**
With the partition columns declared here, `DataSourceV2Strategy` removes the
`FilterExec` of a scalar-subquery filter on them at planning. From then on only
`PushDownUtils.replanWithRuntimeFilters` applies that filter, and only when it
is handed this scan. A plugin node that replaces `BatchScanExec` and its scan
loses it. For example, the [Spark 5.0 shim of
`GpuBatchScanExec`](https://github.com/NVIDIA/cudf-spark/blob/33dcf0491730a8118e9bf064d2b36379d14f63ce/sql-plugin/src/main/spark500/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala#L97-L106)
in the RAPIDS Accelerator calls `replanWithRuntimeFilters` with its
`GpuParquetScan`, which extends `FileScan` but not this trait:
- `pushRuntimeFilters` returns `false` for that scan, so the node reads
every planned partition.
- No `FilterExec` is left above the node, so `SELECT id FROM fact WHERE part
= (SELECT max(dim_id) FROM dim)` returns the rows of every partition.
- Its
[`TagScanForRuntimeFiltering`](https://github.com/NVIDIA/cudf-spark/blob/33dcf0491730a8118e9bf064d2b36379d14f63ce/sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/TagScanForRuntimeFiltering.scala)
keeps a scan on the CPU only when it implements `SupportsRuntimeV2Filtering`,
so it does not catch these scans.
I read this code but did not run it. A DPP filter is safe there, because the
join still applies it. The silent case comes from the fully pushed declaration,
and the `FilterExec` it removes only compares one column with a literal per
row, which is cheap next to reading the files. Would it be worth keeping that
`FilterExec` for the file scans, by returning nothing from
`fullyPushedFilterAttributes()`? If you prefer to keep the declaration, the
description could say that a node replacing `BatchScanExec` now has to apply
`runtimeFilters` itself.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.datasources.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute,
BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference,
NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters,
the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns
(SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their
scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what
`partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning
or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with
SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new
AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both
the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets
runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice
and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would
hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was
stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to
the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation
output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it
groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition
values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan
--
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition
filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] =
filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from
what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one
listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the
survivors group into
+ * partitions is decided here again, for their own size. Where each file was
cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these
filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition
returned satisfies
+ * the expressions even though one partition can still hold files from
several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the
attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside
the read partition
+ // schema has no value to evaluate against, so fail loudly rather than
return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a
subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a =>
fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime
filter over its " +
+ s"read partition columns
${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType,
a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
+ .flatMap(_.files.filter(file => predicate.eval(file.partitionValues)))
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+ // Pack the survivors instead of handing back the planned partitions with
files removed:
+ // `partitions` sized those bins for the whole file set, so a scan that
pruned most of its files
+ // would run the same few tasks over a fraction of the data. V1 packs the
pruned set the same
+ // way, from `dynamicallySelectedPartitions`.
+ val openCostInBytes = conf.filesOpenCostInBytes
+ val maxSplitBytes = FilePartition.maxSplitBytes(
+ sparkSession, keptFiles.map(_.length + openCostInBytes).sum)
Review Comment:
**4. [Nit] The repack counts one open cost per split, not per file**
This sums `length + openCostInBytes` over the surviving splits. V1 sums
`getLen + openCostInBytes` over the surviving files
(`FilePartition.maxSplitBytes` over `PartitionDirectory`s), and so does
`FileScan.partitions`. A file that `partitions` cut into k splits adds k open
costs here instead of one, so `maxSplitBytes` comes out higher by up to
`(splits - files) * openCostInBytes / minPartitionNum`.
The effect on the task count is small, but the comment above says this packs
the way V1 does. `PartitionedFile.fileSize` holds the whole file length, so
counting each file once matches it:
```scala
val totalBytes = keptFiles.map(f => f.filePath -> f.fileSize).distinct
.map(_._2 + openCostInBytes).sum
val maxSplitBytes = FilePartition.maxSplitBytes(sparkSession, totalBytes)
```
If the runtime path moves to pruning before splitting, as suggested in
[today's
review](https://github.com/apache/spark/pull/55579#pullrequestreview-5733830527),
this goes away with it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.datasources.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute,
BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference,
NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters,
the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns
(SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their
scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what
`partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning
or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with
SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new
AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both
the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets
runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice
and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would
hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was
stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to
the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation
output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it
groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition
values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan
--
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition
filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] =
filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from
what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one
listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the
survivors group into
+ * partitions is decided here again, for their own size. Where each file was
cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these
filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition
returned satisfies
+ * the expressions even though one partition can still hold files from
several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the
attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside
the read partition
+ // schema has no value to evaluate against, so fail loudly rather than
return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a
subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a =>
fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime
filter over its " +
+ s"read partition columns
${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType,
a.nullable)
+ }
+ val predicate = Predicate.createInterpreted(bound)
+ val keptFiles = plannedPartitions
+ .flatMap(_.files.filter(file => predicate.eval(file.partitionValues)))
+ .sortBy(_.length)(implicitly[Ordering[Long]].reverse)
+ // Pack the survivors instead of handing back the planned partitions with
files removed:
+ // `partitions` sized those bins for the whole file set, so a scan that
pruned most of its files
+ // would run the same few tasks over a fraction of the data. V1 packs the
pruned set the same
+ // way, from `dynamicallySelectedPartitions`.
+ val openCostInBytes = conf.filesOpenCostInBytes
+ val maxSplitBytes = FilePartition.maxSplitBytes(
+ sparkSession, keptFiles.map(_.length + openCostInBytes).sum)
+ FilePartition.getFilePartitions(sparkSession, keptFiles,
maxSplitBytes).toArray
+ }
+
+ // Read when the scan is built, as `FileScan` reads it for the names
`partitions` matches, so the
+ // runtime filters bind the same way even if the session's setting changes
before they arrive.
+ private val isCaseSensitive = conf.caseSensitiveAnalysis
+
+ private def normalizePartitionName(name: String): String = {
+ if (isCaseSensitive) name else name.toLowerCase(Locale.ROOT)
+ }
Review Comment:
**6. [Nit] `normalizePartitionName` copies `FileScan.normalizeName`**
`isCaseSensitive` and `normalizePartitionName` repeat the private
`isCaseSensitive` and `normalizeName` of `FileScan`. The binding is only right
while the two copies agree, since `partitions` matches the names with the other
one. Making `FileScan.normalizeName` `protected` and calling it here would keep
a single rule.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RuntimePartitionPruningFileScan.scala:
##########
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.datasources.v2
+
+import java.util.Locale
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute,
BoundReference, Expression, Predicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference,
NamedReference}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+
+/**
+ * A [[FileScan]] that prunes its partition directories with runtime filters,
the DPP filters and
+ * scalar subquery filters Spark derives over its partition columns
(SPARK-30628). The built-in
+ * file formats mix this in; a [[FileScan]] that does not extend one of their
scans keeps its
+ * previous behavior.
+ *
+ * The runtime-filter path returns plain [[FilePartition]]s packed from what
`partitions` planned,
+ * so a scan that overrides `planInputPartitions()` or reports a partitioning
or an ordering should
+ * not mix this in.
+ */
+trait RuntimePartitionPruningFileScan extends FileScan with
SupportsRuntimeCatalystFiltering {
+
+ @transient private lazy val plannedPartitionsRef = new
AtomicReference[Seq[FilePartition]]()
+
+ /**
+ * The partitions `partitions` planned, stored once per scan instance. Both
the plain read path
+ * and the runtime-filter path go through this, so a scan node that gets
runtime filters lists the
+ * files once and filters that listing, rather than reading the index twice
and risking two
+ * snapshots. A subclass still customizes `partitions`.
+ *
+ * `partitions` runs outside the scan's monitor, which a `lazy val` would
hold while it runs. Two
+ * threads asking first at once may both list; both get the result that was
stored.
+ */
+ private def plannedPartitions: Seq[FilePartition] = {
+ val ref = plannedPartitionsRef
+ Option(ref.get()).getOrElse {
+ ref.compareAndSet(null, partitions)
+ ref.get()
+ }
+ }
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ plannedPartitions.toArray
+ }
+
+ /**
+ * The partition columns Spark can derive a runtime filter on, restricted to
the ones
+ * `readSchema()` still exposes: a reference missing from the scan relation
output fails to
+ * resolve, and a pushed-down aggregate keeps only the partition columns it
groups by.
+ *
+ * A filter over one of them is evaluated against each file's partition
values, so the scan
+ * evaluates it in full and Spark does not evaluate it again after the scan
--
+ * `FileScanBuilder.pushFilters` already keeps compile-time partition
filters out of the post-scan
+ * filters for the same reason.
+ */
+ override def filterAttributes(): Array[NamedReference] = {
+ val readFields = readSchema().fieldNames.map(normalizePartitionName).toSet
+ readPartitionSchema.fieldNames
+ .filter(name => readFields.contains(normalizePartitionName(name)))
+ .map(FieldReference.column)
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] =
filterAttributes()
+
+ /**
+ * Plans the files whose partition values satisfy `expressions`, taken from
what `partitions`
+ * planned and bin-packed again for the surviving set.
+ *
+ * Filtering what was planned rather than listing the files again keeps one
listing per scan, and
+ * the files a subclass excluded in `partitions` stay excluded. How the
survivors group into
+ * partitions is decided here again, for their own size. Where each file was
cut into splits is
+ * not: `partitions` cut them for the whole set it listed, before these
filters existed.
+ *
+ * A partition value is fixed within a file, so every row of every partition
returned satisfies
+ * the expressions even though one partition can still hold files from
several partition
+ * directories.
+ */
+ override def planInputPartitionsWithRuntimeFilters(
+ expressions: Array[Expression]): Array[InputPartition] = {
+ val fieldIndex = readPartitionSchema.fieldNames.zipWithIndex
+ .map { case (name, i) => normalizePartitionName(name) -> i }.toMap
+ // Evaluating these is the scan's only chance to apply them, since the
attributes are declared
+ // fully pushed and Spark drops the post-scan filter. A reference outside
the read partition
+ // schema has no value to evaluate against, so fail loudly rather than
return wrong rows. Spark
+ // screens against `filterAttributes()` before it gets here, which is a
subset of that schema.
+ val notApplicable = expressions.filterNot(
+ _.references.forall(a =>
fieldIndex.contains(normalizePartitionName(a.name))))
+ if (notApplicable.nonEmpty) {
+ throw SparkException.internalError("A file scan can only apply a runtime
filter over its " +
+ s"read partition columns
${readPartitionSchema.fieldNames.mkString("[", ", ", "]")}, " +
+ s"got ${notApplicable.mkString(", ")}")
+ }
+ val bound = expressions.reduce(And).transform {
+ case a: Attribute =>
+ BoundReference(fieldIndex(normalizePartitionName(a.name)), a.dataType,
a.nullable)
Review Comment:
**5. [Nit] The binding takes the type and nullability from the attribute**
V1 binds the same kind of predicate with the partition schema's type and
`nullable = true` (`FileSourceScanExec.dynamicallySelectedPartitions` and
`PartitioningAwareFileIndex.prunePartitions`), since that is what the partition
values hold. Here both come from the filter's attribute. Today they agree,
because the partition columns of a V2 file table are always nullable and have
the index's types, so I could not build a failing query.
If an attribute ever arrived non-nullable, though, the interpreted
`BoundReference` would read a null partition value of an int column as 0, and
with the filter fully pushed nothing would catch it. Taking both from the
schema removes the dependency:
```scala
case a: Attribute =>
val i = fieldIndex(normalizePartitionName(a.name))
BoundReference(i, readPartitionSchema(i).dataType, nullable = true)
```
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,787 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector
+
+import java.io.File
+
+import scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+import org.apache.spark.sql.{DataFrame, QueryTest, Row}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
DynamicPruningExpression, DynamicPruningSubquery, EqualTo, Expression, Literal}
+import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.FilterExec
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper,
DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec,
DataSourceV2ScanRelation, FileScan, RuntimePartitionPruningFileScan}
+import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{IntegerType, LongType}
+
+/**
+ * Tests for SPARK-30628: subquery partition pruning and DPP for V2 file
sources.
+ *
+ * These tests exercise V2 file scans (e.g. ParquetScan) -- distinct from the
in-memory V2
+ * catalog used by [[org.apache.spark.sql.DynamicPartitionPruningV2Suite]] and
from the V2
+ * iterative-pushdown work covered by
[[DataSourceV2EnhancedRuntimePartitionFilterSuite]].
+ */
+abstract class DataSourceV2FileSourceDPPSuiteBase extends QueryTest
+ with SharedSparkSession
+ with AdaptiveSparkPlanHelper {
+
+ // Standard V2-only conf for DPP tests on file sources. AQE is set by the
concrete subclasses,
+ // which mix in DisableAdaptiveExecutionSuite /
EnableAdaptiveExecutionSuite: the DPP subquery is
+ // planned by PlanDynamicPruningFilters without AQE and by
PlanAdaptiveDynamicPruningFilters with
+ // it, and those are two separate rules.
+ private def withDppV2Conf[T](thunk: => T): T = {
+ withSQLConf(
+ SQLConf.USE_V1_SOURCE_LIST.key -> "",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key ->
"2")(thunk)
+ }
+
+ // Writes a partitioned `fact` and an unpartitioned `dim` table in `format`
and registers them
+ // as temp views. `fact` has 100 rows across 10 partitions (10 rows each);
`dim` has 10 rows
+ // with `dim_id` and `dim_val` both in [0, 10). The read schema is given
explicitly because csv
+ // infers strings, and a cast on the join key changes whether DPP fires at
all.
+ private def writeFactAndDim(dir: File, format: String = "parquet"): Unit = {
+ val factPath = new File(dir, "fact").getCanonicalPath
+ val dimPath = new File(dir, "dim").getCanonicalPath
+ spark.range(100)
+ .selectExpr("id", "id % 10 AS part")
+ .write.format(format).partitionBy("part").save(factPath)
+ spark.read.format(format).schema("id long, part int").load(factPath)
+ .createOrReplaceTempView("fact")
+ spark.range(10)
+ .selectExpr("id AS dim_id", "id AS dim_val")
+ .write.format(format).save(dimPath)
+ spark.read.format(format).schema("dim_id long, dim_val long").load(dimPath)
+ .createOrReplaceTempView("dim")
Review Comment:
**8. [Nit] The temp views are never dropped**
`writeFactAndDim` registers `fact` and `dim` with `createOrReplaceTempView`,
and so do the tests for `fact5`, `fact_other`, `fact_str`, `t` and `dim_int`,
as well as `withFactAndDim` in `DataSourceV2FileScanSubclassSuite`. None of
them is dropped, so after a test the views still point at its deleted temp
directory. A later test that forgot to write its fixture would then fail on
missing files instead of a missing view. The new `ExplainSuite` test already
uses `withTempView`, which would work here too.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2FileSourceDPPSuite.scala:
##########
@@ -0,0 +1,787 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector
+
+import java.io.File
+
+import scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+import org.apache.spark.sql.{DataFrame, QueryTest, Row}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
DynamicPruningExpression, DynamicPruningSubquery, EqualTo, Expression, Literal}
+import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan}
+import org.apache.spark.sql.connector.read.InputPartition
+import org.apache.spark.sql.execution.FilterExec
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper,
DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.datasources.FilePartition
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec,
DataSourceV2ScanRelation, FileScan, RuntimePartitionPruningFileScan}
+import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{IntegerType, LongType}
+
+/**
+ * Tests for SPARK-30628: subquery partition pruning and DPP for V2 file
sources.
+ *
+ * These tests exercise V2 file scans (e.g. ParquetScan) -- distinct from the
in-memory V2
+ * catalog used by [[org.apache.spark.sql.DynamicPartitionPruningV2Suite]] and
from the V2
+ * iterative-pushdown work covered by
[[DataSourceV2EnhancedRuntimePartitionFilterSuite]].
+ */
+abstract class DataSourceV2FileSourceDPPSuiteBase extends QueryTest
+ with SharedSparkSession
+ with AdaptiveSparkPlanHelper {
+
+ // Standard V2-only conf for DPP tests on file sources. AQE is set by the
concrete subclasses,
+ // which mix in DisableAdaptiveExecutionSuite /
EnableAdaptiveExecutionSuite: the DPP subquery is
+ // planned by PlanDynamicPruningFilters without AQE and by
PlanAdaptiveDynamicPruningFilters with
+ // it, and those are two separate rules.
+ private def withDppV2Conf[T](thunk: => T): T = {
+ withSQLConf(
+ SQLConf.USE_V1_SOURCE_LIST.key -> "",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key ->
"2")(thunk)
+ }
+
+ // Writes a partitioned `fact` and an unpartitioned `dim` table in `format`
and registers them
+ // as temp views. `fact` has 100 rows across 10 partitions (10 rows each);
`dim` has 10 rows
+ // with `dim_id` and `dim_val` both in [0, 10). The read schema is given
explicitly because csv
+ // infers strings, and a cast on the join key changes whether DPP fires at
all.
+ private def writeFactAndDim(dir: File, format: String = "parquet"): Unit = {
+ val factPath = new File(dir, "fact").getCanonicalPath
+ val dimPath = new File(dir, "dim").getCanonicalPath
+ spark.range(100)
+ .selectExpr("id", "id % 10 AS part")
+ .write.format(format).partitionBy("part").save(factPath)
+ spark.read.format(format).schema("id long, part int").load(factPath)
+ .createOrReplaceTempView("fact")
+ spark.range(10)
+ .selectExpr("id AS dim_id", "id AS dim_val")
+ .write.format(format).save(dimPath)
+ spark.read.format(format).schema("dim_id long, dim_val long").load(dimPath)
+ .createOrReplaceTempView("dim")
+ }
+
+ // The BatchScanExec reading `fact`, the only partitioned table of the two.
The lookup descends
+ // into AQE query stages, so it finds the scan whichever stage it ended up
in.
+ private def factScanOf(df: DataFrame): BatchScanExec = {
+ val executedPlan = df.queryExecution.executedPlan
+ collectFirst(executedPlan) {
+ case b: BatchScanExec if b.scan.isInstanceOf[FileScan] &&
+ b.scan.asInstanceOf[FileScan].readPartitionSchema.nonEmpty => b
+ }.getOrElse(fail("no fact BatchScanExec found in:\n" +
executedPlan.treeString))
+ }
+
+ private def fileScanOf(df: DataFrame): RuntimePartitionPruningFileScan =
+ factScanOf(df).scan.asInstanceOf[RuntimePartitionPruningFileScan]
+
+ private def collectDppFilters(plan: LogicalPlan): Seq[Expression] = {
+ plan.collect {
+ case f: Filter if
f.condition.exists(_.isInstanceOf[DynamicPruningSubquery]) => f.condition
+ }
+ }
+
+ test("DPP inserts DynamicPruningSubquery for V2 parquet partitioned table") {
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql(
+ """SELECT f.id FROM fact f JOIN dim d
+ |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin)
+ val optimized = df.queryExecution.optimizedPlan
+ assert(collectDppFilters(optimized).nonEmpty,
+ "expected DynamicPruningSubquery in optimized plan over V2 parquet,
got plan:\n" +
+ optimized.treeString)
+ }
+ }
+ }
+
+ test("DPP does not fire when join key is a data column (non-partition)") {
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ // Join on f.id (data column), not f.part (partition column). `f.part`
stays in the
+ // projection so the fact scan still reads a partition column: without
it the scan declares
+ // nothing filterable for a reason that has nothing to do with the
join key.
+ val df = sql(
+ """SELECT f.id, f.part FROM fact f JOIN dim d
+ |ON f.id = d.dim_id WHERE d.dim_val = 7""".stripMargin)
+ val optimized = df.queryExecution.optimizedPlan
+ assert(collectDppFilters(optimized).isEmpty,
+ "DPP should not fire on non-partition join keys, got plan:\n" +
+ optimized.treeString)
+ assert(df.collect().length === 1)
+ // And the scan is one that could have been pruned, had the join key
been its partition
+ // column: a filterAttributes() reporting every read column would make
DPP fire above, and
+ // the fact scan would read fewer than all 100 rows.
+ assert(fileScanOf(df).filterAttributes().flatMap(_.fieldNames).toSeq
=== Seq("part"))
+ val numOutputRows = factScanOf(df).metrics("numOutputRows").value
+ assert(numOutputRows === 100,
+ s"expected the fact scan to read every row, got $numOutputRows")
+ }
+ }
+ }
+
+ test("DPP fires when partitioned fact is on the right side of the join") {
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ // dim on left, fact on right; join on partition column
+ val df = sql(
+ """SELECT f.id FROM dim d JOIN fact f
+ |ON d.dim_id = f.part WHERE d.dim_val = 7""".stripMargin)
+ val optimized = df.queryExecution.optimizedPlan
+ assert(collectDppFilters(optimized).nonEmpty,
+ "expected DPP on right-side partitioned fact, got plan:\n" +
+ optimized.treeString)
+ checkAnswer(df, (0 until 10).map(i => Row(i * 10 + 7)))
+ val numOutputRows = factScanOf(df).metrics("numOutputRows").value
+ assert(numOutputRows === 10,
+ s"expected the fact scan to read only partition 7, got
$numOutputRows rows")
+ }
+ }
+ }
+
+ test("DPP filter reaches BatchScanExec.runtimeFilters as
DynamicPruningExpression") {
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql(
+ """SELECT f.id FROM fact f JOIN dim d
+ |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin)
+ df.collect()
+ val factScan = factScanOf(df)
+ // A DPP filter that degraded to true is also a
DynamicPruningExpression, so match on the
+ // payload: only a subquery child means the filter can actually prune
anything.
+ val dppFilters = factScan.runtimeFilters.collect {
+ case d @ DynamicPruningExpression(child) if
!child.isInstanceOf[Literal] => d
+ }
+ assert(dppFilters.nonEmpty,
+ "expected the fact scan's runtimeFilters to carry a
DynamicPruningExpression over a " +
+ s"subquery, got ${factScan.runtimeFilters}")
+ }
+ }
+ }
+
+ test("DPP prunes input partitions at runtime for v2 parquet") {
+ // Fact has 100 rows across 10 partitions (10 rows each). DPP filter
d.dim_val = 7 selects
+ // exactly partition 7 -> 10 rows. Without runtime pruning, all 100 rows
are read.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql(
+ """SELECT f.id FROM fact f JOIN dim d
+ |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin)
+ val rows = df.collect()
+ // Correctness: 10 fact rows match (partition 7), each multiplied by 1
dim row.
+ assert(rows.length == 10,
+ s"expected 10 rows after join+filter, got ${rows.length}")
+ // Pruning: fact scan should have read only the 10 rows in partition 7.
+ val numOutputRows = factScanOf(df).metrics("numOutputRows").value
+ assert(numOutputRows == 10,
+ s"expected fact scan to read 10 rows after DPP, got $numOutputRows.
plan:\n" +
+ df.queryExecution.executedPlan.treeString)
+ }
+ }
+ }
+
+ test("DPP that prunes every partition leaves the scan with no input
partition") {
+ // `fact5` has partitions 0..4 while the dim row selects 7, so the build
side is not empty and
+ // the surviving partition set is. That combination is what reaches the
zero-partition
+ // DataSourceRDD: BatchScanExec's EmptyRDDWithPartitions branch keys off
SinglePartition, which
+ // a file scan never reports.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val factPath = new File(dir, "fact5").getCanonicalPath
+ spark.range(50).selectExpr("id", "id % 5 AS part")
+ .write.format("parquet").partitionBy("part").save(factPath)
+ spark.read.format("parquet").schema("id long, part int").load(factPath)
+ .createOrReplaceTempView("fact5")
+ val df = sql(
+ """SELECT f.id FROM fact5 f JOIN dim d
+ |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin)
+ assert(df.collect().isEmpty)
+ val factScan = factScanOf(df)
+ assert(factScan.filteredPartitions.flatten.isEmpty,
+ s"expected no input partition, got
${factScan.filteredPartitions.flatten.size}")
+ // What the assertion above does not see: BatchScanExec turning an
empty partition set into
+ // an RDD with one empty partition still satisfies it.
+ assert(factScan.inputRDD.partitions.isEmpty,
+ s"expected an RDD with no partition, got
${factScan.inputRDD.partitions.length}")
+ val numOutputRows = factScan.metrics("numOutputRows").value
+ assert(numOutputRows == 0,
+ s"expected the fact scan to read nothing, got $numOutputRows rows")
+ }
+ }
+ }
+
+ test("a DPP filter a union branch cannot apply is dropped instead of failing
the query") {
+ // PartitionPruning resolves a union key through the first branch only,
and the filter it
+ // inserts is then pushed positionally into every branch. So the
unpartitioned branch's scan
+ // is offered a filter over a column it never declared filterable.
Dropping it costs nothing,
+ // since that branch cannot prune on it and the join re-applies it;
letting it through means an
+ // INTERNAL_ERROR from a query that runs today.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val otherPath = new File(dir, "fact_other").getCanonicalPath
+ spark.range(100).selectExpr("id", "id % 10 AS other")
+ .write.format("parquet").save(otherPath)
+ spark.read.format("parquet").schema("id long, other
long").load(otherPath)
+ .createOrReplaceTempView("fact_other")
+ val df = sql(
+ """SELECT u.id FROM (
+ | SELECT id, part FROM fact UNION ALL SELECT id, other FROM
fact_other
+ |) u JOIN dim d ON u.part = d.dim_id WHERE d.dim_val =
7""".stripMargin)
+ val expected = (0 until 10).map(i => Row(i * 10 + 7))
+ checkAnswer(df, expected ++ expected)
+ // The filter did reach the branch that cannot use it in the optimized
plan, which is what
+ // makes this a regression test rather than a query that happens to
work.
+ val optimizedBranchFilters = df.queryExecution.optimizedPlan.collect {
+ case f: Filter if
f.condition.exists(_.isInstanceOf[DynamicPruningSubquery]) &&
+ f.child.output.exists(_.name == "other") => f
+ }
+ assert(optimizedBranchFilters.size == 1,
+ "expected the DPP filter over the branch that cannot apply it, got
plan:\n" +
+ df.queryExecution.optimizedPlan.treeString)
+ // The strategy dropped it at planning, so the scan of that branch has
no runtime filter.
+ val otherScans = collect(df.queryExecution.executedPlan) {
+ case b: BatchScanExec if b.output.exists(_.name == "other") => b
+ }
+ assert(otherScans.size == 1, s"expected one scan of fact_other, got
$otherScans")
+ assert(otherScans.head.runtimeFilters.isEmpty,
+ s"expected no runtime filter on fact_other, got
${otherScans.head.runtimeFilters}")
+ // And the branch that can use it still pruned.
+ val numOutputRows = factScanOf(df).metrics("numOutputRows").value
+ assert(numOutputRows == 10,
+ s"expected the partitioned branch to read 10 rows, got
$numOutputRows")
+ }
+ }
+ }
+
+ test("a partition value the runtime filter cannot evaluate fails the query")
{
+ // A partition value the filter cannot be evaluated against has to fail
the query: keeping the
+ // file would return rows that nothing filters, since the filter is
declared fully pushed. What
+ // this pins is that the failure is not swallowed. The exposure is the
same one a compile-time
+ // partition filter over the same column has, and it is why this path must
not gain the
+ // eval-failure tolerance the iterative PartitionPredicate path has. That
no FilterExec is left
+ // above the scan is pinned separately.
+ withDppV2Conf {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val strPath = new File(dir, "fact_str").getCanonicalPath
+ spark.range(10)
+ .selectExpr("id", "CASE WHEN id = 3 THEN 'abc' ELSE CAST(id AS
STRING) END AS part")
+ .write.format("parquet").partitionBy("part").save(strPath)
+ spark.read.format("parquet").schema("id long, part
string").load(strPath)
+ .createOrReplaceTempView("fact_str")
+ val e = intercept[SparkThrowable] {
+ sql(
+ """SELECT f.id FROM fact_str f
+ |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val =
7)""".stripMargin)
+ .collect()
+ }
+ assert(e.getCondition === "CAST_INVALID_INPUT",
+ s"expected the failing cast to surface, got ${e.getCondition}")
+ }
+ }
+ }
+ }
+
+ test("a subclass that narrows the file set in partitions keeps it under a
runtime filter") {
+ // The runtime-filter path filters the files `partitions` returned, so a
file an override of it
+ // drops stays dropped. Listing the files again instead would hand it back
as soon as a filter
+ // fired, and with the filter declared fully pushed nothing above the scan
would remove it.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT f.id, f.part FROM fact f")
+ df.collect()
+ val scan = fileScanOf(df).asInstanceOf[ParquetScan]
+ // Keeps only partition directory 3's files and counts how often the
scan asks for its
+ // partitions.
+ var partitionsCalls = 0
+ val narrowed = new ParquetScan(scan.sparkSession, scan.hadoopConf,
scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema,
scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters,
scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ partitionsCalls += 1
+ super.partitions.map { part =>
+ part.copy(files = part.files.filter(_.partitionValues.getInt(0)
== 3))
+ }.filter(_.files.nonEmpty)
+ }
+ }
+ def partValues(parts: Array[InputPartition]): Seq[Int] =
+
parts.flatMap(_.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0)))
+ .distinct.toSeq
+ val partAttr = AttributeReference("part", IntegerType)()
+ // The plain read path first, as planning does through
`supportsColumnar`.
+ assert(partValues(narrowed.planInputPartitions()) === Seq(3))
+ // A filter selecting a directory the override dropped comes back with
nothing, and the
+ // calls after it, plain and filtered, still find the directory it
kept: no call narrows
+ // what the next one starts from.
+ assert(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(7)))).isEmpty)
+ assert(partValues(narrowed.planInputPartitions()) === Seq(3))
+ assert(partValues(narrowed.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(partAttr, Literal(3))))) === Seq(3))
+ // All four calls above shared one listing, which is what keeps the
plan and the execution
+ // on one snapshot of the file index.
+ assert(partitionsCalls === 1, s"expected one listing, got
$partitionsCalls")
+ }
+ }
+ }
+
+ test("partitions runs without holding the scan's monitor") {
+ // An override of `partitions` may hand work to another thread and wait
for it. If the memo
+ // held the scan's monitor meanwhile, that thread would block when it
tried to initialize one
+ // of the scan's lazy vals.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT f.id, f.part FROM fact f")
+ df.collect()
+ val scan = fileScanOf(df).asInstanceOf[ParquetScan]
+ val partAttr = AttributeReference("part", IntegerType)()
+ Seq[ParquetScan => Array[InputPartition]](
+ _.planInputPartitions(),
+ _.planInputPartitionsWithRuntimeFilters(Array(EqualTo(partAttr,
Literal(3))))
+ ).foreach { firstCall =>
+ val heldLock = mutable.ArrayBuffer.empty[Boolean]
+ val probe = new ParquetScan(scan.sparkSession, scan.hadoopConf,
scan.fileIndex,
+ scan.dataSchema, scan.readDataSchema, scan.readPartitionSchema,
scan.pushedFilters,
+ scan.options, scan.pushedAggregate, scan.partitionFilters,
scan.dataFilters,
+ scan.pushedVariantExtractions) {
+ override protected def partitions: Seq[FilePartition] = {
+ heldLock += Thread.holdsLock(this)
+ super.partitions
+ }
+ }
+ firstCall(probe)
+ assert(heldLock.toSeq === Seq(false))
+ }
+ }
+ }
+ }
+
+ test("scalar subquery on partition column of v2 parquet prunes partitions") {
+ // Scalar subquery (SELECT max(dim_id) FROM dim WHERE dim_val = 7)
evaluates to 7.
+ // WHERE f.part = 7 should select exactly partition 7 -> 10 rows. Without
partition
+ // pruning, the scan reads all 100 rows and the post-filter keeps 10.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql(
+ """SELECT f.id FROM fact f
+ |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val =
7)""".stripMargin)
+ val rows = df.collect()
+ assert(rows.length == 10,
+ s"expected 10 rows after scalar-subquery filter, got ${rows.length}")
+ val numOutputRows = factScanOf(df).metrics("numOutputRows").value
+ assert(numOutputRows == 10,
+ s"expected fact scan to read 10 rows after scalar-subquery pruning,
got " +
+ s"$numOutputRows. plan:\n" +
df.queryExecution.executedPlan.treeString)
+ }
+ }
+ }
+
+ test("filterAttributes reports the read partition columns") {
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT f.id, f.part FROM fact f")
+ df.collect()
+ val declared =
fileScanOf(df).filterAttributes().flatMap(_.fieldNames).toSeq
+ assert(declared == Seq("part"),
+ s"expected the partition column to be declared runtime-filterable,
got $declared")
+ }
+ }
+ }
+
+ test("filterAttributes reports nothing once a pushed aggregate drops the
partition columns") {
+ // With an aggregate pushed down, ParquetScan.readSchema() returns only
the aggregate schema, so
+ // the partition columns are not in the scan relation output. Spark
resolves every declared
+ // runtime-filter attribute against that output and fails the query if one
is missing, so
+ // filterAttributes() must not report them.
+ withDppV2Conf {
+ withSQLConf(SQLConf.PARQUET_AGGREGATE_PUSHDOWN_ENABLED.key -> "true") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT count(part) FROM fact")
+ checkAnswer(df, Row(100))
+ val declared =
fileScanOf(df).filterAttributes().flatMap(_.fieldNames).toSeq
+ assert(declared.isEmpty,
+ s"expected no runtime-filterable attribute under a pushed
aggregate, got $declared")
+ }
+ }
+ }
+ }
+
+ test("filterAttributes keeps a partition column a pushed aggregate groups
by") {
+ // Aggregate pushdown only happens when the GROUP BY set covers every
partition column, and
+ // those columns go into the pushed-down schema, so they stay in the scan
relation output and
+ // stay runtime-filterable. This is the other side of the intersection
from the test above.
+ withDppV2Conf {
+ withSQLConf(SQLConf.PARQUET_AGGREGATE_PUSHDOWN_ENABLED.key -> "true") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT part, count(*) FROM fact GROUP BY part")
+ checkAnswer(df, (0 until 10).map(p => Row(p, 10)))
+ val scan = fileScanOf(df)
+ // Without the pushdown the assertion below says nothing: `part` is
in the read partition
+ // schema either way. Pin that the aggregate really went into the
scan.
+ assert(scan.asInstanceOf[ParquetScan].pushedAggregate.isDefined,
+ s"expected the aggregate to be pushed into the scan, got
${scan.description()}")
+ val declared = scan.filterAttributes().flatMap(_.fieldNames).toSeq
+ assert(declared == Seq("part"),
+ s"expected the group-by partition column to stay declared, got
$declared")
+ }
+ }
+ }
+ }
+
+ test("scalar subquery on a partition column is not re-evaluated after the
scan") {
+ // The file scan declares its partition columns in
fullyPushedFilterAttributes: applying such a
+ // filter selects partition directories, so every row the scan returns
satisfies it and Spark
+ // drops the post-scan FilterExec, the same treatment a compile-time
partition filter gets.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql(
+ """SELECT f.id FROM fact f
+ |WHERE f.part = (SELECT max(dim_id) FROM dim WHERE dim_val =
7)""".stripMargin)
+ val rows = df.collect()
+ assert(rows.length == 10, s"expected 10 rows, got ${rows.length}")
+ val executedPlan = df.queryExecution.executedPlan
+ assert(collect(executedPlan) { case f: FilterExec => f }.isEmpty,
+ "a fully pushed partition filter should leave no FilterExec above
the scan, got:\n" +
+ executedPlan.treeString)
+ }
+ }
+ }
+
+ test("a runtime filter is ANDed with the compile-time partition filters, not
substituted") {
+ // The two halves are applied in different places: `partitions` puts
`partitionFilters` through
+ // `fileIndex.listFiles`, and `planInputPartitionsWithRuntimeFilters`
filters what that planned.
+ // Both are the scan's only evaluator: FileScanBuilder.pushFilters returns
just the data filters
+ // as post-scan filters, and a partition-column runtime filter is declared
fully pushed, so
+ // dropping either half returns rows that nothing filters out.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ // min(dim_id) over dim_val > 6 is 7, so this keeps partitions 7 and
8: 7 from the runtime
+ // filter's lower bound, 8 from the compile-time upper bound.
+ val df = sql(
+ """SELECT f.id FROM fact f
+ |WHERE f.part <= 8
+ | AND f.part >= (SELECT min(dim_id) FROM dim WHERE dim_val >
6)""".stripMargin)
+ checkAnswer(df, (0 until 10).flatMap(i => Seq(Row(i * 10 + 7), Row(i *
10 + 8))))
+ val numOutputRows = factScanOf(df).metrics("numOutputRows").value
+ assert(numOutputRows == 20,
+ s"expected the fact scan to read only partitions 7 and 8, got
$numOutputRows rows")
+ }
+ }
+ }
+
+ test("runtime partition pruning survives a mixed-case partition column
reference") {
+ // `f.PART` is resolved to the schema's `part` before it ever becomes a
runtime filter, so this
+ // pins the end-to-end path rather than any name matching: what the scan
receives is identical
+ // to the lower-case query's. The test below pins how the scan matches the
name.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql(
+ """SELECT f.id FROM fact f
+ |WHERE f.PART = (SELECT max(dim_id) FROM dim WHERE dim_val =
7)""".stripMargin)
+ val rows = df.collect()
+ assert(rows.length == 10, s"expected 10 rows, got ${rows.length}")
+ val numOutputRows = factScanOf(df).metrics("numOutputRows").value
+ assert(numOutputRows == 10,
+ s"expected the fact scan to read 10 rows, got $numOutputRows")
+ }
+ }
+ }
+
+ test("a runtime filter the scan cannot apply is rejected rather than
silently dropped") {
+ // Unreachable through SQL: PushDownUtils screens a filter against the
attributes the scan
+ // declared before handing it over. Calling the scan directly is what pins
the guard, which has
+ // to fail rather than drop the filter, since a fully pushed predicate
dropped here would be
+ // evaluated nowhere. The reference is matched the way the rest of the
scan matches partition
+ // column names, so a mixed-case reference binds rather than being
rejected.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql("SELECT f.id, f.part FROM fact f")
+ df.collect()
+ val scan = fileScanOf(df)
+ val e = intercept[SparkException] {
+ scan.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(AttributeReference("id", LongType)(), Literal(1L))))
+ }
+ assert(e.getCondition === "INTERNAL_ERROR")
+ assert(e.getMessage.contains("can only apply a runtime filter over its
read partition"),
+ s"unexpected message: ${e.getMessage}")
+
+ val mixedCase = scan.planInputPartitionsWithRuntimeFilters(
+ Array(EqualTo(AttributeReference("PART", IntegerType)(),
Literal(7))))
+ assert(mixedCase.flatMap(
+
_.asInstanceOf[FilePartition].files.map(_.partitionValues.getInt(0))).distinct.toSeq
+ === Seq(7))
+ }
+ }
+ }
+
+ test("a fully pushed filter stays exact when one FilePartition spans
partition directories") {
+ // The fully-pushed contract used to ask for the attribute's value to be
fixed within every
+ // InputPartition. A file scan does not satisfy that: FilePartition packs
by size and can hold
+ // files from several partition directories. What makes the declaration
sound is that the
+ // predicate chose those directories, so every row in the partition
satisfies it even though the
+ // value varies. Pack the surviving directories into one FilePartition to
exercise that.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1g",
+ SQLConf.FILES_MIN_PARTITION_NUM.key -> "1",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "0") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ // min(dim_id) over dim_val > 7 is 8, so this keeps partitions 8 and
9.
+ val df = sql(
+ """SELECT f.id FROM fact f
+ |WHERE f.part >= (SELECT min(dim_id) FROM dim WHERE dim_val >
7)""".stripMargin)
+ checkAnswer(df, (0 until 10).flatMap(i => Seq(Row(i * 10 + 8), Row(i
* 10 + 9))))
+ val factScan = factScanOf(df)
+ assert(factScan.filteredPartitions.flatten.size == 1,
+ "expected both surviving partition directories to pack into one
FilePartition, got " +
+ s"${factScan.filteredPartitions.flatten.size}")
+ val numOutputRows = factScan.metrics("numOutputRows").value
+ assert(numOutputRows == 20,
+ s"expected that one FilePartition to carry both directories' 20
rows, got " +
+ s"$numOutputRows")
+ assert(collect(df.queryExecution.executedPlan) { case f: FilterExec
=> f }.isEmpty,
+ "expected the partition filter to be fully pushed, leaving no
FilterExec")
+ }
+ }
+ }
+ }
+
+ test("runtime pruning packs the surviving files for their own size") {
+ // `partitions` sized its bins for the whole file set, so handing those
bins back with files
+ // removed leaves a heavily pruned scan running the few tasks the unpruned
set needed. Selecting
+ // the same partitions at compile time packs for the pruned set, so the
two layouts must agree.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1g",
+ SQLConf.FILES_MIN_PARTITION_NUM.key -> "3",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "0") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val expected = (0 until 10).flatMap(i => Seq(Row(i * 10 + 8, 8),
Row(i * 10 + 9, 9)))
+ // min(dim_id) over dim_val > 7 is 8, so this keeps partitions 8 and
9.
+ val runtime = sql(
+ """SELECT f.id, f.part FROM fact f
+ |WHERE f.part >= (SELECT min(dim_id) FROM dim WHERE dim_val >
7)""".stripMargin)
+ checkAnswer(runtime, expected)
+ val compileTime = sql("SELECT f.id, f.part FROM fact f WHERE f.part
>= 8")
+ checkAnswer(compileTime, expected)
+ // What gives the comparison teeth: the whole file set packs two or
more of the surviving
+ // files into one bin, so handing those bins back with the other
files removed cannot
+ // produce as many partitions as packing the survivors does. That
needs the fixture to
+ // write more than one file per partition directory, and this is the
assertion that fails
+ // if it stops doing so.
+ val unpruned = sql("SELECT f.id, f.part FROM fact f")
+ unpruned.collect()
+ val unprunedPartitions =
factScanOf(unpruned).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val binsHoldingSurvivors =
+
unprunedPartitions.count(_.files.exists(_.partitionValues.getInt(0) >= 8))
+ val survivorsWhenUnpruned =
+
unprunedPartitions.flatMap(_.files).count(_.partitionValues.getInt(0) >= 8)
+ assert(binsHoldingSurvivors < survivorsWhenUnpruned,
+ s"expected the whole set to pack $survivorsWhenUnpruned surviving
files into fewer " +
+ s"than $survivorsWhenUnpruned bins, got $binsHoldingSurvivors")
+ val compileTimePartitions =
factScanOf(compileTime).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val runtimePartitions =
factScanOf(runtime).filteredPartitions.flatten
+ .map(_.asInstanceOf[FilePartition])
+ val runtimeFiles = runtimePartitions.flatMap(_.files).size
+ // The teeth above were measured on the unpruned scan; this is what
ties them to this arm.
+ assert(runtimeFiles === survivorsWhenUnpruned,
+ s"expected this arm to read the $survivorsWhenUnpruned files the
unpruned packing " +
+ s"held, got $runtimeFiles")
+ // Both arms must see the same files: the runtime path cannot
re-split, so a compile-time
+ // maxSplitBytes below one file's length would cut that arm's files
finer and the counts
+ // would differ for a reason this test is not about.
+ assert(runtimeFiles === compileTimePartitions.flatMap(_.files).size,
+ s"expected both arms to read $runtimeFiles files, got " +
+ s"${compileTimePartitions.flatMap(_.files).size} at compile
time")
+ // The two arms pack from differently ordered input: `partitions`
sorts within a partition
+ // directory, this path sorts the survivors globally, and Next Fit
Decreasing is sensitive
+ // to that. What makes the counts comparable is that under this conf
no two surviving
+ // files fit in one bin, so both arms give one partition per file
whatever the order.
+ assert(runtimePartitions.size === runtimeFiles,
+ s"expected one partition per surviving file, got
${runtimePartitions.size} " +
+ s"for $runtimeFiles files")
+ assert(runtimePartitions.size === compileTimePartitions.size,
+ "expected the runtime-pruned scan to be packed like the
compile-time one " +
+ s"(${compileTimePartitions.size} partitions), got
${runtimePartitions.size}")
+ }
+ }
+ }
+ }
+
+ test("a DPP filter degraded to true leaves the v2 file scan unpruned") {
+ // With reuse-only DPP and no broadcast to reuse, the subquery is replaced
by true before
+ // execution. What this pins is that such a filter still reaches the scan
and the query stays
+ // correct; the row count only records that nothing was pruned, which is
what a true predicate
+ // does whether or not `catalystRuntimeFilters` drops it.
+ withDppV2Conf {
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ val df = sql(
+ """SELECT f.id FROM fact f JOIN dim d
+ |ON f.part = d.dim_id WHERE d.dim_val = 7""".stripMargin)
+ val rows = df.collect()
+ assert(rows.length == 10, s"expected 10 rows after join+filter, got
${rows.length}")
+ val factScan = factScanOf(df)
+ assert(factScan.runtimeFilters.exists {
+ case DynamicPruningExpression(l: Literal) => l ==
Literal.TrueLiteral
+ case _ => false
+ }, s"expected a degraded DPP filter, got ${factScan.runtimeFilters}")
+ val numOutputRows = factScan.metrics("numOutputRows").value
+ assert(numOutputRows == 100,
+ s"expected the fact scan to read all 100 rows, got $numOutputRows")
+ }
+ }
+ }
+ }
+
+ test("a runtime filter binds to the read partition columns when only some of
them are read") {
+ // `partitions` projects each file's partition values down to
`readPartitionSchema` when the
+ // query reads a strict subset of the partition columns, so a filter on
`p2` has to read slot 0
+ // of those values. Binding by the table's partition schema would read
slot 1, which in the
+ // projected values holds `p3`: rows that satisfy the filter are lost, and
with the filter
+ // fully pushed, rows that do not are returned.
+ withDppV2Conf {
+ withTempDir { dir =>
+ val tPath = new File(dir, "t").getCanonicalPath
+ spark.range(100).selectExpr("id", "id % 2 AS p1", "id % 5 AS p2", "id
% 3 AS p3")
+ .write.format("parquet").partitionBy("p1", "p2", "p3").save(tPath)
+ spark.read.format("parquet").schema("id long, p1 int, p2 int, p3
int").load(tPath)
+ .createOrReplaceTempView("t")
+ writeFactAndDim(dir)
+ val expected = (0 until 100).filter(_ % 5 == 2).map(i => Row(i.toLong,
i % 3))
+ Seq(
+ """SELECT id, p3 FROM t
+ |WHERE p2 = (SELECT max(dim_id) FROM dim WHERE dim_val =
2)""".stripMargin,
+ """SELECT t.id, t.p3 FROM t JOIN dim d
+ |ON t.p2 = d.dim_id WHERE d.dim_val = 2""".stripMargin).foreach {
query =>
+ val df = sql(query)
+ checkAnswer(df, expected)
+ val scan = factScanOf(df)
+
assert(scan.scan.asInstanceOf[FileScan].readPartitionSchema.fieldNames.toSeq ===
+ Seq("p2", "p3"), "the scan must read a strict subset of the
partition columns")
+ val numOutputRows = scan.metrics("numOutputRows").value
+ assert(numOutputRows == expected.size,
+ s"expected the scan to read ${expected.size} rows, got
$numOutputRows")
+ }
+ }
+ }
+ }
+
+ test("SPARK-38148: no DPP over a partition column a static equality already
fixes") {
+ // `FileScanBuilder` pushes the static partition filter fully into the
scan, so no `Filter`
+ // above the scan carries it; the cleanup has to read it from the scan
relation's pushed
+ // filters, as it reads it from the `Filter` over a V1 relation. V1 is the
reference.
+ withDppV2Conf {
+ withTempDir { dir =>
+ writeFactAndDim(dir)
+ // An int join key, so the pruning key is `part` itself rather than a
cast of it.
+ val dimIntPath = new File(dir, "dim_int").getCanonicalPath
+ spark.range(10).selectExpr("CAST(id AS INT) AS dim_id", "CAST(id AS
INT) AS dim_val")
+ .write.format("parquet").save(dimIntPath)
+ // (condition, the partitions it keeps, whether DPP stays, whether V1
is checked). The first
+ // six are the SPARK-38148 cases. The last two do not fix the key, so
V2 keeps the DPP; they
+ // are checked under V2 only, because V1 reads an equality anywhere in
its filter and drops
+ // the DPP.
+ Seq[(String, Int => Boolean, Boolean, Boolean)](
+ ("f.part = 3", _ == 3, false, true),
+ ("3 = f.part", _ == 3, false, true),
+ ("f.part <=> 3", _ == 3, false, true),
+ ("3 <=> f.part", _ == 3, false, true),
+ ("f.part > 3", _ > 3, true, true),
+ ("5 > f.part", _ < 5, true, true),
+ ("f.part <> 3", _ != 3, true, false),
+ ("(f.part = 3 OR f.part = 4)", p => p == 3 || p == 4, true, false)
+ ).foreach { case (condition, keeps, hasDpp, checkV1) =>
+ val query =
+ s"""SELECT /*+ broadcast(d) */ f.id FROM fact f JOIN dim_int d
+ |ON f.part = d.dim_id WHERE $condition AND d.dim_val >
2""".stripMargin
+ val expected = (0 until 100).filter { i =>
+ val part = i % 10
+ part > 2 && keeps(part)
+ }.map(i => Row(i.toLong))
+ (if (checkV1) Seq("parquet", "") else Seq("")).foreach { v1Sources =>
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> v1Sources) {
+ // A temp view keeps the relation it was created with, so create
both under this conf.
+ spark.read.format("parquet").schema("id long, part int")
+ .load(new File(dir,
"fact").getCanonicalPath).createOrReplaceTempView("fact")
+
spark.read.format("parquet").load(dimIntPath).createOrReplaceTempView("dim_int")
+ val df = sql(query)
+ val optimized = df.queryExecution.optimizedPlan
+
assert(optimized.exists(_.isInstanceOf[DataSourceV2ScanRelation]) ===
+ v1Sources.isEmpty, s"wrong source version for
'$v1Sources':\n$optimized")
+ assert(collectDppFilters(optimized).nonEmpty === hasDpp,
+ s"`$condition` with v1 sources '$v1Sources': expected DPP
$hasDpp, got\n$optimized")
+ checkAnswer(df, expected)
+ }
+ }
+ }
+ }
+ }
+ }
+
+ Seq("orc", "json", "csv").foreach { format =>
Review Comment:
**7. [Nit] Nothing runs a runtime filter on Avro or text**
This loop covers parquet, orc, json and csv, but `AvroScan` and `TextScan`
extend `RuntimePartitionPruningFileScan` too, and no test gives either of them
a DPP or scalar-subquery filter. Text needs its own fixture, since the text
writer takes a single string column, and Avro would go under `connector/avro`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [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]
