You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "peter-toth (via GitHub)" <gi...@apache.org> on 2026/10/22 17:28:36 UTC
[PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
peter-toth opened a new pull request, #59715:
URL: https://github.com/apache/spark/pull/59715
### What changes were proposed in this pull request?
A follow-up to SPARK-59620, which lets the V1 Parquet reader apply a join's runtime bloom filter while it reads. This PR gives data source V2 scans the same. A connector such as Iceberg can then read the bloom's key columns first, and skip the other columns of the rows the bloom rejects.
The API, all `@Evolving`:
- `SupportsStorageFilters`, a `Scan` mix-in. `storageFilterAttributes()` names the columns the scan's readers can evaluate a storage filter on.
- `StorageFilter`, an opaque predicate that Spark implements. It has `references()`, `eval(InternalRow)` and `int select(ColumnarBatch, int[] selected)`, which writes the positions of the rows it keeps. The default `select` loops over `eval`.
- `StorageFilterNotAppliedException`, error condition `STORAGE_FILTER_NOT_APPLIED`. `eval` and `select` throw it when the filter cannot be applied to a row.
- `StorageFilterRequest`, with `filter()` and `appendsCheckedColumn()`.
- `PartitionReaderFactory.createReader(partition, request)` and `createColumnarReader(partition, request)`. The factory of a scan that declares storage-filter columns must implement the one for its read mode.
- `Scan.toBatch` now says it may be called more than once for a scan offered a storage filter, as for runtime filtering.
Planning and execution:
- `DataSourceV2Strategy` offers a top-level runtime bloom over declared columns, under the new `spark.sql.sources.v2.storageFilterPushdown.enabled`. The per-conjunct rules are shared with `FileSourceStrategy` through the new `StorageFilterPlanning`.
- The bloom stays in the post-scan `Filter`. Where whole-stage codegen is on and the scan's output still fits it with the column, the scan gets a boolean checked column last, and the `Filter` evaluates the bloom as `IF(checked, true, bloom)`. A `Project` drops the column.
- `BatchScanExec` ANDs the storage filters into one `StorageFilter`, and `DataSourceRDD` hands each task's readers a `StorageFilterRequest`. Each task deserializes its own filter.
The design decisions:
- **Applying the filter is optional per unit,** such as a file or a row group. Spark keeps the predicate above the scan. A reader that cannot use it on a unit reads that unit as a plain read does, and marks its rows unchecked. When the filter throws `StorageFilterNotAppliedException`, the reader reads the rest of the unit that way. Spark's filter throws it for an error evaluating it, but not for a task kill or a fatal error.
- **The checked column says which rows the reader filtered.** It is requested in the task, through the request, since the scan's read schema is fixed before `InjectRuntimeFilter` adds the bloom. It is requested only where dropping it copies no extra row. A columnar scan out of whole-stage codegen would copy every row in the `Project` that drops it. A row-based V2 scan has that `Project` anyway, for the UnsafeRow conversion.
- **The filter is delivered in the task, not at planning.** The reader factory does not carry it. The same API can later hand each partition its own bloom, for storage-partitioned joins.
- **One combined filter per scan, opaque to the connector.** Spark implements it, so a connector does not reproduce Spark's bloom format or hash. Only runtime blooms are offered for now. Offering more is a planner change.
- **`select` returns positions, not a bitmap.** Iceberg's Spark reader already wraps a filtered batch as `ColumnVectorWithFilter(delegate, int[] rowIdMapping)`.
`BatchScanExec` gets a `storageFilters` field, so the positional patterns over it in eight test files gain one `_`.
### Why are the changes needed?
SPARK-59620 saves IO only in Spark's own Parquet reader. A table read through a V2 connector, such as Iceberg, still reads every projected column of the rows a runtime bloom then drops. With this API a connector can do what the V1 reader does. The Iceberg side is a follow-up in that project.
### Does this PR introduce _any_ user-facing change?
Yes:
- a new public connector API, listed above;
- with `spark.sql.sources.v2.storageFilterPushdown.enabled` on, which is off by default:
- a V2 scan that implements it lists the offered filters under `StorageFilters` in its description and in `EXPLAIN FORMATTED`;
- its output gains the checked column, which a `Project` drops;
- the post-scan `Filter` skips the bloom on the rows the reader marks checked.
Nothing changes for a scan that does not implement `SupportsStorageFilters`.
### How was this patch tested?
New `StorageFilterPushDownSuite`, over a V2 test source whose reader applies the filter row by row or through `select`:
- a row-based and a columnar scan, with AQE on and off, each get the bloom and the checked column;
- the `Filter` skips the bloom on the rows a reader marks checked, even rows it did not filter;
- a reader that gives the filter up returns its rows unchecked, and the `Filter` applies the bloom to them;
- a filter that throws `StorageFilterNotAppliedException` is given up, and its rows come back unchecked;
- a scan too wide for whole-stage codegen with the column, or without whole-stage codegen, gets the bloom without it;
- nothing is offered with the conf off, to a scan that declares no column or only another one, or to a scan with a column named like the checked column;
- a non-deterministic conjunct is not offered;
- the filters show in the scan description and in `EXPLAIN FORMATTED`, and keep a plain scan's canonicalized plan apart, while a twin with fresh exprIds keeps the same result.
New `StorageFilterImplSuite`: the positions `select` writes, the not-applied answer, that a task kill propagates, and the internal error for a column outside the scan output. `ParquetStorageFilterSuite` passes over the shared planning rules.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code with Claude Opus 5.5
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4372009016
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -187,10 +191,16 @@ class DataSourceRDD(
castPartition(split).inputPartition.iterator.flatMap { inputPartition =>
val iter = if (columnarReads) {
- val batchReader = partitionReaderFactory.createColumnarReader(inputPartition)
+ val batchReader = storageFilterRequest match {
+ case Some(request) => partitionReaderFactory.createColumnarReader(inputPartition, request)
+ case None => partitionReaderFactory.createColumnarReader(inputPartition)
+ }
new PartitionIterator[ColumnarBatch](batchReader, _.numRows, taskState)
} else {
- val rowReader = partitionReaderFactory.createReader(inputPartition)
+ val rowReader = storageFilterRequest match {
+ case Some(request) => partitionReaderFactory.createReader(inputPartition, request)
Review Comment:
**2. [Low] Nothing checks that a reader appends the checked column it was asked for**
When `appendsCheckedColumn()` is true, the scan's `output` ends with the checked column. But nothing checks that the reader returned here actually appends it. An easy mistake is a stub that takes the request and delegates to `createReader(partition)`. Then:
- For an `UnsafeRow` of `n` fields, the Filter's generated code reads ordinal `n`. `assertIndexIsValid` is a Java `assert`, which is off in production, so `isNullAt(n)` and `getBoolean(n)` read the bytes after the fixed-length region, and the bloom is skipped on arbitrary rows.
- A `GenericInternalRow`, or `batch.column(n)` in `ColumnarToRowExec`'s generated code, throws an `ArrayIndexOutOfBoundsException` that names neither the connector nor the request.
Today the join still filters, so only the contract breaks, but the failure is hard to trace back. Could the first row or batch of each such reader be checked against `output.length`, the way `StorageFilterImpl.eval` checks its own row width, and fail with an error that names the factory class?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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.internal.connector
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.{SparkEnv, SparkException, TaskContext}
+import org.apache.spark.executor.Executor
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.internal.config.KILL_ON_FATAL_ERROR_DEPTH
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to
+ * it. The binding is built on first use, so an instance created on the driver can be evaluated on
+ * an executor.
+ */
+class StorageFilterImpl private (
+ private val catalystExpr: CatalystExpression,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes.
+ private val numReferences = attrs.length
+
+ // Whether this instance has logged a failure, so a task logs only its first.
+ @transient private var loggedFailure = false
+
+ @transient private lazy val boundPredicate: InternalRow => Boolean = {
+ val bound = BindReferences.bindReference(catalystExpr, attrs)
+ CatalystPredicate.create(bound).eval
+ }
+
+ /**
+ * Wraps a non-fatal error in [[StorageFilterNotAppliedException]], unless the task is being
+ * killed or a fatal error is in its cause chain. Each task logs the first one. A row of another
+ * width than [[references]] is a reader bug, which fails the task, since the predicate would read
+ * the wrong values.
+ */
+ @throws[StorageFilterNotAppliedException]
+ override def eval(row: InternalRow): Boolean = {
+ // Outside the `try`, so a failure to build the predicate, or a row of another width, fails the
+ // task.
+ if (row.numFields != numReferences) {
+ throw SparkException.internalError(
+ s"A storage filter over ${references.mkString("[", ", ", "]")} takes rows of " +
+ s"$numReferences values, but got ${row.numFields}")
+ }
+ val predicate = boundPredicate
+ try {
+ predicate(row)
+ } catch {
+ case NonFatal(e) if !StorageFilterImpl.taskKilled && !StorageFilterImpl.hasFatalCause(e) =>
Review Comment:
**1. [Medium] `eval` does not follow the V1 error rules it says it mirrors**
The comment at L120-121 and the PR description say a fatal error and a task kill go out "as from the V1 reader", but this `catch` only wraps a `NonFatal` error and lets everything else out as it was thrown. `ParquetStorageFilter.rethrowIfMustPropagate(e, fromRead = false)` differs in four cases:
| Error from the predicate | V1 reader | `StorageFilterImpl.eval` |
|---|---|---|
| `NonFatal` wrapper around a fatal error | throws the fatal error itself | throws the wrapper |
| bare `InternalError` | throws it in an `UnknownError` | throws it as is |
| `NonFatal` error while a kill is pending | throws a checked `SparkException` around it | throws it as is |
| `LinkageError`, `SparkOutOfMemoryError`, `InterruptedException` without a kill | gives the filter up | fails the task |
V1 converts the first three on purpose (ParquetStorageFilter.scala:228-230 and 276-280). Spark's own `FilePartitionReader` takes any `RuntimeException | IOException | InternalError` from a reader for a corrupt file under `spark.sql.files.ignoreCorruptFiles=true` (FilePartitionReader.scala:74-77, DataSourceUtils.scala:273). So take a `FileScan`-based source that adopts this API, with a per-file reader that catches only `StorageFilterNotAppliedException` as the Javadoc asks. If the bloom raises `IllegalStateException(cause = OutOfMemoryError)` or an `InternalError`, that reader logs "Skipped the rest of the content in the corrupted file" and silently drops the rest of the file, and the executor keeps running after the OOM. The same evaluation in the post-scan `Filter` fails the task.
No built-in scan implements the mix-in yet, so this is latent. But the two suites already pin opposite outcomes for the same input: StorageFilterImplSuite.scala:112-125 asserts `e eq error`, and a raw `SparkArithmeticException` under a kill, while ParquetStorageFilterSuite.scala:645-659 asserts a checked `SparkException` and the wrapped `OutOfMemoryError` itself.
Could `eval` catch `Throwable` and apply the `fromRead = false` half of `rethrowIfMustPropagate` (with `throwIfKilled` and `fatalErrorDepth`) before it wraps the error, so that both paths share one rule? If the narrower rule is intended, the comment at L120-121, the PR description and the `@throws` text of `StorageFilter.eval` should say so.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/filter/StorageFilterNotAppliedException.java:
##########
@@ -0,0 +1,52 @@
+/*
+ * 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.expressions.filter;
+
+import java.util.Map;
+
+import org.apache.spark.SparkThrowable;
+import org.apache.spark.SparkThrowableHelper;
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * Thrown by a {@link StorageFilter} that cannot be applied to a row, for instance because
+ * evaluating it raised an error. The reader catches it and reads the rest of the unit it was
+ * filtering, such as a file or a row group, without the filter. The cause is the error the filter
+ * met.
+ * <p>
+ * It is checked, so a reader handles it where it can fall back to reading without the filter. A
Review Comment:
**5. [Nit] "It is checked" does not help a Scala reader**
Scala does not enforce checked exceptions. So a Scala `PartitionReader` that calls `filter.eval` or `select` without a `catch` compiles without a warning. On the first evaluation error it then fails the task with `STORAGE_FILTER_NOT_APPLIED`, instead of reading the unit without the filter. The test readers have to remember the `catch` by hand (StorageFilterTestSource.scala:222 and :275). Maybe add a sentence here and in `docs/sql-v2-data-sources.md` that a reader written in Scala must catch it itself.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/PartitionReaderFactory.java:
##########
@@ -56,10 +56,40 @@ default PartitionReader<ColumnarBatch> createColumnarReader(InputPartition parti
"DATA_SOURCE_COLUMNAR_READER_NOT_SUPPORTED", Map.of("class", getClass().getName()));
}
+ /**
+ * Returns a row-based partition reader that may apply the storage filter of {@code request} while
+ * it reads the given {@link InputPartition}. Spark calls this instead of
+ * {@link #createReader(InputPartition)} for a {@link SupportsStorageFilters} scan it offered a
+ * filter, so the factory of such a scan must implement it for row-based reads.
+ *
+ * @since 5.0.0
+ */
+ default PartitionReader<InternalRow> createReader(
Review Comment:
**3. [Nit] A factory without the new overload is found only inside the tasks**
Spark calls this overload because the `Scan` declares `storageFilterAttributes()`, while the factory comes from the `Batch`. Take a factory that does not override the overload for the scan's read mode, e.g. a delegating or metrics-wrapping one. It fails here in every task, after the bloom's build side has run, and only after `spark.task.maxFailures` retries. Small test tables rarely trigger `InjectRuntimeFilter`, so this can show up first in production.
`BatchScanExec` knows `readerFactory` and `supportsColumnar` on the driver. So `inputRDD` could check once, e.g. through `getDeclaringClass` of the overload, and fail with the same error before any task starts, while EXPLAIN keeps working. The existing `createColumnarReader(partition)` default behaves the same way, so this is only a suggestion.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -200,10 +223,20 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
// rewritten it to TrueLiteral by the time we get here.
val runtimeFilters = dynamicFilters ++ scalarSubqueryFilters
- val batchExec = BatchScanExec(relation.output, relation.scan, runtimeFilters,
- relation.ordering, relation.relation.table, relation.keyGroupedPartitioning)
- DataSourceV2Strategy.withProjectAndFilter(
- project, postScanFilters.diff(fullyPushedRuntimeFilters),
+ // Whether the scan gets the checked column does not depend on whether it is columnar. Asking
+ // that can plan the scan's partitions, which the node built afterwards would plan again. A
+ // row-based scan gets a Project for the UnsafeRow conversion anyway, so dropping the column
+ // there copies nothing more.
+ val filtersAbove = postScanFilters.diff(fullyPushedRuntimeFilters)
+ val storageFilters = storageFiltersFor(relation, filtersAbove)
+ val checked = StorageFilterPlanning.checkedColumn(
+ session.sessionState.conf, storageFilters, relation.output)
+ val batchExec = BatchScanExec(relation.output ++ checked, relation.scan, runtimeFilters,
Review Comment:
**7. [Low] The subquery rules now re-plan every storage-filtered scan**
With the bloom's `ScalarSubquery` in `storageFilters`, the node built here no longer survives planning:
- `PlanSubqueries` or `PlanAdaptiveSubqueries` rewrites it.
- Under AQE, `ReuseAdaptiveSubquery` rewrites it again at stage creation, after `EnsureRequirements` (AdaptiveSparkPlanExec.scala:195, ReuseAdaptiveSubquery.scala:38).
Each copy computes `batch`, `inputPartitions` and `readerFactory` again. `fastEquals` forces `batch` through `equals`, and `RemoveRedundantProjects`, `ApplyColumnarRulesAndInsertTransitions` or `EnsureRequirements` force the partitions. With AQE on and the default `PARTITION_DEFINED` mode, a scan offered a bloom runs `toBatch()` 4 times and `planInputPartitions()` / `createReaderFactory()` 3 times. With the conf off, each runs 2 times. Every AQE re-plan of such a scan pays the extra copy again. The comment at L226-229 avoids planning a probe node, but the node it keeps is replaced right away.
Two consequences:
- `Batch.planInputPartitions` still promises a single call (Batch.java:40), and only the doc of `Scan.toBatch` was updated.
- Under a storage-partitioned join, the `ReuseAdaptiveSubquery` copy lands under a `GroupPartitionsExec` that was planned over the previous node. `checkChildStillMatches` then compares the re-planned keys by value (GroupPartitionsExec.scala:127 and :352). A connector whose second planning returns other keys fails with INTERNAL_ERROR, and that doc names only `PlanAdaptiveDynamicPruningFilters` as a rule that rebuilds `BatchScanExec`.
The test source returns `this` from `toBatch` and a fixed partition, so no test can see the extra calls. Could the planned `batch`, partitions and factory carry over to copies of the same `scan`, e.g. through a holder passed in `otherCopyArgs`? If not, the two docs above could name the new rebuilders.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -200,10 +223,20 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
// rewritten it to TrueLiteral by the time we get here.
val runtimeFilters = dynamicFilters ++ scalarSubqueryFilters
- val batchExec = BatchScanExec(relation.output, relation.scan, runtimeFilters,
- relation.ordering, relation.relation.table, relation.keyGroupedPartitioning)
- DataSourceV2Strategy.withProjectAndFilter(
- project, postScanFilters.diff(fullyPushedRuntimeFilters),
+ // Whether the scan gets the checked column does not depend on whether it is columnar. Asking
+ // that can plan the scan's partitions, which the node built afterwards would plan again. A
+ // row-based scan gets a Project for the UnsafeRow conversion anyway, so dropping the column
+ // there copies nothing more.
+ val filtersAbove = postScanFilters.diff(fullyPushedRuntimeFilters)
+ val storageFilters = storageFiltersFor(relation, filtersAbove)
+ val checked = StorageFilterPlanning.checkedColumn(
+ session.sessionState.conf, storageFilters, relation.output)
+ val batchExec = BatchScanExec(relation.output ++ checked, relation.scan, runtimeFilters,
+ relation.ordering, relation.relation.table, relation.keyGroupedPartitioning,
+ storageFilters, checked)
+ // With the checked column the scan's output differs from `project`, so a Project drops it.
+ DataSourceV2Strategy.withProjectAndFilter(project,
+ StorageFilterPlanning.skipCheckedRows(filtersAbove, storageFilters, checked),
Review Comment:
**10. [Low] The bloom's subquery now occurs twice in the plan**
The same `ScalarSubquery` now sits in the scan's `storageFilters` and, through this call, in the post-scan `Filter`.
- Without AQE, `PlanSubqueries` calls `QueryExecution.prepareExecutedPlan` once per occurrence (subquery.scala:456). So the bloom's creation side is physically planned twice, including the partition planning of its own scans.
- With `spark.sql.execution.reuseSubquery=false`, nothing dedups the two `SubqueryExec`s, so the bloom aggregation runs twice. With AQE, its final stage runs twice.
Running it twice needs a non-default conf, so this may be acceptable. If so, could the PR description say so next to its note that the task binary carries the bloom twice?
##########
docs/sql-v2-data-sources.md:
##########
@@ -319,7 +319,10 @@ for data locality.
is serialized to executors and creates a
[`PartitionReader`](api/java/org/apache/spark/sql/connector/read/PartitionReader.html) for each
input partition. The reader iterates over rows (or `ColumnarBatch` instances for columnar sources)
-and is closed after consumption.
+and is closed after consumption. For a `SupportsStorageFilters` scan that Spark offered a storage
+filter, `createReader(partition, request)` or `createColumnarReader(partition, request)` creates the
+reader instead. The `StorageFilterRequest` carries the filter, and whether the reader appends a
+column that marks the rows it applied the filter to.
Review Comment:
**15. [Low] The guide says the column marks the rows the filter was applied to**
The API docs say the mark is true only on a row the filter was evaluated on and kept, and false on a row kept although the filter rejected it: StorageFilterRequest.java:40-42, `StorageFilter.eval` and `SupportsStorageFilters`. "the rows it applied the filter to" invites a reader that filters per page or per batch to mark every row it evaluated, including returned rows the bloom rejected. The post-scan `If(checked, true, bloom)` then skips the bloom on exactly those rows. Today only the join, which applies the bloom again, hides that.
Maybe: "... and whether the reader appends a boolean column that is true only on the rows the filter was evaluated on and kept."
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,314 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GreaterThan, If, Literal, NumericEvalContext, Rand}
+import org.apache.spark.sql.execution.{BaseLimitExec, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{DoubleType, LongType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
+ withSQLConf(bloomJoinConf ++ conf: _*) {
+ withTable("bf_build") {
+ spark.range(0, 30).selectExpr("id AS bk", "id AS bv")
+ .write.format("parquet").saveAsTable("bf_build")
+ body
+ }
+ }
+ }
+
+ // A join whose build side is small and carries a selective predicate, which is the shape
+ // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`.
+ private def bloomJoin(probe: DataFrame): DataFrame = {
+ val build = spark.table("bf_build").where("bv = 5")
+ probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*)
+ }
+
+ private def v2Probe(options: (String, String)*): DataFrame =
+ spark.read.format(sourceName).options(options.toMap).load()
+
+ private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME
+
+ private def scanOf(plan: SparkPlan): BatchScanExec =
+ collect(plan) { case b: BatchScanExec => b }.headOption
+ .getOrElse(fail(s"no BatchScanExec in plan:\n$plan"))
+
+ // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows
+ // the reader marks as checked.
+ private def bloomsAbove(plan: SparkPlan): (Int, Int) = {
+ val conditions = collect(plan) { case f: FilterExec => f.condition }
+ val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum
+ val skipped = conditions.map(_.collect {
+ case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain)
+ if a.name == checkedName => 1
+ }.sum).sum
+ (all - skipped, skipped)
+ }
+
+ // The rows the Filter on the scan returned, counted before the join.
+ private def rowsOutOfFilter(plan: SparkPlan): Long = {
+ val filters = collect(plan) {
+ case f: FilterExec if f.exists(_.isInstanceOf[BatchScanExec]) => f
+ }
+ assert(filters.size == 1, s"expected one Filter on the scan:\n$plan")
+ filters.head.metrics("numOutputRows").value
+ }
+
+ // Runs the bloom join over the v2 source read with `options`, and returns its rows and plan.
+ private def runBloomJoin(options: (String, String)*): (Set[(Long, Long)], SparkPlan) = {
+ StorageFilterTestSource.reset()
+ val df = bloomJoin(v2Probe(options: _*))
+ (df.collect().map(r => (r.getLong(0), r.getLong(1))).toSet, df.queryExecution.executedPlan)
+ }
+
+ // The build side keeps `bv = 5`, so the join returns that one row.
+ private val joinResult = Set((5L, 50L))
+
+ for (columnar <- Seq(false, true); aqe <- Seq(false, true)) {
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind scan gets the checked column, and the Filter skips the bloom on checked rows " +
+ s"(AQE $aqe)") {
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val (result, plan) = runBloomJoin("columnar" -> columnar.toString)
+ val scan = scanOf(plan)
+ assert(scan.supportsColumnar == columnar, s"the scan must read $kind:\n$plan")
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(scan.checkedColumn.isDefined && scan.checkedColumn == scan.output.lastOption,
+ s"the scan must get the checked column, last:\n$plan")
+ assert(bloomsAbove(plan) == (0, 1),
+ s"the Filter must skip the bloom on checked rows:\n$plan")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(),
+ s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " +
+ s"${StorageFilterTestSource.rowsReturned.get()}")
+ val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output }
+ assert(probeSide.exists(!_.exists(_.name == checkedName)),
+ s"the Project above the scan must drop the checked column; got $probeSide")
+ assert(StorageFilterTestSource.filteredReaders.get() == 1,
+ s"the one partition's reader must get the filter; " +
+ s"got ${StorageFilterTestSource.filteredReaders.get()}")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 1,
+ "the reader must be asked for the checked column")
+ assert(result == joinResult, s"got $result")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ s"the reader must evaluate every row; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(false, true).foreach { columnar =>
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind reader that gives the filter up marks its rows unchecked, and the Filter " +
+ "filters them") {
+ withBloomJoin() {
+ val (result, plan) = runBloomJoin("columnar" -> columnar.toString, "giveUp" -> "true")
+ assert(bloomsAbove(plan) == (0, 1), s"the plan is that of a filtering reader:\n$plan")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == 0,
+ s"the reader must evaluate nothing; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() == StorageFilterTestSource.numRows,
+ s"the reader must return every row; got ${StorageFilterTestSource.rowsReturned.get()}")
+ assert(rowsOutOfFilter(plan) < StorageFilterTestSource.numRows,
+ s"the Filter must apply the bloom to unchecked rows; got ${rowsOutOfFilter(plan)}")
+ assert(result == joinResult, s"got $result")
+ }
+ }
+
+ test(s"the Filter takes a $kind reader's word on the rows it marks checked") {
+ // This reader breaks the contract on purpose: it applies no filter but marks every row
+ // checked. Only skipping the bloom on those rows keeps all of them.
+ withBloomJoin() {
+ val (result, plan) =
+ runBloomJoin("columnar" -> columnar.toString, "markAllChecked" -> "true")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.numRows,
+ s"the Filter must keep every row marked checked; got ${rowsOutOfFilter(plan)}")
+ assert(result == joinResult, s"the join must still filter; got $result")
+ }
+ }
+ }
+
+ test("a filter that cannot be applied is given up, and its rows come back unchecked") {
Review Comment:
**14. [Nit] This test exercises only the fixture**
It builds the same ANSI `1 / k > 0` filter as `StorageFilterImplSuite.positiveInverse`, and drives `StorageFilterTestColumnarReader` by hand. StorageFilterImplSuite's "an error evaluating the filter answers that it is not applied" already covers the filter. So the only new code this test covers is the fixture's own `catch`, and it would pass even if `DataSourceRDD` never called the new overload. The row reader's identical `catch` (StorageFilterTestSource.scala:222-226) never runs. Yet the PR description counts this test as coverage of "a filter that throws `StorageFilterNotAppliedException` is given up".
Smaller things in this suite:
- Every caller of `runBloomJoin` asserts `result == joinResult` right after it.
- The "nothing offered" assertion pair is copied into two loops.
- `checkedName` (L72) uses the V1 alias `FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME` rather than `SupportsStorageFilters.CHECKED_COLUMN_NAME`.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,314 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GreaterThan, If, Literal, NumericEvalContext, Rand}
+import org.apache.spark.sql.execution.{BaseLimitExec, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{DoubleType, LongType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
+ withSQLConf(bloomJoinConf ++ conf: _*) {
+ withTable("bf_build") {
+ spark.range(0, 30).selectExpr("id AS bk", "id AS bv")
+ .write.format("parquet").saveAsTable("bf_build")
+ body
+ }
+ }
+ }
+
+ // A join whose build side is small and carries a selective predicate, which is the shape
+ // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`.
+ private def bloomJoin(probe: DataFrame): DataFrame = {
+ val build = spark.table("bf_build").where("bv = 5")
+ probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*)
+ }
+
+ private def v2Probe(options: (String, String)*): DataFrame =
+ spark.read.format(sourceName).options(options.toMap).load()
+
+ private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME
+
+ private def scanOf(plan: SparkPlan): BatchScanExec =
+ collect(plan) { case b: BatchScanExec => b }.headOption
+ .getOrElse(fail(s"no BatchScanExec in plan:\n$plan"))
+
+ // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows
+ // the reader marks as checked.
+ private def bloomsAbove(plan: SparkPlan): (Int, Int) = {
+ val conditions = collect(plan) { case f: FilterExec => f.condition }
+ val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum
+ val skipped = conditions.map(_.collect {
+ case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain)
+ if a.name == checkedName => 1
+ }.sum).sum
+ (all - skipped, skipped)
+ }
+
+ // The rows the Filter on the scan returned, counted before the join.
+ private def rowsOutOfFilter(plan: SparkPlan): Long = {
+ val filters = collect(plan) {
+ case f: FilterExec if f.exists(_.isInstanceOf[BatchScanExec]) => f
+ }
+ assert(filters.size == 1, s"expected one Filter on the scan:\n$plan")
+ filters.head.metrics("numOutputRows").value
+ }
+
+ // Runs the bloom join over the v2 source read with `options`, and returns its rows and plan.
+ private def runBloomJoin(options: (String, String)*): (Set[(Long, Long)], SparkPlan) = {
+ StorageFilterTestSource.reset()
+ val df = bloomJoin(v2Probe(options: _*))
+ (df.collect().map(r => (r.getLong(0), r.getLong(1))).toSet, df.queryExecution.executedPlan)
+ }
+
+ // The build side keeps `bv = 5`, so the join returns that one row.
+ private val joinResult = Set((5L, 50L))
+
+ for (columnar <- Seq(false, true); aqe <- Seq(false, true)) {
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind scan gets the checked column, and the Filter skips the bloom on checked rows " +
+ s"(AQE $aqe)") {
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val (result, plan) = runBloomJoin("columnar" -> columnar.toString)
+ val scan = scanOf(plan)
+ assert(scan.supportsColumnar == columnar, s"the scan must read $kind:\n$plan")
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(scan.checkedColumn.isDefined && scan.checkedColumn == scan.output.lastOption,
+ s"the scan must get the checked column, last:\n$plan")
+ assert(bloomsAbove(plan) == (0, 1),
+ s"the Filter must skip the bloom on checked rows:\n$plan")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(),
+ s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " +
+ s"${StorageFilterTestSource.rowsReturned.get()}")
+ val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output }
+ assert(probeSide.exists(!_.exists(_.name == checkedName)),
+ s"the Project above the scan must drop the checked column; got $probeSide")
+ assert(StorageFilterTestSource.filteredReaders.get() == 1,
+ s"the one partition's reader must get the filter; " +
+ s"got ${StorageFilterTestSource.filteredReaders.get()}")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 1,
+ "the reader must be asked for the checked column")
+ assert(result == joinResult, s"got $result")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ s"the reader must evaluate every row; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(false, true).foreach { columnar =>
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind reader that gives the filter up marks its rows unchecked, and the Filter " +
+ "filters them") {
+ withBloomJoin() {
+ val (result, plan) = runBloomJoin("columnar" -> columnar.toString, "giveUp" -> "true")
+ assert(bloomsAbove(plan) == (0, 1), s"the plan is that of a filtering reader:\n$plan")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == 0,
+ s"the reader must evaluate nothing; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() == StorageFilterTestSource.numRows,
+ s"the reader must return every row; got ${StorageFilterTestSource.rowsReturned.get()}")
+ assert(rowsOutOfFilter(plan) < StorageFilterTestSource.numRows,
+ s"the Filter must apply the bloom to unchecked rows; got ${rowsOutOfFilter(plan)}")
+ assert(result == joinResult, s"got $result")
+ }
+ }
+
+ test(s"the Filter takes a $kind reader's word on the rows it marks checked") {
+ // This reader breaks the contract on purpose: it applies no filter but marks every row
+ // checked. Only skipping the bloom on those rows keeps all of them.
+ withBloomJoin() {
+ val (result, plan) =
+ runBloomJoin("columnar" -> columnar.toString, "markAllChecked" -> "true")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.numRows,
+ s"the Filter must keep every row marked checked; got ${rowsOutOfFilter(plan)}")
+ assert(result == joinResult, s"the join must still filter; got $result")
+ }
+ }
+ }
+
+ test("a filter that cannot be applied is given up, and its rows come back unchecked") {
+ // `1 / k` under ANSI throws at k = 0, the partition's first row, so the reader gives the
+ // filter up before it keeps or drops anything.
+ val k = AttributeReference("k", LongType)()
+ val inverse = Divide(Literal(1.0), Cast(k, DoubleType), NumericEvalContext(EvalMode.ANSI))
+ val filter = StorageFilterImpl(GreaterThan(inverse, Literal(0.0)), Seq(k))
+ val reader = new StorageFilterTestColumnarReader(Some(filter), appendsChecked = true)
+ val rows = mutable.ArrayBuffer.empty[(Long, Boolean)]
+ try {
+ while (reader.next()) {
+ val batch = reader.get()
+ (0 until batch.numRows()).foreach { i =>
+ rows += ((batch.column(0).getLong(i), batch.column(2).getBoolean(i)))
+ }
+ }
+ } finally {
+ reader.close()
+ }
+ assert(rows == (0L until StorageFilterTestSource.numRows.toLong).map(_ -> false),
+ s"expected every row, unchecked; got ${rows.take(5)} and ${rows.size} rows in all")
+ }
+
+ // Under a limit of two fields, the checked column would make three.
+ Seq(
+ ("too wide for whole-stage codegen", SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> "2"),
+ ("without whole-stage codegen", SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false")
+ ).foreach { case (when, conf) =>
+ test(s"a scan $when gets the bloom without the checked column") {
+ withBloomJoin(conf) {
+ val (result, plan) = runBloomJoin("columnar" -> "true")
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(scan.checkedColumn.isEmpty, s"the scan must not get the checked column:\n$plan")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 0,
+ "the reader must not be asked for the checked column")
+ assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan")
+ assert(result == joinResult, s"got $result")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(
+ ("with the conf off", Seq(SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "false"), Nil),
+ ("to a scan that declares no column", Nil, Seq("declared" -> "")),
+ ("to a scan that declares only another column", Nil, Seq("declared" -> "v")),
+ ("to a scan that declares only a nested field", Nil, Seq("declared" -> "k.x")),
+ ("to a scan with a column named like the checked column", Nil,
+ Seq("valueName" -> checkedName.toUpperCase(java.util.Locale.ROOT)))
+ ).foreach { case (when, conf, options) =>
+ test(s"nothing is offered $when") {
+ withBloomJoin(conf: _*) {
+ val (result, plan) = runBloomJoin(options :+ ("columnar" -> "true"): _*)
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.isEmpty && scan.checkedColumn.isEmpty,
+ s"the scan must get no storage filter:\n$plan")
+ assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan")
+ assert(StorageFilterTestSource.filteredReaders.get() == 0,
+ "no reader may get a storage filter")
+ assert(result == joinResult, s"got $result")
+ }
+ }
+ }
+
+ test("the storage filter shows in the scan description and separates canonicalized plans") {
+ withBloomJoin() {
+ val plan = bloomJoin(v2Probe("columnar" -> "true")).queryExecution.executedPlan
+ val scan = scanOf(plan)
+ assert(scan.simpleString(100).contains("StorageFilters:"),
+ s"the scan description must name the filters: ${scan.simpleString(100)}")
+ assert(scan.verboseStringWithOperatorId().contains("StorageFilters:"),
+ s"EXPLAIN FORMATTED must name the filters: ${scan.verboseStringWithOperatorId()}")
+
+ intercept[IllegalArgumentException](scan.copy(storageFilters = Nil))
+ val plain = scan.copy(output = scan.output.init, storageFilters = Nil, checkedColumn = None)
+ assert(scan != plain, "equals must see storageFilters")
Review Comment:
**13. [Low] These asserts do not show that `equals` and `doCanonicalize` see `storageFilters`**
`plain` also drops `checkedColumn`, and `equals` compares `checkedColumn.isDefined` on its own. So both asserts hold even with `storageFilters` removed from `equals` or `doCanonicalize`. The case the message names, a scan with storage filters and no checked column (too wide, or codegen off), is never compared, and `hashCode` is not checked.
Other tests would still catch such a removal today. Still, could this compare two scans that differ only in `storageFilters`, both with `checkedColumn = None` and the same output, and also assert on `hashCode`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -127,6 +127,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
tableSpec.withNewLocation(newLoc)
}
+ /**
+ * The conjuncts of `filters` offered to the scan as storage filters, see
+ * `SupportsStorageFilters` and `StorageFilterPlanning.offeredConjuncts`. Only runtime blooms for
+ * now, as on the file-source side.
+ */
+ private def storageFiltersFor(
+ relation: DataSourceV2ScanRelation,
+ filters: Seq[Expression]): Seq[Expression] = {
+ if (!session.sessionState.conf.v2StorageFilterPushdownEnabled) return Nil
+ if (relation.storageFilterAttrs.isEmpty) return Nil
+ // A scan that took a pushed limit, offset, top-N or sample, or an aggregate or join, would have
+ // its readers filter ahead of it. `V2ScanRelationPushDown` leaves every such scan, and any scan
+ // a later pushdown builds, not mergeable.
+ if (!relation.mergeableScan) return Nil
Review Comment:
**18. [Low] Gating on `mergeableScan` couples storage filters to scan merging**
`mergeableScan` means "a rebuilt scan can reproduce this scan's pushdowns" (`hasBlockingPushdown`, read by `PlanMerger`). This check uses it as "no pushdown decides which rows the scan returns". The two sets agree today, but not by construction:
- A scan with pushed variant extraction keeps its rows. Yet `buildScanWithPushedVariants` leaves it not mergeable (V2ScanRelationPushDown.scala:959), so `SELECT f.k, variant_get(f.v, '$.a', 'int') FROM fact f JOIN dim d ON f.k = d.k WHERE d.x = 5` gets no storage filter.
- `hasBlockingPushdown`'s maintenance rule sorts a pushdown by whether a merge can reproduce it. Suppose a later change makes a row-selecting pushdown reproducible, e.g. an identical pushed top-N or a seeded sample. The scan then becomes mergeable and gets the bloom ahead of that pushdown, which then picks other rows. Only doc sentences and the four limit / offset / top-N / sample tests guard against that.
A dedicated flag on `DataSourceV2ScanRelation` would decouple the two. It would default to false, and the same build sites would set it from the pushdown's row semantics.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala:
##########
@@ -29,22 +29,37 @@ import org.apache.spark.sql.connector.catalog.Table
import org.apache.spark.sql.connector.read._
import org.apache.spark.sql.connector.write.RowLevelOperation.Command.DELETE
import org.apache.spark.sql.connector.write.RowLevelOperationTable
-import org.apache.spark.sql.execution.EmptyRDDWithPartitions
+import org.apache.spark.sql.execution.{EmptyRDDWithPartitions, ScalarSubquery => ExecScalarSubquery}
import org.apache.spark.sql.execution.metric.{SQLLastAttemptMetrics, SQLMetric, SQLMetrics}
+import org.apache.spark.sql.internal.connector.{StorageFilterImpl, StorageFilterRequestImpl}
import org.apache.spark.util.ArrayImplicits._
/**
* Physical plan node for scanning a batch of data from a data source v2.
+ *
+ * @param storageFilters Predicates the scan's readers may evaluate, to skip reading the other
+ * columns of the rows they reject, see [[SupportsStorageFilters]]. Non-empty
+ * only for a scan that declared their columns. Each task's readers get them
+ * ANDed into one filter, in a [[StorageFilterRequest]]. They stay in the
+ * post-scan `Filter` too.
+ * @param checkedColumn The column the readers mark the rows they checked in, last in `output`, when
+ * the planner asked for one, see `StorageFilterPlanning.checkedColumn`.
*/
case class BatchScanExec(
output: Seq[AttributeReference],
@transient scan: Scan,
runtimeFilters: Seq[Expression],
ordering: Option[Seq[SortOrder]] = None,
@transient table: Table,
- keyGroupedPartitioning: Option[Seq[Expression]] = None
+ keyGroupedPartitioning: Option[Seq[Expression]] = None,
+ storageFilters: Seq[Expression] = Nil,
+ checkedColumn: Option[AttributeReference] = None
Review Comment:
**19. [Nit] `checkedColumn` duplicates `output.last`**
Main code reads only `checkedColumn.isDefined` (L86, L92, L125), and the strategy uses its local `checked` for `skipCheckedRows`. Three things exist only to keep the two copies equal: the `require` (L60), the extra normalization in `doCanonicalize` (L151) and the test twin's remapping. `appendsCheckedColumn: Boolean = false` with `require(!appendsCheckedColumn || storageFilters.nonEmpty)` would carry the same information, with fewer places to keep in sync. If wanted, it could also check that the last output column is a `BooleanType`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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.internal.connector
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.{SparkEnv, SparkException, TaskContext}
+import org.apache.spark.executor.Executor
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.internal.config.KILL_ON_FATAL_ERROR_DEPTH
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to
+ * it. The binding is built on first use, so an instance created on the driver can be evaluated on
+ * an executor.
+ */
+class StorageFilterImpl private (
+ private val catalystExpr: CatalystExpression,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes.
+ private val numReferences = attrs.length
Review Comment:
**20. [Nit] Small cleanups in `StorageFilterImpl`**
The column list is rendered three ways (L68, L79 and L99), and the comment "a runtime bloom's literal would print every byte of the bloom" appears twice (L78 and L98). `toString` already is the safe rendering, so the internal error and the warning could interpolate `$this`, and the comment could stay once, on `toString`. Also, `numReferences` is just `attrs.length`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala:
##########
@@ -98,13 +116,22 @@ case class BatchScanExec(
override lazy val readerFactory: PartitionReaderFactory = batch.createReaderFactory()
+ // What every task's readers get, the storage filters ANDed into one. A runtime bloom arrives as
+ // a ScalarSubquery over the build side, which SparkPlan has computed by the time this is forced.
+ // The readers get its literal, so the task does not carry the subquery's plan.
+ @transient private lazy val storageFilterRequest: Option[StorageFilterRequest] =
+ storageFilters.map(_.transform { case s: ExecScalarSubquery => s.toLiteral })
Review Comment:
**6. [Medium] A merged bloom ships the whole merged struct once per conjunct**
Only the `ScalarSubquery` node is replaced here. Take a join on two declared keys:
1. `InjectRuntimeFilter` injects one bloom per key (InjectRuntimeFilter.scala:418).
2. `MergeSubplans` runs in the next batch whenever subquery reuse is on, which is the default. It merges the two into `might_contain(GetStructField(ScalarSubquery(merged), i), xxhash64(k_i))`, and both conjuncts still pass `_.isInstanceOf[BloomFilterMightContain]`.
3. Each `ScalarSubquery`'s `toLiteral` (subquery.scala:112-114) is then a separate `UnsafeRow` view of the whole `struct<bloom_1, bloom_2>`, and Java serialization writes each one in full.
So the request carries `2 x (b1 + b2)` bytes instead of `b1 + b2`. That is up to 32 MiB at the default `spark.sql.optimizer.runtime.bloomFilter.maxNumBits`. With `n` merged blooms, every task gets `n` times their total. The post-scan `Filter` ships only the deserialized bloom it uses.
Could the substitution also fold each bloom argument to its own binary literal on the driver, e.g. `Literal(bloomFilterExpression.eval(), BinaryType)` once the subquery is replaced? A test with a two-key join would cover this shape too (item 12).
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/filter/StorageFilter.java:
##########
@@ -0,0 +1,111 @@
+/*
+ * 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.expressions.filter;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.read.Scan;
+import org.apache.spark.sql.connector.read.StorageFilterRequest;
+import org.apache.spark.sql.connector.read.SupportsStorageFilters;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+/**
+ * A row-level predicate a scan may evaluate inside its readers, to skip reading the other columns
+ * of the rows it rejects. The reader reads the columns of {@link #references()} first, evaluates
+ * the filter, and reads the remaining columns only for the rows it keeps.
+ * <p>
+ * The filter is a predicate on the rows the scan returns. Spark offers none to a scan that took a
+ * pushed operation that decides which rows it returns, such as a limit, an offset, a top-N, a
+ * sample or an aggregate.
+ * <p>
+ * Spark implements the predicate, and the reader only supplies the values. For example, a scan
+ * does not have to reproduce the Spark sketch and hash of a runtime bloom filter. The values must
+ * be the ones the scan returns for those columns, in the types of {@link Scan#readSchema()}, nulls
+ * included.
+ * <p>
+ * Applying it is optional, see {@link SupportsStorageFilters}.
+ * <p>
+ * A row's result depends only on its own values. So with one reference, a reader may evaluate each
+ * distinct value once, such as each entry of a dictionary, and reuse the result for every row
+ * holding that value. A dictionary entry is evaluated as the value the scan returns for it, decoded
+ * and converted.
+ * <p>
+ * Spark hands the filter to each task in a {@link StorageFilterRequest}. The readers of a task
+ * share one instance, which is not thread-safe.
+ *
+ * @since 5.0.0
+ */
+@Evolving
+public abstract class StorageFilter extends Predicate {
Review Comment:
**4. [Nit] Does `StorageFilter` need to be a `Predicate`?**
A connector gets a `StorageFilter` only through `StorageFilterRequest.filter()`, never in a `Predicate[]`. `PartitionPredicate` is different: it travels through `SupportsRuntimeV2Filtering.filter` and `pushPredicates`. Extending `Predicate` brings two things here:
- `name() == "STORAGE_FILTER"` with no children, which a generic V2 predicate translator cannot handle;
- `GeneralScalarExpression`'s name-and-children equality, under which any two filters are equal. That is why `StorageFilterImpl` overrides `equals` and `hashCode`, which nothing calls.
A standalone `public abstract class StorageFilter implements Serializable` with `references()`, `eval` and `select` would avoid both. This is new `@Evolving` API, and removing a superclass later breaks binary compatibility, so it may be worth settling now.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -200,10 +223,20 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
// rewritten it to TrueLiteral by the time we get here.
val runtimeFilters = dynamicFilters ++ scalarSubqueryFilters
- val batchExec = BatchScanExec(relation.output, relation.scan, runtimeFilters,
- relation.ordering, relation.relation.table, relation.keyGroupedPartitioning)
- DataSourceV2Strategy.withProjectAndFilter(
- project, postScanFilters.diff(fullyPushedRuntimeFilters),
+ // Whether the scan gets the checked column does not depend on whether it is columnar. Asking
+ // that can plan the scan's partitions, which the node built afterwards would plan again. A
+ // row-based scan gets a Project for the UnsafeRow conversion anyway, so dropping the column
+ // there copies nothing more.
+ val filtersAbove = postScanFilters.diff(fullyPushedRuntimeFilters)
+ val storageFilters = storageFiltersFor(relation, filtersAbove)
+ val checked = StorageFilterPlanning.checkedColumn(
Review Comment:
**8. [Low] A row-based scan loses the checked column where it would cost nothing**
The gate in `StorageFilterPlanning.checkedColumn` is meant for columnar scans: whole-stage codegen on, and the output plus the column within `spark.sql.codegen.maxFields`. As the comment above says, a row-based scan always gets the `ProjectExec` for the UnsafeRow conversion (`needsUnsafeConversion`, kept by RemoveRedundantProjects.scala:108).
Take a row-based scan whose reader drops rejected rows, with `spark.sql.codegen.wholeStage=false` or with an output already over `maxFields`. It gets no column, so the post-scan `Filter` evaluates `xxhash64` and the bloom probes again on every row the reader kept. The same Project would have dropped the column at no extra cost. At exactly `maxFields` the column does push the Filter and the Project out of codegen, so the gate is right there.
`scan.columnarSupportMode()` answers without planning partitions for `SUPPORTED` and `UNSUPPORTED`. Could those two decide by read mode, and `PARTITION_DEFINED` keep today's gate? The doc of `StorageFilterRequest.appendsCheckedColumn` would then need to say so.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -127,6 +127,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
tableSpec.withNewLocation(newLoc)
}
+ /**
+ * The conjuncts of `filters` offered to the scan as storage filters, see
+ * `SupportsStorageFilters` and `StorageFilterPlanning.offeredConjuncts`. Only runtime blooms for
+ * now, as on the file-source side.
+ */
+ private def storageFiltersFor(
+ relation: DataSourceV2ScanRelation,
+ filters: Seq[Expression]): Seq[Expression] = {
+ if (!session.sessionState.conf.v2StorageFilterPushdownEnabled) return Nil
+ if (relation.storageFilterAttrs.isEmpty) return Nil
+ // A scan that took a pushed limit, offset, top-N or sample, or an aggregate or join, would have
+ // its readers filter ahead of it. `V2ScanRelationPushDown` leaves every such scan, and any scan
+ // a later pushdown builds, not mergeable.
+ if (!relation.mergeableScan) return Nil
+ // Over the whole table, as on the file-source side, since `SupportsStorageFilters` reserves
+ // the name.
+ if (StorageFilterPlanning.clashesWithCheckedColumn(relation.relation.output.map(_.name))) {
+ return Nil
+ }
+ StorageFilterPlanning.offeredConjuncts(filters, relation.storageFilterAttrs,
+ AttributeSet(relation.output), _.isInstanceOf[BloomFilterMightContain])
Review Comment:
**11. [Nit] Metadata columns count as columns left to skip**
`readColumns` is the whole `relation.output`, metadata columns included, while the file-source side passes only the data columns (`readDataColumns`). Take a scan that reads only the key plus, say, `_file` or `_pos`. It passes the "something left to skip" rule, and pays for the checked column, the request and the extra planning with no IO saved. V1 would offer nothing there.
Whether that is a waste depends on the connector's metadata columns. So either of these would do: leave the relation's metadata attributes out, matched by exprId against `relation.relation.output`, or add a sentence to the `offeredConjuncts` doc that the rule is an approximation.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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.internal.connector
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.{SparkEnv, SparkException, TaskContext}
+import org.apache.spark.executor.Executor
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.internal.config.KILL_ON_FATAL_ERROR_DEPTH
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to
+ * it. The binding is built on first use, so an instance created on the driver can be evaluated on
+ * an executor.
+ */
+class StorageFilterImpl private (
+ private val catalystExpr: CatalystExpression,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes.
+ private val numReferences = attrs.length
+
+ // Whether this instance has logged a failure, so a task logs only its first.
+ @transient private var loggedFailure = false
+
+ @transient private lazy val boundPredicate: InternalRow => Boolean = {
+ val bound = BindReferences.bindReference(catalystExpr, attrs)
+ CatalystPredicate.create(bound).eval
Review Comment:
**9. [Low] The bloom's literal stays alive for the whole task**
Once `boundPredicate` is built, the generated predicate holds only the deserialized bloom. But `catalystExpr` keeps the bloom's binary literal reachable from the task's `DataSourceRDD` until the task ends. At the default 8 MiB bloom, that is another 8 MiB per running task, about 128 MiB on a 16-slot executor, and more with item 6.
The PR description lists this under possible later work. V1 already handles it in a few lines: `ParquetStorageFilter.preparedPredicate` drops its expression once `Predicate.create` returns (ParquetStorageFilter.scala:120-127). Only `equals` and `hashCode` read `catalystExpr` afterwards, so they could use a hash computed up front, or the fields that remain.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterTestSource.scala:
##########
@@ -0,0 +1,303 @@
+/*
+ * 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.util
+import java.util.OptionalLong
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.connector.catalog.{SupportsRead, Table, TableCapability, TableProvider}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference, SortOrder, Transform}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.types.{BooleanType, LongType, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+/**
+ * A read-only v2 source over a fixed table of `(k, v)` longs that implements
+ * [[SupportsStorageFilters]]. Its reader gets the filter in the task and drops the rows the filter
+ * rejects. Its options:
+ * - `columnar`: reads in batches, through `StorageFilter.select`, rather than row by row;
+ * - `giveUp`: reads without the filter, as a reader does for a unit it cannot use the filter on.
+ * Here the unit is the one partition;
+ * - `markAllChecked`: reads without the filter but marks every row checked, which breaks the
+ * contract on purpose;
+ * - `declared`: the comma-separated storage-filter columns, `k` by default, a dot separating a
+ * nested field;
+ * - `valueName`: the name of the value column, `v` by default;
+ * - `push`: the comma-separated pushed operations it takes whole without applying them, of
+ * `limit`, `offset`, `topN` and `sample`, for planning tests only.
+ *
+ * A filter that throws `StorageFilterNotAppliedException` is given up for the rest of the
+ * partition too. The counters below record what the readers got, evaluated and returned.
+ */
+class StorageFilterTestSource extends TableProvider {
+
+ override def inferSchema(options: CaseInsensitiveStringMap): StructType =
+ new StructType().add("k", LongType).add(options.getOrDefault("valueName", "v"), LongType)
+
+ override def getTable(
+ schema: StructType,
+ partitioning: Array[Transform],
+ properties: util.Map[String, String]): Table = new StorageFilterTestTable(schema)
+}
+
+object StorageFilterTestSource {
+
+ /** Rows the source serves, `(k, k * 10)` for k in 0 until numRows. */
+ val numRows: Int = 200
+
+ /** Rows a batch holds at most in a columnar read. */
+ val batchSize: Int = 64
+
+ /** How many readers were created with a storage-filter request. */
+ val filteredReaders = new AtomicInteger(0)
+
+ /** How many of those requests asked for the checked column. */
+ val checkedColumnRequests = new AtomicInteger(0)
+
+ /** Rows the reader evaluated the storage filter against. */
+ val rowsEvaluated = new AtomicLong(0)
+
+ /** Rows the reader returned. */
+ val rowsReturned = new AtomicLong(0)
+
+ def reset(): Unit = {
+ filteredReaders.set(0)
+ checkedColumnRequests.set(0)
+ rowsEvaluated.set(0)
+ rowsReturned.set(0)
+ }
+}
+
+class StorageFilterTestTable(tableSchema: StructType) extends Table with SupportsRead {
+
+ override def name(): String = "storage_filter_test"
+
+ override def schema(): StructType = tableSchema
+
+ override def capabilities(): util.Set[TableCapability] =
+ Set(TableCapability.BATCH_READ).asJava
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
+ val factory = StorageFilterTestReaderFactory(
+ columnar = options.getBoolean("columnar", false),
+ giveUp = options.getBoolean("giveUp", false),
+ marksAllChecked = options.getBoolean("markAllChecked", false))
+ val declared = options.getOrDefault("declared", "k").split(",").filter(_.nonEmpty).toSeq
+ val scan = () => new StorageFilterTestScan(tableSchema, factory, declared)
+ val pushes = options.getOrDefault("push", "").split(",").filter(_.nonEmpty).toSet
+ if (pushes.nonEmpty) new StorageFilterTestPushingScanBuilder(scan, pushes) else () => scan()
+ }
+}
+
+/**
+ * Takes the pushed operations named in `pushes` whole. Its scan does not apply them, so it serves
+ * planning tests only.
+ */
+class StorageFilterTestPushingScanBuilder(scan: () => Scan, pushes: Set[String])
+ extends SupportsPushDownLimit
+ with SupportsPushDownOffset
+ with SupportsPushDownTopN
+ with SupportsPushDownTableSample {
+
+ override def build(): Scan = scan()
+ override def pushLimit(limit: Int): Boolean = pushes("limit")
+ override def pushOffset(offset: Int): Boolean = pushes("offset")
+ override def pushTopN(orders: Array[SortOrder], limit: Int): Boolean = pushes("topN")
+ override def pushTableSample(
+ lowerBound: Double,
+ upperBound: Double,
+ withReplacement: Boolean,
+ seed: Long,
+ sampleMethod: SampleMethod): Boolean = pushes("sample")
+ override def isPartiallyPushed: Boolean = false
+}
+
+class StorageFilterTestScan(
+ tableSchema: StructType,
+ factory: StorageFilterTestReaderFactory,
+ declared: Seq[String])
+ extends Scan with Batch with SupportsReportStatistics with SupportsStorageFilters {
+
+ override def readSchema(): StructType = tableSchema
+
+ override def toBatch: Batch = this
+
+ // Over the suite's application-side scan size threshold, so InjectRuntimeFilter considers this
+ // side worth filtering.
+ override def estimateStatistics(): Statistics = new Statistics {
+ override def sizeInBytes(): OptionalLong = OptionalLong.of(10L * 1024 * 1024)
+ override def numRows(): OptionalLong = OptionalLong.of(StorageFilterTestSource.numRows.toLong)
+ }
+
+ override def storageFilterAttributes(): Array[NamedReference] =
+ declared.map(name => FieldReference(name.split('.').toSeq)).toArray[NamedReference]
+
+ override def planInputPartitions(): Array[InputPartition] =
+ Array(StorageFilterTestPartition)
+
+ override def createReaderFactory(): PartitionReaderFactory = factory
+}
+
+case object StorageFilterTestPartition extends InputPartition
+
+case class StorageFilterTestReaderFactory(
+ columnar: Boolean,
+ giveUp: Boolean,
+ marksAllChecked: Boolean)
+ extends PartitionReaderFactory {
+
+ override def supportColumnarReads(partition: InputPartition): Boolean = columnar
+
+ // The filter the reader applies, none when it gives the filter up or only marks the rows.
+ private def filterOf(request: StorageFilterRequest): Option[StorageFilter] = {
+ StorageFilterTestSource.filteredReaders.incrementAndGet()
+ if (request.appendsCheckedColumn()) {
+ StorageFilterTestSource.checkedColumnRequests.incrementAndGet()
+ }
+ if (giveUp || marksAllChecked) None else Some(request.filter())
+ }
+
+ override def createReader(partition: InputPartition): PartitionReader[InternalRow] =
+ new StorageFilterTestReader(None, appendsChecked = false)
+
+ override def createReader(
+ partition: InputPartition,
+ request: StorageFilterRequest): PartitionReader[InternalRow] =
+ new StorageFilterTestReader(filterOf(request), request.appendsCheckedColumn(), marksAllChecked)
+
+ override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] =
+ new StorageFilterTestColumnarReader(None, appendsChecked = false)
+
+ override def createColumnarReader(
+ partition: InputPartition,
+ request: StorageFilterRequest): PartitionReader[ColumnarBatch] =
+ new StorageFilterTestColumnarReader(
+ filterOf(request), request.appendsCheckedColumn(), marksAllChecked)
+}
+
+/**
+ * Serves `(k, k * 10)` row by row and drops any row `filter` rejects, with the checked column
+ * last when `appendsChecked`. The filter's row layout is its `references()`, only `k` here.
+ * `filter` is cleared when it cannot be applied. `marksAllChecked` marks every row checked.
+ */
+class StorageFilterTestReader(
+ private var filter: Option[StorageFilter],
+ appendsChecked: Boolean,
+ marksAllChecked: Boolean = false)
+ extends PartitionReader[InternalRow] {
+
+ private var current: Long = -1L
+ private var row: InternalRow = _
+
+ override def next(): Boolean = {
+ var found = false
+ while (!found && current + 1 < StorageFilterTestSource.numRows) {
+ current += 1
+ val key = current
+ val passes = filter.forall { f =>
+ StorageFilterTestSource.rowsEvaluated.incrementAndGet()
+ try f.eval(new GenericInternalRow(Array[Any](key))) catch {
Review Comment:
**12. [Medium] No test covers more than one bloom, more than one reader, or runtime filtering**
Every planner test plans one bloom over `k` on a one-partition scan (`storageFilters.size == 1`, `filteredReaders == 1`). This reader ignores `references()` and builds a one-value key row, and the table has only `(k, v)`. So these paths have no end-to-end test:
- Several offered conjuncts: the `reduceOption(And)` request, its layout ordered by output position, and one `If(checked, true, _)` per conjunct in the Filter. For example, keeping only the first storage filter would pass every test, while the Filter would skip the second bloom on rows never evaluated against it.
- The `MergeSubplans` form `GetStructField(ScalarSubquery(merged), i)` (item 6).
- A scan that also implements `SupportsRuntimeV2Filtering` or `SupportsRuntimeCatalystFiltering` on the key, where a fully pushed bloom must not be offered (`filtersAbove`).
- Several readers of one task sharing one `StorageFilterImpl`, which the `StorageFilter` doc promises. This happens with coalesced partitions or the SPJ k-way merge.
Most of this would be covered by three tests: a two-key join with a probe-side conjunct that is not offered, a runtime-filtering variant of the source, and a task that reads several partitions. The fixture would need a third column, and would need to read the layout from `references()`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -127,6 +127,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
tableSpec.withNewLocation(newLoc)
}
+ /**
+ * The conjuncts of `filters` offered to the scan as storage filters, see
+ * `SupportsStorageFilters` and `StorageFilterPlanning.offeredConjuncts`. Only runtime blooms for
+ * now, as on the file-source side.
+ */
+ private def storageFiltersFor(
+ relation: DataSourceV2ScanRelation,
+ filters: Seq[Expression]): Seq[Expression] = {
+ if (!session.sessionState.conf.v2StorageFilterPushdownEnabled) return Nil
+ if (relation.storageFilterAttrs.isEmpty) return Nil
+ // A scan that took a pushed limit, offset, top-N or sample, or an aggregate or join, would have
+ // its readers filter ahead of it. `V2ScanRelationPushDown` leaves every such scan, and any scan
+ // a later pushdown builds, not mergeable.
+ if (!relation.mergeableScan) return Nil
+ // Over the whole table, as on the file-source side, since `SupportsStorageFilters` reserves
Review Comment:
**16. [Low] The stated reason for this check is circular**
"since `SupportsStorageFilters` reserves the name" explains the check by the reservation it enforces. The real reason is different. On the V2 path the checked column is found by position, but `BatchScanExec.filteredPartitions` passes the widened `output` to `PushDownUtils.replanWithRuntimeFilters`, which resolves names over it: `resolveDataSourceRuntimeFilterRefs` at PushDownUtils.scala:212 and `getPartitionPredicateSchema` at :222.
So without this check, take a table column of that name declared in `filterAttributes()` of a `SupportsRuntimeV2Filtering` scan with iterative pushdown. The query would fail on an ambiguous runtime-filter attribute, or silently drop its iterative partition predicates. d82f19a48ab gave this reason ("since `output` is matched by name later"), and 2a855b5af65 replaced it.
There are two ways to fix it:
- Restore that reason here.
- Pass `output` without the checked column to `replanWithRuntimeFilters`. The V2 path could then drop the check, the reservation and the public `CHECKED_COLUMN_NAME`, and the name would stay internal to `FileFormat`, where V1 needs it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/StorageFilterPlanning.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, Expression, ExpressionSet, If, Literal}
+import org.apache.spark.sql.connector.read.SupportsStorageFilters
+import org.apache.spark.sql.execution.WholeStageCodegenExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.BooleanType
+
+/**
+ * The planning rules file-source and DSv2 scans share for storage filters, see
+ * `FileFormat.buildReaderWithStorageFilters` and `SupportsStorageFilters`. Each strategy decides
+ * on its own whether its scan can take any.
+ */
+object StorageFilterPlanning {
+
+ /** Whether one of `names` is the checked column's, in which case no filter is offered. */
+ def clashesWithCheckedColumn(names: Seq[String]): Boolean =
+ names.exists(_.equalsIgnoreCase(SupportsStorageFilters.CHECKED_COLUMN_NAME))
+
+ /**
+ * The conjuncts of `filters` offered to the scan, each once, as an `ExpressionSet` counts them.
+ * Each one:
+ * - is deterministic, since the post-scan `Filter` takes the reader's word on the rows it
+ * checked;
+ * - references at least one column, and only `keyCandidates`;
+ * - is taken by `accepts`.
+ *
+ * Nothing is offered when the offered conjuncts reference every one of `readColumns`, since the
+ * reader then reads the same columns for the same rows either way.
+ */
+ def offeredConjuncts(
+ filters: Seq[Expression],
+ keyCandidates: AttributeSet,
+ readColumns: AttributeSet,
+ accepts: Expression => Boolean): Seq[Expression] = {
+ val offered = ExpressionSet(filters).toSeq.filter { expr =>
+ val refs = expr.references
+ expr.deterministic && refs.nonEmpty && refs.subsetOf(keyCandidates) && accepts(expr)
+ }
+ if (offered.isEmpty) return Nil
+ val storageKeyAttrs = AttributeSet.fromAttributeSets(offered.map(_.references))
+ if (readColumns.subsetOf(storageKeyAttrs)) Nil else offered
+ }
+
+ /**
+ * The column the reader marks the rows it checked in, added after `output` when a filter is
Review Comment:
**17. [Low] This doc claims more than the gate does**
Two statements here do not hold for both callers:
- "added after `output`": `FileSourceStrategy` passes `readerDataColumns ++ partitionColumns ++ constantMetadataColumns`, but places the column right after `readerDataColumns` (FileSourceStrategy.scala:344 and :354). `FileSourceScanExec` relies on that position (`output.take(requiredSchema.length)`, DataSourceScanExec.scala:818). A maintainer who trusts this line and appends the column last on the V1 path too would break the binding. It would be safer to say that the caller places it.
- "the column never changes how the scan runs", and the PR description's "requested only where dropping it copies no extra row": take a post-scan conjunct that is `CodegenFallback`, e.g. `cardinality(map_filter(m, (a, b) -> b > 0)) > 0`. `CollapseCodegenStages` keeps the Filter out of codegen (WholeStageCodegenExec.scala:1206-1212). The Project that drops the column then becomes its own stage, and copies every surviving row of a columnar scan. Without the column, no Project is planned at all.
The trade-off itself was accepted in SPARK-60034, so only the wording needs to change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #59715:
URL: https://github.com/apache/spark/pull/59715#issuecomment-6442780297
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/ff8da45b33d7fe889ce3d1712e5710b87affe6db
*Posted by `merge_spark_pr.py`*
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4379088695
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala:
##########
@@ -817,8 +817,8 @@ case class FileSourceScanExec(
protected lazy val preparedStorageFilters: Seq[Expression] = {
val requestedDataAttrs = output.take(requiredSchema.length)
storageFilters.map { expr =>
- val subqueryReplaced = expr.transform { case s: execution.ScalarSubquery => s.toLiteral }
- BindReferences.bindReference(subqueryReplaced, requestedDataAttrs)
+ BindReferences.bindReference(
+ StorageFilterPlanning.withSubqueryResults(expr), requestedDataAttrs)
Review Comment:
**3. [Low] No V1 test covers the per-bloom literal in `preparedStorageFilters`**
This changes the V1 Parquet path too. A bloom that `MergeSubplans` merged with another used to ship the whole merged struct to every task, and now ships only its own bytes. But `ParquetStorageFilterSuite` changed only an import and one call, and none of its tests runs a two-key bloom join. The helper's unit test ("a bloom read out of a merged subquery's struct ships as a literal of its own" in `StorageFilterPushDownSuite`) does not go through this call. So if this line went back to `s.toLiteral`, every test would still pass. Could `ParquetStorageFilterSuite` get a two-key counterpart of the new V2 test, one that checks that each bloom reaches the reader once?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/Scan.java:
##########
@@ -70,9 +70,10 @@ default String description() {
* {@link Table} that creates this scan returns {@link TableCapability#BATCH_READ} support in its
* {@link Table#capabilities()}.
* <p>
- * If the scan supports runtime filtering and implements {@link SupportsRuntimeFiltering},
- * this method may be called multiple times. Therefore, implementations can cache some state
- * to avoid planning the job twice.
+ * If the scan supports runtime filtering and implements {@link SupportsRuntimeFiltering}, or
+ * Spark offers it a storage filter through {@link SupportsStorageFilters}, this method may be
Review Comment:
**6. [Nit] `toBatch` can be called more than once for every scan, not only these two kinds**
As your [reply](https://github.com/apache/spark/pull/59715#discussion_r4377264760) to item 7 notes, `planInputPartitions` is already called twice with the conf off. In the default `PARTITION_DEFINED` mode, `DataSourceV2Strategy` forces `batch` through `supportsColumnar` (DataSourceV2Strategy.scala:240). The executed plan is then built from `sparkPlan.clone()` (QueryExecution.scala:393), and that copy calls `toBatch` again. `Batch.planInputPartitions` now says so without a condition. This doc still reads as if a scan with neither runtime filtering nor a storage filter is planned once, so a connector may skip caching it needs. Maybe use the same wording as in `Batch`: "Spark may call this method more than once for the same scan, for instance when it plans the scan again for runtime filtering or for a storage filter."
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,491 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GetStructField, GreaterThan, If, Literal, NumericEvalContext, PredicateHelper, Rand, XxHash64}
+import org.apache.spark.sql.catalyst.plans.logical.ReturnAnswer
+import org.apache.spark.sql.execution.{BaseLimitExec, CoalesceExec, ExecSubqueryExpression, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation}
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{BinaryType, DoubleType, LongType, StructType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper with PredicateHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
Review Comment:
**2. [Low] The new multi-bloom and multi-reader tests run only with AQE off**
`withBloomJoin` turns AQE off by default, and only the first test loop (L127-131) turns it on, over one bloom and one partition. So the two-key join (L329), the task with two readers (L370), the runtime-filtering tests and the missing-column tests all run without AQE, while AQE is the default.
Under AQE the scan that runs is a different node. `PlanAdaptiveSubqueries` gives each occurrence of the bloom's subquery its own `SubqueryExec`, so `ReuseAdaptiveSubquery` turns one of the two into a `ReusedSubqueryExec` (ReuseAdaptiveSubquery.scala:38). For the scan that means a rebuilt `BatchScanExec`, as in item 7. For the two-key join, its `storageFilters` then read the blooms through `GetStructField(ScalarSubquery(ReusedSubqueryExec(...)), i)`. No test checks that `withSubqueryResults` still gives each bloom a literal of its own there, that both conjuncts are still skipped on checked rows, or that the readers of a task still share one filter. Could the two-key and the two-reader tests also run with `aqe <- Seq(false, true)`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -254,6 +279,7 @@ private class PartitionIterator[T](
}
valuePrepared = false
val result = reader.get()
+ check.foreach(_(result))
Review Comment:
**5. [Nit] The width check now runs a capturing lambda per row**
Since `cb01a6f`, this runs for every row of a row-based read. `_(result)` captures `result`, so each call allocates a lambda unless the JIT removes it, and then calls the check function before it reads `numFields`. That is cheap next to reading a row, but it is on the per-row path of every row-based scan that has the checked column. A `match` avoids the capture:
```scala
check match {
case Some(c) => c(result)
case None =>
}
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/StorageFilterPlanning.scala:
##########
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, BloomFilterMightContain, Expression, ExpressionSet, If, Literal}
+import org.apache.spark.sql.execution.{ScalarSubquery => ExecScalarSubquery, WholeStageCodegenExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, BooleanType}
+
+/**
+ * The rules file-source and DSv2 scans share for storage filters, see
+ * `FileFormat.buildReaderWithStorageFilters` and `SupportsStorageFilters`. Each strategy decides
+ * on its own whether its scan can take any.
+ */
+object StorageFilterPlanning {
+
+ /**
+ * The conjuncts of `filters` offered to the scan, each once, as an `ExpressionSet` counts them.
+ * Each one:
+ * - is deterministic, since the post-scan `Filter` takes the reader's word on the rows it
+ * checked;
+ * - references at least one column, and only `keyCandidates`;
+ * - is taken by `accepts`.
+ *
+ * Nothing is offered when the offered conjuncts reference every one of `readColumns`, since the
+ * reader then reads the same columns for the same rows either way.
+ */
+ def offeredConjuncts(
+ filters: Seq[Expression],
+ keyCandidates: AttributeSet,
+ readColumns: AttributeSet,
+ accepts: Expression => Boolean): Seq[Expression] = {
+ val offered = ExpressionSet(filters).toSeq.filter { expr =>
+ val refs = expr.references
+ expr.deterministic && refs.nonEmpty && refs.subsetOf(keyCandidates) && accepts(expr)
+ }
+ if (offered.isEmpty) return Nil
+ val storageKeyAttrs = AttributeSet.fromAttributeSets(offered.map(_.references))
+ if (readColumns.subsetOf(storageKeyAttrs)) Nil else offered
+ }
+
+ /**
+ * The column the reader marks the rows it checked in, when a filter is `offered`. The caller
+ * places it in the scan's output. Only where whole-stage codegen is on and `output` with the
+ * column still fits it, so the column keeps the scan in whole-stage codegen. Elsewhere the
+ * Project that drops the column could copy every surviving row. A post-scan conjunct that cannot
+ * be generated still puts that Project in a stage of its own, the trade-off SPARK-60034 accepted.
+ */
+ def checkedColumn(
+ conf: SQLConf,
+ offered: Seq[Expression],
+ output: Seq[Attribute]): Option[AttributeReference] = {
+ if (offered.isEmpty) return None
+ val checked = AttributeReference(FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME, BooleanType)()
+ Option.when(conf.wholeStageEnabled &&
+ !WholeStageCodegenExec.isTooManyFields(conf, (output :+ checked).toStructType))(checked)
+ }
+
+ /**
+ * `filters` with each of `offered` skipped on the rows `checked` marks, or as they are without
+ * the column. Matched by `ExpressionSet`, as `offered` was built. An `If` keeps two things an
+ * `Or` would lose:
+ * - a null mark leaves the conjunct as it was, an `And` around it included;
+ * - subexpression elimination cannot hoist what the conjunct shares with another one ahead of
+ * the skip.
+ */
+ def skipCheckedRows(
+ filters: Seq[Expression],
+ offered: Seq[Expression],
+ checked: Option[Attribute]): Seq[Expression] = checked match {
+ case Some(c) =>
+ val offeredSet = ExpressionSet(offered)
+ filters.map(f => if (offeredSet.contains(f)) If(c, Literal.TrueLiteral, f) else f)
+ case None => filters
+ }
+
+ /**
+ * `filter` with each scalar subquery replaced by its result, for the readers. Each bloom filter
+ * becomes a literal of its own bytes. A subquery `MergeSubplans` merged returns the blooms of
+ * several filters in one struct, so a bloom read out of it would otherwise carry the whole
+ * struct.
+ */
+ def withSubqueryResults(filter: Expression): Expression = filter.transformUp {
+ case s: ExecScalarSubquery => s.toLiteral
+ case m @ BloomFilterMightContain(bloom, _) if bloom.foldable && !bloom.isInstanceOf[Literal] =>
Review Comment:
**4. [Nit] `withSubqueryResults` folds only a bloom's argument**
The fix for item 6 of my previous review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009034)) special-cases `BloomFilterMightContain`. Any other conjunct that reads a field of a merged subquery still gets the whole struct, and that struct holds every bloom of the merge. Only blooms are offered today, so this costs nothing yet. But the description names the key-column conjuncts of SPARK-60448 as the next step. If those come from a subquery that `MergeSubplans` merges with the blooms, a conjunct such as `k <= GetStructField(merged, 2)` would ship all the blooms again, once per conjunct. Folding the field access covers both cases:
```scala
def withSubqueryResults(filter: Expression): Expression = filter.transformUp {
case s: ExecScalarSubquery => s.toLiteral
case g @ GetStructField(_: Literal, _, _) => Literal(g.eval(), g.dataType)
}
```
Unlike folding any foldable expression here, a `GetStructField` over a literal cannot fail.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/filter/StorageFilter.java:
##########
@@ -0,0 +1,105 @@
+/*
+ * 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.expressions.filter;
+
+import java.io.Serializable;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.read.Scan;
+import org.apache.spark.sql.connector.read.StorageFilterRequest;
+import org.apache.spark.sql.connector.read.SupportsStorageFilters;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+/**
+ * A row-level predicate a scan may evaluate inside its readers, to skip reading the other columns
+ * of the rows it rejects. The reader reads the columns of {@link #references()} first, evaluates
+ * the filter, and reads the remaining columns only for the rows it keeps.
+ * <p>
+ * The filter is a predicate on the rows the scan returns. Spark offers none to a scan that took a
+ * pushed operation that decides which rows it returns, such as a limit, an offset, a top-N, a
+ * sample, an aggregate or a join.
+ * <p>
+ * Spark implements the predicate, and the reader only supplies the values. For example, a scan
+ * does not have to reproduce the Spark sketch and hash of a runtime bloom filter. The values must
+ * be the ones the scan returns for those columns, in the types of {@link Scan#readSchema()}, nulls
+ * included.
+ * <p>
+ * Applying it is optional, see {@link SupportsStorageFilters}.
+ * <p>
+ * A row's result depends only on its own values. So with one reference, a reader may evaluate each
+ * distinct value once, such as each entry of a dictionary, and reuse the result for every row
+ * holding that value. A dictionary entry is evaluated as the value the scan returns for it, decoded
+ * and converted.
+ * <p>
+ * Spark hands the filter to each task in a {@link StorageFilterRequest}. The readers of a task
+ * share one instance, which is not thread-safe. A reader must not copy or serialize it, since an
+ * evaluated instance may have released what a copy would be built from.
+ *
+ * @since 5.0.0
+ */
+@Evolving
+public abstract class StorageFilter implements Serializable {
Review Comment:
**8. [Nit] Should `StorageFilter` move to `connector.read`?**
Since the change for item 4 ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009027)), `StorageFilter` is a standalone class, not a `Predicate`. Yet it stays in `connector.expressions.filter`, with the V2 expressions `Predicate`, `AlwaysTrue` and `PartitionPredicate`. A connector gets it only through `StorageFilterRequest`, which sits in `connector.read` next to `SupportsStorageFilters`, and never inside an expression tree. Both classes are new, so moving them, with `StorageFilterNotAppliedException`, costs nothing before a release and would be hard after one. Keeping them where they are is fine too if you prefer the filter types together.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterTestSource.scala:
##########
@@ -0,0 +1,444 @@
+/*
+ * 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.{ByteArrayOutputStream, ObjectOutputStream}
+import java.util
+import java.util.OptionalLong
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow}
+import org.apache.spark.sql.connector.catalog.{MetadataColumn, SupportsMetadataColumns, SupportsRead, Table, TableCapability, TableProvider}
+import org.apache.spark.sql.connector.expressions.{Expressions, FieldReference, NamedReference, SortOrder, Transform}
+import org.apache.spark.sql.connector.expressions.filter.{Predicate, StorageFilter, StorageFilterNotAppliedException}
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.sql.types.{BooleanType, DataType, LongType, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+/**
+ * A read-only v2 source over a fixed table of longs that implements [[SupportsStorageFilters]]:
+ * the key `k`, the value `v` holding `k * 10`, and with `wide` also `w` holding `k * 100`. Its
+ * reader gets the filter in the task and drops the rows the filter rejects. It reads the columns
+ * Spark requires, and builds the filter's row from the filter's `references()`. Its options:
+ * - `columnar`: reads in batches, through `StorageFilter.select`, rather than row by row;
+ * - `giveUp`: reads without the filter, as a reader does for a unit it cannot use the filter on.
+ * Here the unit is the partition;
+ * - `markAllChecked`: reads without the filter but marks every row checked, which breaks the
+ * contract on purpose;
+ * - `dropChecked`: `true` never appends the checked column, and `later` appends it to the first
+ * row or batch only. Either breaks the contract on purpose;
+ * - `declared`: the comma-separated storage-filter columns, `k` by default, a dot separating a
+ * nested field;
+ * - `valueName`: the name of the value column, `v` by default;
+ * - `wide`: adds the column `w`;
+ * - `numPartitions`: splits the keys into that many input partitions, 1 by default;
+ * - `metadataColumn`: adds the metadata column `_pos`, which holds the key;
+ * - `runtimeFiltering`: `v2` makes the scan a `SupportsRuntimeV2Filtering` one over every column,
+ * with iterative pushdown on a table partitioned by `k`. `catalyst` makes it a
+ * `SupportsRuntimeCatalystFiltering` one that evaluates the filters on `k` fully. Neither
+ * applies the filters it takes, so they serve planning tests only;
+ * - `push`: the comma-separated pushed operations it takes whole without applying them, of
+ * `limit`, `offset`, `topN` and `sample`, for planning tests only.
+ *
+ * A filter that throws `StorageFilterNotAppliedException` is given up for the rest of the
+ * partition too. The counters below record what the readers got, evaluated and returned.
+ */
+class StorageFilterTestSource extends TableProvider {
+
+ override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+ val schema = new StructType().add("k", LongType).add(options.getOrDefault("valueName", "v"),
+ LongType)
+ if (options.getBoolean("wide", false)) schema.add("w", LongType) else schema
+ }
+
+ override def getTable(
+ schema: StructType,
+ partitioning: Array[Transform],
+ properties: util.Map[String, String]): Table = {
+ val options = new CaseInsensitiveStringMap(properties)
+ if (options.getBoolean("metadataColumn", false)) {
+ new StorageFilterTestTable(schema, options) with SupportsMetadataColumns {
+ override def metadataColumns(): Array[MetadataColumn] = Array(new MetadataColumn {
+ override def name(): String = StorageFilterTestSource.posName
+ override def dataType(): DataType = LongType
+ })
+ }
+ } else {
+ new StorageFilterTestTable(schema, options)
+ }
+ }
+}
+
+object StorageFilterTestSource {
+
+ /** Rows the source serves, `k` from 0 until numRows. */
+ val numRows: Int = 200
+
+ /** Rows a batch holds at most in a columnar read. */
+ val batchSize: Int = 64
+
+ /** The name of the metadata column the `metadataColumn` option adds. */
+ val posName: String = "_pos"
+
+ /** How many readers were created with a storage-filter request. */
+ val filteredReaders = new AtomicInteger(0)
+
+ /** How many of those requests asked for the checked column. */
+ val checkedColumnRequests = new AtomicInteger(0)
+
+ /** Rows the reader evaluated the storage filter against. */
+ val rowsEvaluated = new AtomicLong(0)
+
+ /** Rows the reader returned. */
+ val rowsReturned = new AtomicLong(0)
+
+ /** The serialized size of the last filter a reader got. */
+ val filterBytes = new AtomicLong(0)
+
+ /** The references of the filters the readers got, and the identities of those filters. */
+ val filterReferences = mutable.Set.empty[Seq[String]]
+ val filterIdentities = mutable.Set.empty[Int]
+
+ /** How many times a runtime-filtering scan was handed filters. */
+ val runtimeFilterCalls = new AtomicInteger(0)
+
+ def reset(): Unit = {
+ filteredReaders.set(0)
+ checkedColumnRequests.set(0)
+ rowsEvaluated.set(0)
+ rowsReturned.set(0)
+ filterBytes.set(0)
+ filterReferences.synchronized(filterReferences.clear())
+ filterIdentities.synchronized(filterIdentities.clear())
+ runtimeFilterCalls.set(0)
+ }
+
+ /** The value of column `name` in the row of `key`. Any column but `k`, `w` and `_pos` is `v`. */
+ def valueOf(name: String, key: Long): Long = name match {
+ case "k" | `posName` => key
+ case "w" => key * 100
+ case _ => key * 10
+ }
+
+ /** The columns of the row the filter evaluates, in that order. */
+ def columnsOf(filter: StorageFilter): Seq[String] = filter.references.map(_.fieldNames.head).toSeq
+
+ // Records the filter a reader got, before the reader evaluates it.
+ private[connector] def record(filter: StorageFilter): Unit = {
+ val bytes = new ByteArrayOutputStream()
+ val out = new ObjectOutputStream(bytes)
+ try out.writeObject(filter) finally out.close()
Review Comment:
**7. [Nit] The test source serializes the shared filter, against the new doc**
`filterOf` calls `record(request.filter())` for every reader (L284), and this line serializes the filter. The `StorageFilter` doc now says "A reader must not copy or serialize it" (StorageFilter.java:52). In a task with two readers, the second one records the filter after the first has evaluated it and released its expression. `filterBytes`, "The serialized size of the last filter a reader got" (L117), then drops to a size without the bloom. The two-key test reads one partition, so its assertion at StorageFilterPushDownSuite.scala:361 holds today. With `numPartitions -> 2` and a `coalesce(1)`, it would pass even if each bloom shipped twice. Recording the size only for the first reader, e.g. with `filterBytes.compareAndSet(0, bytes.size())`, keeps it right.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -204,6 +221,13 @@ class DataSourceRDD(
}
}
+ private def checkWidth(expected: Int, actual: Int): Unit = {
+ if (actual != expected) {
Review Comment:
**1. [Low] The width check passes a misplaced checked column and fails a wider batch**
This compares only the number of columns, with `!=`. Two cases come out differently from what the `StorageFilterRequest.appendsCheckedColumn` doc and the error message say:
- A reader that adds the checked column in front of the partition or metadata columns it appends last returns `n + 1` columns, so the check passes. `ColumnarToRowExec` binds `output` to `batch.column(i)` by position (Columnar.scala:149-152 and :175). The first of those columns then reads the boolean vector, and the checked attribute reads the last of them. That fails with an NPE, or, when the types happen to match, returns shifted values and skips the bloom on whatever rows that last column marks, with no error.
- A reader that appends the column at position `n`, as documented, and returns more columns after it now fails with `DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING`, although the column is there. Without a request, the plan above the scan reads only the first `output.length` columns of such a batch or row.
For a columnar read, also checking `batch.column(expected - 1).dataType() == BooleanType` would catch the first case at one call per batch. For the second, either `actual < expected`, or a message that says the width is wrong rather than that the column is missing, would do. This is the check from item 2 of my previous review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009016)).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "spark-test (via GitHub)" <gi...@apache.org>.
spark-test commented on PR #59715:
URL: https://github.com/apache/spark/pull/59715#issuecomment-6351596552
The AI review requested by @HyukjinKwon did not finish. A committer can add the `ai-review-requested` label again to retry.
<!-- SPARK_OSS_REVIEW_REQUEST 33993372670 -->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4377263148
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -187,10 +191,16 @@ class DataSourceRDD(
castPartition(split).inputPartition.iterator.flatMap { inputPartition =>
val iter = if (columnarReads) {
- val batchReader = partitionReaderFactory.createColumnarReader(inputPartition)
+ val batchReader = storageFilterRequest match {
+ case Some(request) => partitionReaderFactory.createColumnarReader(inputPartition, request)
+ case None => partitionReaderFactory.createColumnarReader(inputPartition)
+ }
new PartitionIterator[ColumnarBatch](batchReader, _.numRows, taskState)
} else {
- val rowReader = partitionReaderFactory.createReader(inputPartition)
+ val rowReader = storageFilterRequest match {
+ case Some(request) => partitionReaderFactory.createReader(inputPartition, request)
Review Comment:
Thanks, fixed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788) and [`cb01a6f`](https://github.com/apache/spark/pull/59715/commits/cb01a6f1c54dbd6e29153c6c9d6f1753a731a999). `DataSourceRDD` checks every row or batch of a reader asked for the column against the scan's output width. A reader that leaves the column out fails the task with `DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING`, which names the factory class. `StorageFilterRequest.appendsCheckedColumn` states it. A test covers both read modes, with a reader that never appends the column and one that stops after its first row or batch.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/PartitionReaderFactory.java:
##########
@@ -56,10 +56,40 @@ default PartitionReader<ColumnarBatch> createColumnarReader(InputPartition parti
"DATA_SOURCE_COLUMNAR_READER_NOT_SUPPORTED", Map.of("class", getClass().getName()));
}
+ /**
+ * Returns a row-based partition reader that may apply the storage filter of {@code request} while
+ * it reads the given {@link InputPartition}. Spark calls this instead of
+ * {@link #createReader(InputPartition)} for a {@link SupportsStorageFilters} scan it offered a
+ * filter, so the factory of such a scan must implement it for row-based reads.
+ *
+ * @since 5.0.0
+ */
+ default PartitionReader<InternalRow> createReader(
Review Comment:
I'd leave this one out. The default of `createColumnarReader(partition)` fails the same way. A factory that overrides the overload but forwards to a delegate without it would also pass a check through `getDeclaringClass`.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/filter/StorageFilter.java:
##########
@@ -0,0 +1,111 @@
+/*
+ * 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.expressions.filter;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.read.Scan;
+import org.apache.spark.sql.connector.read.StorageFilterRequest;
+import org.apache.spark.sql.connector.read.SupportsStorageFilters;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+/**
+ * A row-level predicate a scan may evaluate inside its readers, to skip reading the other columns
+ * of the rows it rejects. The reader reads the columns of {@link #references()} first, evaluates
+ * the filter, and reads the remaining columns only for the rows it keeps.
+ * <p>
+ * The filter is a predicate on the rows the scan returns. Spark offers none to a scan that took a
+ * pushed operation that decides which rows it returns, such as a limit, an offset, a top-N, a
+ * sample or an aggregate.
+ * <p>
+ * Spark implements the predicate, and the reader only supplies the values. For example, a scan
+ * does not have to reproduce the Spark sketch and hash of a runtime bloom filter. The values must
+ * be the ones the scan returns for those columns, in the types of {@link Scan#readSchema()}, nulls
+ * included.
+ * <p>
+ * Applying it is optional, see {@link SupportsStorageFilters}.
+ * <p>
+ * A row's result depends only on its own values. So with one reference, a reader may evaluate each
+ * distinct value once, such as each entry of a dictionary, and reuse the result for every row
+ * holding that value. A dictionary entry is evaluated as the value the scan returns for it, decoded
+ * and converted.
+ * <p>
+ * Spark hands the filter to each task in a {@link StorageFilterRequest}. The readers of a task
+ * share one instance, which is not thread-safe.
+ *
+ * @since 5.0.0
+ */
+@Evolving
+public abstract class StorageFilter extends Predicate {
Review Comment:
Thanks, changed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788). `StorageFilter` is now a standalone abstract class that implements `Serializable`. `NAME`, the constructor, and `StorageFilterImpl`'s `equals` and `hashCode` are gone.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/filter/StorageFilterNotAppliedException.java:
##########
@@ -0,0 +1,52 @@
+/*
+ * 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.expressions.filter;
+
+import java.util.Map;
+
+import org.apache.spark.SparkThrowable;
+import org.apache.spark.SparkThrowableHelper;
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * Thrown by a {@link StorageFilter} that cannot be applied to a row, for instance because
+ * evaluating it raised an error. The reader catches it and reads the rest of the unit it was
+ * filtering, such as a file or a row group, without the filter. The cause is the error the filter
+ * met.
+ * <p>
+ * It is checked, so a reader handles it where it can fall back to reading without the filter. A
Review Comment:
Thanks, added in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788), in the exception's Javadoc and in `docs/sql-v2-data-sources.md`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala:
##########
@@ -98,13 +116,22 @@ case class BatchScanExec(
override lazy val readerFactory: PartitionReaderFactory = batch.createReaderFactory()
+ // What every task's readers get, the storage filters ANDed into one. A runtime bloom arrives as
+ // a ScalarSubquery over the build side, which SparkPlan has computed by the time this is forced.
+ // The readers get its literal, so the task does not carry the subquery's plan.
+ @transient private lazy val storageFilterRequest: Option[StorageFilterRequest] =
+ storageFilters.map(_.transform { case s: ExecScalarSubquery => s.toLiteral })
Review Comment:
Thanks, fixed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788). I measured it first. In a two-key join, a storage filter with one offered 1 MB bloom shipped 2.1 MB, and with both blooms 4.2 MB instead of 2.1 MB. `StorageFilterPlanning.withSubqueryResults` now folds each bloom argument to its own binary literal. The file-source side had the same problem in `FileSourceScanExec.preparedStorageFilters`, so it uses the same helper. A new two-key join test checks that the filter ships each bloom once.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -200,10 +223,20 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
// rewritten it to TrueLiteral by the time we get here.
val runtimeFilters = dynamicFilters ++ scalarSubqueryFilters
- val batchExec = BatchScanExec(relation.output, relation.scan, runtimeFilters,
- relation.ordering, relation.relation.table, relation.keyGroupedPartitioning)
- DataSourceV2Strategy.withProjectAndFilter(
- project, postScanFilters.diff(fullyPushedRuntimeFilters),
+ // Whether the scan gets the checked column does not depend on whether it is columnar. Asking
+ // that can plan the scan's partitions, which the node built afterwards would plan again. A
+ // row-based scan gets a Project for the UnsafeRow conversion anyway, so dropping the column
+ // there copies nothing more.
+ val filtersAbove = postScanFilters.diff(fullyPushedRuntimeFilters)
+ val storageFilters = storageFiltersFor(relation, filtersAbove)
+ val checked = StorageFilterPlanning.checkedColumn(
+ session.sessionState.conf, storageFilters, relation.output)
+ val batchExec = BatchScanExec(relation.output ++ checked, relation.scan, runtimeFilters,
Review Comment:
Thanks. The mechanism is right. With the test source our counts differ a bit: 6 `toBatch` and 3 `planInputPartitions` calls with AQE on, and 3 and 2 with the conf off. Both consequences already hold on master, though:
- `planInputPartitions` is called twice with the conf off too;
- SPARK-56467 puts the same subquery into `runtimeFilters`, so `ReuseAdaptiveSubquery` copies a runtime-filtered scan under a `GroupPartitionsExec` the same way.
So I'd leave carrying the planned `Batch` over to the copies for a separate change, since it touches every V2 scan, runtime filtering included. I listed it under the possible follow-ups. In [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788) and [`cb01a6f`](https://github.com/apache/spark/pull/59715/commits/cb01a6f1c54dbd6e29153c6c9d6f1753a731a999) `Batch.planInputPartitions` says Spark may call it more than once, and that a source whose planning is costly can cache the partitions in its `Scan` and plan them again once runtime filtering narrows the scan. The `GroupPartitionsExec` doc names `ReuseAdaptiveSubquery` too.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -200,10 +223,20 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
// rewritten it to TrueLiteral by the time we get here.
val runtimeFilters = dynamicFilters ++ scalarSubqueryFilters
- val batchExec = BatchScanExec(relation.output, relation.scan, runtimeFilters,
- relation.ordering, relation.relation.table, relation.keyGroupedPartitioning)
- DataSourceV2Strategy.withProjectAndFilter(
- project, postScanFilters.diff(fullyPushedRuntimeFilters),
+ // Whether the scan gets the checked column does not depend on whether it is columnar. Asking
+ // that can plan the scan's partitions, which the node built afterwards would plan again. A
+ // row-based scan gets a Project for the UnsafeRow conversion anyway, so dropping the column
+ // there copies nothing more.
+ val filtersAbove = postScanFilters.diff(fullyPushedRuntimeFilters)
+ val storageFilters = storageFiltersFor(relation, filtersAbove)
+ val checked = StorageFilterPlanning.checkedColumn(
Review Comment:
I'd leave this one. The column saves one `xxhash64` and one bloom probe per row the reader returned, and those are the rows the bloom kept. The gain also needs a connector that overrides `columnarSupportMode`. Iceberg does not, in any of its Spark versions, so its scans are `PARTITION_DEFINED`. One gate shared with the file-source side seems simpler to me.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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.internal.connector
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.{SparkEnv, SparkException, TaskContext}
+import org.apache.spark.executor.Executor
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.internal.config.KILL_ON_FATAL_ERROR_DEPTH
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to
+ * it. The binding is built on first use, so an instance created on the driver can be evaluated on
+ * an executor.
+ */
+class StorageFilterImpl private (
+ private val catalystExpr: CatalystExpression,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes.
+ private val numReferences = attrs.length
+
+ // Whether this instance has logged a failure, so a task logs only its first.
+ @transient private var loggedFailure = false
+
+ @transient private lazy val boundPredicate: InternalRow => Boolean = {
+ val bound = BindReferences.bindReference(catalystExpr, attrs)
+ CatalystPredicate.create(bound).eval
Review Comment:
Thanks, fixed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788). `StorageFilterImpl` drops its expression once it has built the predicate, as `ParquetStorageFilter` does. With `equals` and `hashCode` gone (item 4), nothing else reads it. So an evaluated instance cannot be copied, which the `StorageFilter` doc states since [`cb01a6f`](https://github.com/apache/spark/pull/59715/commits/cb01a6f1c54dbd6e29153c6c9d6f1753a731a999), and a unit test pins. The 8 MiB is the cap, though. A bloom built without a row count is 1 MiB.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -200,10 +223,20 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
// rewritten it to TrueLiteral by the time we get here.
val runtimeFilters = dynamicFilters ++ scalarSubqueryFilters
- val batchExec = BatchScanExec(relation.output, relation.scan, runtimeFilters,
- relation.ordering, relation.relation.table, relation.keyGroupedPartitioning)
- DataSourceV2Strategy.withProjectAndFilter(
- project, postScanFilters.diff(fullyPushedRuntimeFilters),
+ // Whether the scan gets the checked column does not depend on whether it is columnar. Asking
+ // that can plan the scan's partitions, which the node built afterwards would plan again. A
+ // row-based scan gets a Project for the UnsafeRow conversion anyway, so dropping the column
+ // there copies nothing more.
+ val filtersAbove = postScanFilters.diff(fullyPushedRuntimeFilters)
+ val storageFilters = storageFiltersFor(relation, filtersAbove)
+ val checked = StorageFilterPlanning.checkedColumn(
+ session.sessionState.conf, storageFilters, relation.output)
+ val batchExec = BatchScanExec(relation.output ++ checked, relation.scan, runtimeFilters,
+ relation.ordering, relation.relation.table, relation.keyGroupedPartitioning,
+ storageFilters, checked)
+ // With the checked column the scan's output differs from `project`, so a Project drops it.
+ DataSourceV2Strategy.withProjectAndFilter(project,
+ StorageFilterPlanning.skipCheckedRows(filtersAbove, storageFilters, checked),
Review Comment:
Thanks. I measured the non-AQE case with `reuseSubquery=false`. Exchange reuse still reuses the bloom's partial-aggregate shuffle, so only the final aggregation runs twice (3 jobs instead of 2). The file-source side and master's runtime-filter routing do the same. The description now says so, next to the note on the task binary.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -127,6 +127,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
tableSpec.withNewLocation(newLoc)
}
+ /**
+ * The conjuncts of `filters` offered to the scan as storage filters, see
+ * `SupportsStorageFilters` and `StorageFilterPlanning.offeredConjuncts`. Only runtime blooms for
+ * now, as on the file-source side.
+ */
+ private def storageFiltersFor(
+ relation: DataSourceV2ScanRelation,
+ filters: Seq[Expression]): Seq[Expression] = {
+ if (!session.sessionState.conf.v2StorageFilterPushdownEnabled) return Nil
+ if (relation.storageFilterAttrs.isEmpty) return Nil
+ // A scan that took a pushed limit, offset, top-N or sample, or an aggregate or join, would have
+ // its readers filter ahead of it. `V2ScanRelationPushDown` leaves every such scan, and any scan
+ // a later pushdown builds, not mergeable.
+ if (!relation.mergeableScan) return Nil
+ // Over the whole table, as on the file-source side, since `SupportsStorageFilters` reserves
+ // the name.
+ if (StorageFilterPlanning.clashesWithCheckedColumn(relation.relation.output.map(_.name))) {
+ return Nil
+ }
+ StorageFilterPlanning.offeredConjuncts(filters, relation.storageFilterAttrs,
+ AttributeSet(relation.output), _.isInstanceOf[BloomFilterMightContain])
Review Comment:
Thanks, fixed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788) and [`cb01a6f`](https://github.com/apache/spark/pull/59715/commits/cb01a6f1c54dbd6e29153c6c9d6f1753a731a999). The rule leaves out the relation's metadata columns, so it counts the data columns, as the file-source side does. They are matched by name, since a variant-extraction scan's output has fresh exprIds. A test reads `k` with a metadata column, and then with `v` as well.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterTestSource.scala:
##########
@@ -0,0 +1,303 @@
+/*
+ * 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.util
+import java.util.OptionalLong
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.connector.catalog.{SupportsRead, Table, TableCapability, TableProvider}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference, SortOrder, Transform}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.types.{BooleanType, LongType, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+/**
+ * A read-only v2 source over a fixed table of `(k, v)` longs that implements
+ * [[SupportsStorageFilters]]. Its reader gets the filter in the task and drops the rows the filter
+ * rejects. Its options:
+ * - `columnar`: reads in batches, through `StorageFilter.select`, rather than row by row;
+ * - `giveUp`: reads without the filter, as a reader does for a unit it cannot use the filter on.
+ * Here the unit is the one partition;
+ * - `markAllChecked`: reads without the filter but marks every row checked, which breaks the
+ * contract on purpose;
+ * - `declared`: the comma-separated storage-filter columns, `k` by default, a dot separating a
+ * nested field;
+ * - `valueName`: the name of the value column, `v` by default;
+ * - `push`: the comma-separated pushed operations it takes whole without applying them, of
+ * `limit`, `offset`, `topN` and `sample`, for planning tests only.
+ *
+ * A filter that throws `StorageFilterNotAppliedException` is given up for the rest of the
+ * partition too. The counters below record what the readers got, evaluated and returned.
+ */
+class StorageFilterTestSource extends TableProvider {
+
+ override def inferSchema(options: CaseInsensitiveStringMap): StructType =
+ new StructType().add("k", LongType).add(options.getOrDefault("valueName", "v"), LongType)
+
+ override def getTable(
+ schema: StructType,
+ partitioning: Array[Transform],
+ properties: util.Map[String, String]): Table = new StorageFilterTestTable(schema)
+}
+
+object StorageFilterTestSource {
+
+ /** Rows the source serves, `(k, k * 10)` for k in 0 until numRows. */
+ val numRows: Int = 200
+
+ /** Rows a batch holds at most in a columnar read. */
+ val batchSize: Int = 64
+
+ /** How many readers were created with a storage-filter request. */
+ val filteredReaders = new AtomicInteger(0)
+
+ /** How many of those requests asked for the checked column. */
+ val checkedColumnRequests = new AtomicInteger(0)
+
+ /** Rows the reader evaluated the storage filter against. */
+ val rowsEvaluated = new AtomicLong(0)
+
+ /** Rows the reader returned. */
+ val rowsReturned = new AtomicLong(0)
+
+ def reset(): Unit = {
+ filteredReaders.set(0)
+ checkedColumnRequests.set(0)
+ rowsEvaluated.set(0)
+ rowsReturned.set(0)
+ }
+}
+
+class StorageFilterTestTable(tableSchema: StructType) extends Table with SupportsRead {
+
+ override def name(): String = "storage_filter_test"
+
+ override def schema(): StructType = tableSchema
+
+ override def capabilities(): util.Set[TableCapability] =
+ Set(TableCapability.BATCH_READ).asJava
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
+ val factory = StorageFilterTestReaderFactory(
+ columnar = options.getBoolean("columnar", false),
+ giveUp = options.getBoolean("giveUp", false),
+ marksAllChecked = options.getBoolean("markAllChecked", false))
+ val declared = options.getOrDefault("declared", "k").split(",").filter(_.nonEmpty).toSeq
+ val scan = () => new StorageFilterTestScan(tableSchema, factory, declared)
+ val pushes = options.getOrDefault("push", "").split(",").filter(_.nonEmpty).toSet
+ if (pushes.nonEmpty) new StorageFilterTestPushingScanBuilder(scan, pushes) else () => scan()
+ }
+}
+
+/**
+ * Takes the pushed operations named in `pushes` whole. Its scan does not apply them, so it serves
+ * planning tests only.
+ */
+class StorageFilterTestPushingScanBuilder(scan: () => Scan, pushes: Set[String])
+ extends SupportsPushDownLimit
+ with SupportsPushDownOffset
+ with SupportsPushDownTopN
+ with SupportsPushDownTableSample {
+
+ override def build(): Scan = scan()
+ override def pushLimit(limit: Int): Boolean = pushes("limit")
+ override def pushOffset(offset: Int): Boolean = pushes("offset")
+ override def pushTopN(orders: Array[SortOrder], limit: Int): Boolean = pushes("topN")
+ override def pushTableSample(
+ lowerBound: Double,
+ upperBound: Double,
+ withReplacement: Boolean,
+ seed: Long,
+ sampleMethod: SampleMethod): Boolean = pushes("sample")
+ override def isPartiallyPushed: Boolean = false
+}
+
+class StorageFilterTestScan(
+ tableSchema: StructType,
+ factory: StorageFilterTestReaderFactory,
+ declared: Seq[String])
+ extends Scan with Batch with SupportsReportStatistics with SupportsStorageFilters {
+
+ override def readSchema(): StructType = tableSchema
+
+ override def toBatch: Batch = this
+
+ // Over the suite's application-side scan size threshold, so InjectRuntimeFilter considers this
+ // side worth filtering.
+ override def estimateStatistics(): Statistics = new Statistics {
+ override def sizeInBytes(): OptionalLong = OptionalLong.of(10L * 1024 * 1024)
+ override def numRows(): OptionalLong = OptionalLong.of(StorageFilterTestSource.numRows.toLong)
+ }
+
+ override def storageFilterAttributes(): Array[NamedReference] =
+ declared.map(name => FieldReference(name.split('.').toSeq)).toArray[NamedReference]
+
+ override def planInputPartitions(): Array[InputPartition] =
+ Array(StorageFilterTestPartition)
+
+ override def createReaderFactory(): PartitionReaderFactory = factory
+}
+
+case object StorageFilterTestPartition extends InputPartition
+
+case class StorageFilterTestReaderFactory(
+ columnar: Boolean,
+ giveUp: Boolean,
+ marksAllChecked: Boolean)
+ extends PartitionReaderFactory {
+
+ override def supportColumnarReads(partition: InputPartition): Boolean = columnar
+
+ // The filter the reader applies, none when it gives the filter up or only marks the rows.
+ private def filterOf(request: StorageFilterRequest): Option[StorageFilter] = {
+ StorageFilterTestSource.filteredReaders.incrementAndGet()
+ if (request.appendsCheckedColumn()) {
+ StorageFilterTestSource.checkedColumnRequests.incrementAndGet()
+ }
+ if (giveUp || marksAllChecked) None else Some(request.filter())
+ }
+
+ override def createReader(partition: InputPartition): PartitionReader[InternalRow] =
+ new StorageFilterTestReader(None, appendsChecked = false)
+
+ override def createReader(
+ partition: InputPartition,
+ request: StorageFilterRequest): PartitionReader[InternalRow] =
+ new StorageFilterTestReader(filterOf(request), request.appendsCheckedColumn(), marksAllChecked)
+
+ override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] =
+ new StorageFilterTestColumnarReader(None, appendsChecked = false)
+
+ override def createColumnarReader(
+ partition: InputPartition,
+ request: StorageFilterRequest): PartitionReader[ColumnarBatch] =
+ new StorageFilterTestColumnarReader(
+ filterOf(request), request.appendsCheckedColumn(), marksAllChecked)
+}
+
+/**
+ * Serves `(k, k * 10)` row by row and drops any row `filter` rejects, with the checked column
+ * last when `appendsChecked`. The filter's row layout is its `references()`, only `k` here.
+ * `filter` is cleared when it cannot be applied. `marksAllChecked` marks every row checked.
+ */
+class StorageFilterTestReader(
+ private var filter: Option[StorageFilter],
+ appendsChecked: Boolean,
+ marksAllChecked: Boolean = false)
+ extends PartitionReader[InternalRow] {
+
+ private var current: Long = -1L
+ private var row: InternalRow = _
+
+ override def next(): Boolean = {
+ var found = false
+ while (!found && current + 1 < StorageFilterTestSource.numRows) {
+ current += 1
+ val key = current
+ val passes = filter.forall { f =>
+ StorageFilterTestSource.rowsEvaluated.incrementAndGet()
+ try f.eval(new GenericInternalRow(Array[Any](key))) catch {
Review Comment:
Thanks, added in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788). The test source can now have a third column, prunes columns, and builds the filter's row from `references()`. It can also have several partitions, a metadata column and runtime filtering. New tests:
- a two-key join, where one merged subquery returns both blooms. The scan gets one filter over both columns, ships each bloom once, and a conjunct over the third column stays in the `Filter`;
- a `SupportsRuntimeV2Filtering` scan with iterative pushdown, which gets the bloom and takes it as a runtime filter too;
- a `SupportsRuntimeCatalystFiltering` scan that evaluates the bloom fully, which is offered nothing;
- `coalesce(1)` over two partitions, whose readers share the one filter their task deserialized.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,314 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GreaterThan, If, Literal, NumericEvalContext, Rand}
+import org.apache.spark.sql.execution.{BaseLimitExec, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{DoubleType, LongType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
+ withSQLConf(bloomJoinConf ++ conf: _*) {
+ withTable("bf_build") {
+ spark.range(0, 30).selectExpr("id AS bk", "id AS bv")
+ .write.format("parquet").saveAsTable("bf_build")
+ body
+ }
+ }
+ }
+
+ // A join whose build side is small and carries a selective predicate, which is the shape
+ // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`.
+ private def bloomJoin(probe: DataFrame): DataFrame = {
+ val build = spark.table("bf_build").where("bv = 5")
+ probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*)
+ }
+
+ private def v2Probe(options: (String, String)*): DataFrame =
+ spark.read.format(sourceName).options(options.toMap).load()
+
+ private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME
+
+ private def scanOf(plan: SparkPlan): BatchScanExec =
+ collect(plan) { case b: BatchScanExec => b }.headOption
+ .getOrElse(fail(s"no BatchScanExec in plan:\n$plan"))
+
+ // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows
+ // the reader marks as checked.
+ private def bloomsAbove(plan: SparkPlan): (Int, Int) = {
+ val conditions = collect(plan) { case f: FilterExec => f.condition }
+ val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum
+ val skipped = conditions.map(_.collect {
+ case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain)
+ if a.name == checkedName => 1
+ }.sum).sum
+ (all - skipped, skipped)
+ }
+
+ // The rows the Filter on the scan returned, counted before the join.
+ private def rowsOutOfFilter(plan: SparkPlan): Long = {
+ val filters = collect(plan) {
+ case f: FilterExec if f.exists(_.isInstanceOf[BatchScanExec]) => f
+ }
+ assert(filters.size == 1, s"expected one Filter on the scan:\n$plan")
+ filters.head.metrics("numOutputRows").value
+ }
+
+ // Runs the bloom join over the v2 source read with `options`, and returns its rows and plan.
+ private def runBloomJoin(options: (String, String)*): (Set[(Long, Long)], SparkPlan) = {
+ StorageFilterTestSource.reset()
+ val df = bloomJoin(v2Probe(options: _*))
+ (df.collect().map(r => (r.getLong(0), r.getLong(1))).toSet, df.queryExecution.executedPlan)
+ }
+
+ // The build side keeps `bv = 5`, so the join returns that one row.
+ private val joinResult = Set((5L, 50L))
+
+ for (columnar <- Seq(false, true); aqe <- Seq(false, true)) {
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind scan gets the checked column, and the Filter skips the bloom on checked rows " +
+ s"(AQE $aqe)") {
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val (result, plan) = runBloomJoin("columnar" -> columnar.toString)
+ val scan = scanOf(plan)
+ assert(scan.supportsColumnar == columnar, s"the scan must read $kind:\n$plan")
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(scan.checkedColumn.isDefined && scan.checkedColumn == scan.output.lastOption,
+ s"the scan must get the checked column, last:\n$plan")
+ assert(bloomsAbove(plan) == (0, 1),
+ s"the Filter must skip the bloom on checked rows:\n$plan")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(),
+ s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " +
+ s"${StorageFilterTestSource.rowsReturned.get()}")
+ val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output }
+ assert(probeSide.exists(!_.exists(_.name == checkedName)),
+ s"the Project above the scan must drop the checked column; got $probeSide")
+ assert(StorageFilterTestSource.filteredReaders.get() == 1,
+ s"the one partition's reader must get the filter; " +
+ s"got ${StorageFilterTestSource.filteredReaders.get()}")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 1,
+ "the reader must be asked for the checked column")
+ assert(result == joinResult, s"got $result")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ s"the reader must evaluate every row; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(false, true).foreach { columnar =>
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind reader that gives the filter up marks its rows unchecked, and the Filter " +
+ "filters them") {
+ withBloomJoin() {
+ val (result, plan) = runBloomJoin("columnar" -> columnar.toString, "giveUp" -> "true")
+ assert(bloomsAbove(plan) == (0, 1), s"the plan is that of a filtering reader:\n$plan")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == 0,
+ s"the reader must evaluate nothing; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() == StorageFilterTestSource.numRows,
+ s"the reader must return every row; got ${StorageFilterTestSource.rowsReturned.get()}")
+ assert(rowsOutOfFilter(plan) < StorageFilterTestSource.numRows,
+ s"the Filter must apply the bloom to unchecked rows; got ${rowsOutOfFilter(plan)}")
+ assert(result == joinResult, s"got $result")
+ }
+ }
+
+ test(s"the Filter takes a $kind reader's word on the rows it marks checked") {
+ // This reader breaks the contract on purpose: it applies no filter but marks every row
+ // checked. Only skipping the bloom on those rows keeps all of them.
+ withBloomJoin() {
+ val (result, plan) =
+ runBloomJoin("columnar" -> columnar.toString, "markAllChecked" -> "true")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.numRows,
+ s"the Filter must keep every row marked checked; got ${rowsOutOfFilter(plan)}")
+ assert(result == joinResult, s"the join must still filter; got $result")
+ }
+ }
+ }
+
+ test("a filter that cannot be applied is given up, and its rows come back unchecked") {
+ // `1 / k` under ANSI throws at k = 0, the partition's first row, so the reader gives the
+ // filter up before it keeps or drops anything.
+ val k = AttributeReference("k", LongType)()
+ val inverse = Divide(Literal(1.0), Cast(k, DoubleType), NumericEvalContext(EvalMode.ANSI))
+ val filter = StorageFilterImpl(GreaterThan(inverse, Literal(0.0)), Seq(k))
+ val reader = new StorageFilterTestColumnarReader(Some(filter), appendsChecked = true)
+ val rows = mutable.ArrayBuffer.empty[(Long, Boolean)]
+ try {
+ while (reader.next()) {
+ val batch = reader.get()
+ (0 until batch.numRows()).foreach { i =>
+ rows += ((batch.column(0).getLong(i), batch.column(2).getBoolean(i)))
+ }
+ }
+ } finally {
+ reader.close()
+ }
+ assert(rows == (0L until StorageFilterTestSource.numRows.toLong).map(_ -> false),
+ s"expected every row, unchecked; got ${rows.take(5)} and ${rows.size} rows in all")
+ }
+
+ // Under a limit of two fields, the checked column would make three.
+ Seq(
+ ("too wide for whole-stage codegen", SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> "2"),
+ ("without whole-stage codegen", SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false")
+ ).foreach { case (when, conf) =>
+ test(s"a scan $when gets the bloom without the checked column") {
+ withBloomJoin(conf) {
+ val (result, plan) = runBloomJoin("columnar" -> "true")
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(scan.checkedColumn.isEmpty, s"the scan must not get the checked column:\n$plan")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 0,
+ "the reader must not be asked for the checked column")
+ assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan")
+ assert(result == joinResult, s"got $result")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(
+ ("with the conf off", Seq(SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "false"), Nil),
+ ("to a scan that declares no column", Nil, Seq("declared" -> "")),
+ ("to a scan that declares only another column", Nil, Seq("declared" -> "v")),
+ ("to a scan that declares only a nested field", Nil, Seq("declared" -> "k.x")),
+ ("to a scan with a column named like the checked column", Nil,
+ Seq("valueName" -> checkedName.toUpperCase(java.util.Locale.ROOT)))
+ ).foreach { case (when, conf, options) =>
+ test(s"nothing is offered $when") {
+ withBloomJoin(conf: _*) {
+ val (result, plan) = runBloomJoin(options :+ ("columnar" -> "true"): _*)
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.isEmpty && scan.checkedColumn.isEmpty,
+ s"the scan must get no storage filter:\n$plan")
+ assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan")
+ assert(StorageFilterTestSource.filteredReaders.get() == 0,
+ "no reader may get a storage filter")
+ assert(result == joinResult, s"got $result")
+ }
+ }
+ }
+
+ test("the storage filter shows in the scan description and separates canonicalized plans") {
+ withBloomJoin() {
+ val plan = bloomJoin(v2Probe("columnar" -> "true")).queryExecution.executedPlan
+ val scan = scanOf(plan)
+ assert(scan.simpleString(100).contains("StorageFilters:"),
+ s"the scan description must name the filters: ${scan.simpleString(100)}")
+ assert(scan.verboseStringWithOperatorId().contains("StorageFilters:"),
+ s"EXPLAIN FORMATTED must name the filters: ${scan.verboseStringWithOperatorId()}")
+
+ intercept[IllegalArgumentException](scan.copy(storageFilters = Nil))
+ val plain = scan.copy(output = scan.output.init, storageFilters = Nil, checkedColumn = None)
+ assert(scan != plain, "equals must see storageFilters")
Review Comment:
Thanks, fixed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788). The test now compares two scans that differ only in `storageFilters`, both without the checked column, through `equals`, `hashCode` and `sameResult`. It fails with `storageFilters` taken out of `equals`.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,314 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GreaterThan, If, Literal, NumericEvalContext, Rand}
+import org.apache.spark.sql.execution.{BaseLimitExec, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{DoubleType, LongType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
+ withSQLConf(bloomJoinConf ++ conf: _*) {
+ withTable("bf_build") {
+ spark.range(0, 30).selectExpr("id AS bk", "id AS bv")
+ .write.format("parquet").saveAsTable("bf_build")
+ body
+ }
+ }
+ }
+
+ // A join whose build side is small and carries a selective predicate, which is the shape
+ // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`.
+ private def bloomJoin(probe: DataFrame): DataFrame = {
+ val build = spark.table("bf_build").where("bv = 5")
+ probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*)
+ }
+
+ private def v2Probe(options: (String, String)*): DataFrame =
+ spark.read.format(sourceName).options(options.toMap).load()
+
+ private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME
+
+ private def scanOf(plan: SparkPlan): BatchScanExec =
+ collect(plan) { case b: BatchScanExec => b }.headOption
+ .getOrElse(fail(s"no BatchScanExec in plan:\n$plan"))
+
+ // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows
+ // the reader marks as checked.
+ private def bloomsAbove(plan: SparkPlan): (Int, Int) = {
+ val conditions = collect(plan) { case f: FilterExec => f.condition }
+ val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum
+ val skipped = conditions.map(_.collect {
+ case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain)
+ if a.name == checkedName => 1
+ }.sum).sum
+ (all - skipped, skipped)
+ }
+
+ // The rows the Filter on the scan returned, counted before the join.
+ private def rowsOutOfFilter(plan: SparkPlan): Long = {
+ val filters = collect(plan) {
+ case f: FilterExec if f.exists(_.isInstanceOf[BatchScanExec]) => f
+ }
+ assert(filters.size == 1, s"expected one Filter on the scan:\n$plan")
+ filters.head.metrics("numOutputRows").value
+ }
+
+ // Runs the bloom join over the v2 source read with `options`, and returns its rows and plan.
+ private def runBloomJoin(options: (String, String)*): (Set[(Long, Long)], SparkPlan) = {
+ StorageFilterTestSource.reset()
+ val df = bloomJoin(v2Probe(options: _*))
+ (df.collect().map(r => (r.getLong(0), r.getLong(1))).toSet, df.queryExecution.executedPlan)
+ }
+
+ // The build side keeps `bv = 5`, so the join returns that one row.
+ private val joinResult = Set((5L, 50L))
+
+ for (columnar <- Seq(false, true); aqe <- Seq(false, true)) {
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind scan gets the checked column, and the Filter skips the bloom on checked rows " +
+ s"(AQE $aqe)") {
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val (result, plan) = runBloomJoin("columnar" -> columnar.toString)
+ val scan = scanOf(plan)
+ assert(scan.supportsColumnar == columnar, s"the scan must read $kind:\n$plan")
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(scan.checkedColumn.isDefined && scan.checkedColumn == scan.output.lastOption,
+ s"the scan must get the checked column, last:\n$plan")
+ assert(bloomsAbove(plan) == (0, 1),
+ s"the Filter must skip the bloom on checked rows:\n$plan")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(),
+ s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " +
+ s"${StorageFilterTestSource.rowsReturned.get()}")
+ val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output }
+ assert(probeSide.exists(!_.exists(_.name == checkedName)),
+ s"the Project above the scan must drop the checked column; got $probeSide")
+ assert(StorageFilterTestSource.filteredReaders.get() == 1,
+ s"the one partition's reader must get the filter; " +
+ s"got ${StorageFilterTestSource.filteredReaders.get()}")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 1,
+ "the reader must be asked for the checked column")
+ assert(result == joinResult, s"got $result")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ s"the reader must evaluate every row; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(false, true).foreach { columnar =>
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind reader that gives the filter up marks its rows unchecked, and the Filter " +
+ "filters them") {
+ withBloomJoin() {
+ val (result, plan) = runBloomJoin("columnar" -> columnar.toString, "giveUp" -> "true")
+ assert(bloomsAbove(plan) == (0, 1), s"the plan is that of a filtering reader:\n$plan")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == 0,
+ s"the reader must evaluate nothing; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() == StorageFilterTestSource.numRows,
+ s"the reader must return every row; got ${StorageFilterTestSource.rowsReturned.get()}")
+ assert(rowsOutOfFilter(plan) < StorageFilterTestSource.numRows,
+ s"the Filter must apply the bloom to unchecked rows; got ${rowsOutOfFilter(plan)}")
+ assert(result == joinResult, s"got $result")
+ }
+ }
+
+ test(s"the Filter takes a $kind reader's word on the rows it marks checked") {
+ // This reader breaks the contract on purpose: it applies no filter but marks every row
+ // checked. Only skipping the bloom on those rows keeps all of them.
+ withBloomJoin() {
+ val (result, plan) =
+ runBloomJoin("columnar" -> columnar.toString, "markAllChecked" -> "true")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.numRows,
+ s"the Filter must keep every row marked checked; got ${rowsOutOfFilter(plan)}")
+ assert(result == joinResult, s"the join must still filter; got $result")
+ }
+ }
+ }
+
+ test("a filter that cannot be applied is given up, and its rows come back unchecked") {
Review Comment:
Thanks, fixed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788). No runtime bloom throws, so no query reaches this path today. The test now drives both test readers, so the row reader's `catch` runs too, and the description says it covers the test readers. `runBloomJoin` checks the rows itself, and the two "nothing offered" checks share a helper. `FileFormat`'s name is the only one left, since item 16 removed `SupportsStorageFilters.CHECKED_COLUMN_NAME`.
##########
docs/sql-v2-data-sources.md:
##########
@@ -319,7 +319,10 @@ for data locality.
is serialized to executors and creates a
[`PartitionReader`](api/java/org/apache/spark/sql/connector/read/PartitionReader.html) for each
input partition. The reader iterates over rows (or `ColumnarBatch` instances for columnar sources)
-and is closed after consumption.
+and is closed after consumption. For a `SupportsStorageFilters` scan that Spark offered a storage
+filter, `createReader(partition, request)` or `createColumnarReader(partition, request)` creates the
+reader instead. The `StorageFilterRequest` carries the filter, and whether the reader appends a
+column that marks the rows it applied the filter to.
Review Comment:
Thanks, fixed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788) with your wording.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -127,6 +127,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
tableSpec.withNewLocation(newLoc)
}
+ /**
+ * The conjuncts of `filters` offered to the scan as storage filters, see
+ * `SupportsStorageFilters` and `StorageFilterPlanning.offeredConjuncts`. Only runtime blooms for
+ * now, as on the file-source side.
+ */
+ private def storageFiltersFor(
+ relation: DataSourceV2ScanRelation,
+ filters: Seq[Expression]): Seq[Expression] = {
+ if (!session.sessionState.conf.v2StorageFilterPushdownEnabled) return Nil
+ if (relation.storageFilterAttrs.isEmpty) return Nil
+ // A scan that took a pushed limit, offset, top-N or sample, or an aggregate or join, would have
+ // its readers filter ahead of it. `V2ScanRelationPushDown` leaves every such scan, and any scan
+ // a later pushdown builds, not mergeable.
+ if (!relation.mergeableScan) return Nil
+ // Over the whole table, as on the file-source side, since `SupportsStorageFilters` reserves
Review Comment:
Thanks, fixed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788), the second way. `BatchScanExec` passes its output without the checked column to `replanWithRuntimeFilters`. So the V2 path no longer reserves the name, and `SupportsStorageFilters.CHECKED_COLUMN_NAME` is gone. The name is internal to `FileFormat` again. A test reads a table with a column of that name through a `SupportsRuntimeV2Filtering` scan with iterative pushdown. With the widened output it fails on the ambiguous filter attribute.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/StorageFilterPlanning.scala:
##########
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, Expression, ExpressionSet, If, Literal}
+import org.apache.spark.sql.connector.read.SupportsStorageFilters
+import org.apache.spark.sql.execution.WholeStageCodegenExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.BooleanType
+
+/**
+ * The planning rules file-source and DSv2 scans share for storage filters, see
+ * `FileFormat.buildReaderWithStorageFilters` and `SupportsStorageFilters`. Each strategy decides
+ * on its own whether its scan can take any.
+ */
+object StorageFilterPlanning {
+
+ /** Whether one of `names` is the checked column's, in which case no filter is offered. */
+ def clashesWithCheckedColumn(names: Seq[String]): Boolean =
+ names.exists(_.equalsIgnoreCase(SupportsStorageFilters.CHECKED_COLUMN_NAME))
+
+ /**
+ * The conjuncts of `filters` offered to the scan, each once, as an `ExpressionSet` counts them.
+ * Each one:
+ * - is deterministic, since the post-scan `Filter` takes the reader's word on the rows it
+ * checked;
+ * - references at least one column, and only `keyCandidates`;
+ * - is taken by `accepts`.
+ *
+ * Nothing is offered when the offered conjuncts reference every one of `readColumns`, since the
+ * reader then reads the same columns for the same rows either way.
+ */
+ def offeredConjuncts(
+ filters: Seq[Expression],
+ keyCandidates: AttributeSet,
+ readColumns: AttributeSet,
+ accepts: Expression => Boolean): Seq[Expression] = {
+ val offered = ExpressionSet(filters).toSeq.filter { expr =>
+ val refs = expr.references
+ expr.deterministic && refs.nonEmpty && refs.subsetOf(keyCandidates) && accepts(expr)
+ }
+ if (offered.isEmpty) return Nil
+ val storageKeyAttrs = AttributeSet.fromAttributeSets(offered.map(_.references))
+ if (readColumns.subsetOf(storageKeyAttrs)) Nil else offered
+ }
+
+ /**
+ * The column the reader marks the rows it checked in, added after `output` when a filter is
Review Comment:
Thanks, fixed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788). The doc says the caller places the column, and names a conjunct that cannot be generated, with SPARK-60034's trade-off. The description's sentence is reworded the same way.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -127,6 +127,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
tableSpec.withNewLocation(newLoc)
}
+ /**
+ * The conjuncts of `filters` offered to the scan as storage filters, see
+ * `SupportsStorageFilters` and `StorageFilterPlanning.offeredConjuncts`. Only runtime blooms for
+ * now, as on the file-source side.
+ */
+ private def storageFiltersFor(
+ relation: DataSourceV2ScanRelation,
+ filters: Seq[Expression]): Seq[Expression] = {
+ if (!session.sessionState.conf.v2StorageFilterPushdownEnabled) return Nil
+ if (relation.storageFilterAttrs.isEmpty) return Nil
+ // A scan that took a pushed limit, offset, top-N or sample, or an aggregate or join, would have
+ // its readers filter ahead of it. `V2ScanRelationPushDown` leaves every such scan, and any scan
+ // a later pushdown builds, not mergeable.
+ if (!relation.mergeableScan) return Nil
Review Comment:
Thanks, changed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788) and [`cb01a6f`](https://github.com/apache/spark/pull/59715/commits/cb01a6f1c54dbd6e29153c6c9d6f1753a731a999). `DataSourceV2ScanRelation` has a new `onlyRowWisePushdowns` flag, false by default. The plain and the variant-extraction builds set it from `hasRowSelectingPushdown`, the one list of the pushdowns that decide which rows a scan returns. `hasBlockingPushdown` reads the same list, so a new pushdown listed there blocks merging as well. The storage-filter gate reads the flag, and `mergeableScan` means only scan merging again. A test plans the same relation with the two flags set apart, both ways.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala:
##########
@@ -29,22 +29,37 @@ import org.apache.spark.sql.connector.catalog.Table
import org.apache.spark.sql.connector.read._
import org.apache.spark.sql.connector.write.RowLevelOperation.Command.DELETE
import org.apache.spark.sql.connector.write.RowLevelOperationTable
-import org.apache.spark.sql.execution.EmptyRDDWithPartitions
+import org.apache.spark.sql.execution.{EmptyRDDWithPartitions, ScalarSubquery => ExecScalarSubquery}
import org.apache.spark.sql.execution.metric.{SQLLastAttemptMetrics, SQLMetric, SQLMetrics}
+import org.apache.spark.sql.internal.connector.{StorageFilterImpl, StorageFilterRequestImpl}
import org.apache.spark.util.ArrayImplicits._
/**
* Physical plan node for scanning a batch of data from a data source v2.
+ *
+ * @param storageFilters Predicates the scan's readers may evaluate, to skip reading the other
+ * columns of the rows they reject, see [[SupportsStorageFilters]]. Non-empty
+ * only for a scan that declared their columns. Each task's readers get them
+ * ANDed into one filter, in a [[StorageFilterRequest]]. They stay in the
+ * post-scan `Filter` too.
+ * @param checkedColumn The column the readers mark the rows they checked in, last in `output`, when
+ * the planner asked for one, see `StorageFilterPlanning.checkedColumn`.
*/
case class BatchScanExec(
output: Seq[AttributeReference],
@transient scan: Scan,
runtimeFilters: Seq[Expression],
ordering: Option[Seq[SortOrder]] = None,
@transient table: Table,
- keyGroupedPartitioning: Option[Seq[Expression]] = None
+ keyGroupedPartitioning: Option[Seq[Expression]] = None,
+ storageFilters: Seq[Expression] = Nil,
+ checkedColumn: Option[AttributeReference] = None
Review Comment:
Thanks, changed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788). `BatchScanExec` has `appendsCheckedColumn: Boolean` now. The `require` checks that the last output column is a boolean, and the extra normalization and the twin's remapping are gone.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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.internal.connector
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.{SparkEnv, SparkException, TaskContext}
+import org.apache.spark.executor.Executor
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.internal.config.KILL_ON_FATAL_ERROR_DEPTH
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to
+ * it. The binding is built on first use, so an instance created on the driver can be evaluated on
+ * an executor.
+ */
+class StorageFilterImpl private (
+ private val catalystExpr: CatalystExpression,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes.
+ private val numReferences = attrs.length
Review Comment:
Thanks, fixed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788). The messages interpolate `$this`, and the comment stays once, on `toString`. I kept `numReferences`, since `eval` reads it on every row and `attrs` can be a `List`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun closed pull request #59715: [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans
URL: https://github.com/apache/spark/pull/59715
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on PR #59715:
URL: https://github.com/apache/spark/pull/59715#issuecomment-6398087156
@dongjoon-hyun @szehon-ho @viirya @cloud-fan @aokolnychyi, this is now ready for review. It extends late materialization to data source V2 scans, the next step after the V1 Parquet support in #58895 (SPARK-59620). A DSv2 reader can then evaluate the runtime bloom filter inside the scan, and skip reading the other columns of the rows it drops.
Your feedback would be much appreciated. Since this adds to the DSv2 API, I'm happy to write an SPIP if you think it needs one. The API follows the shape of `PartitionPredicate` (SPARK-55596). It is an opaque predicate that Spark implements, and the connector only supplies the values it reads.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4386748262
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterPredicate.scala:
##########
@@ -0,0 +1,52 @@
+/*
+ * 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.internal.connector
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BasePredicate, BindReferences, Expression, Predicate}
+
+/**
+ * The predicate of a storage filter, built on first use from `boundExpression`, since a generated
+ * one can be awkward to serialize. The file-source Parquet reader and [[StorageFilterImpl]] share
+ * it.
+ */
+class StorageFilterPredicate(private var boundExpression: Expression) extends Serializable {
+
+ @transient lazy val prepared: BasePredicate = {
+ // The expression is released below, so a copy made after a first evaluation would have nothing
+ // to build from. That is named here rather than left to an NPE.
+ if (boundExpression == null) {
+ throw SparkException.internalError("A storage filter was copied after an evaluation " +
Review Comment:
Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). `StorageFilterPredicate` takes the filter's name for this message. That is the `toString` of `StorageFilterImpl`, and the key columns for the V1 filter. The message also says that only a filter that was never evaluated can be copied.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,103 @@
+/*
+ * 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.internal.connector
+
+import org.apache.spark.SparkException
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression => CatalystExpression}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the predicate is bound to
+ * it. The predicate is built on first use, see [[StorageFilterPredicate]].
+ */
+class StorageFilterImpl private (
+ private val predicate: StorageFilterPredicate,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes, read on every row.
+ private val numReferences = attrs.length
+
+ // Whether this instance has logged a failure, so a task logs only its first.
+ @transient private var loggedFailure = false
+
+ /**
+ * Wraps an error in [[StorageFilterNotAppliedException]], unless
+ * [[StorageFilterErrors.rethrowIfMustPropagate]] rethrows it, as the file-source reader does: a
+ * task kill, or a fatal error in its cause chain. Each task logs the first one. A row of another
+ * width than [[references]] is a reader bug, which fails the task, since the predicate would read
+ * the wrong values.
+ */
+ @throws[StorageFilterNotAppliedException]
+ override def eval(row: InternalRow): Boolean = {
+ // Outside the `try`, so a failure to build the predicate, or a row of another width, fails the
+ // task.
+ if (row.numFields != numReferences) {
+ throw SparkException.internalError(
+ s"$this takes rows of $numReferences values, but got ${row.numFields}")
+ }
+ val prepared = predicate.prepared
Review Comment:
Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). `eval` keeps the built predicate in a plain transient field.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,103 @@
+/*
+ * 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.internal.connector
+
+import org.apache.spark.SparkException
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression => CatalystExpression}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the predicate is bound to
+ * it. The predicate is built on first use, see [[StorageFilterPredicate]].
+ */
+class StorageFilterImpl private (
+ private val predicate: StorageFilterPredicate,
+ private val attrs: Seq[Attribute])
Review Comment:
Thanks, changed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). `StorageFilterImpl` keeps the references, and `StorageFilterPredicate.apply` is gone.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,560 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GetStructField, GreaterThan, If, LessThanOrEqual, Literal, NumericEvalContext, PredicateHelper, Rand, XxHash64}
+import org.apache.spark.sql.catalyst.plans.logical.ReturnAnswer
+import org.apache.spark.sql.execution.{BaseLimitExec, CoalesceExec, ExecSubqueryExpression, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation}
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{BinaryType, DoubleType, LongType, StructType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper with PredicateHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
+ withSQLConf(bloomJoinConf ++ conf: _*) {
+ withTable("bf_build") {
+ spark.range(0, 30).selectExpr("id AS bk", "id AS bv")
+ .write.format("parquet").saveAsTable("bf_build")
+ body
+ }
+ }
+ }
+
+ // A join whose build side is small and carries a selective predicate, which is the shape
+ // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`.
+ private def bloomJoin(probe: DataFrame): DataFrame = {
+ val build = spark.table("bf_build").where("bv = 5")
+ probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*)
+ }
+
+ private def v2Probe(options: (String, String)*): DataFrame =
+ spark.read.format(sourceName).options(options.toMap).load()
+
+ private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME
+
+ private def scanOf(plan: SparkPlan): BatchScanExec =
+ collect(plan) { case b: BatchScanExec => b }.headOption
+ .getOrElse(fail(s"no BatchScanExec in plan:\n$plan"))
+
+ // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows
+ // the reader marks as checked.
+ private def bloomsAbove(plan: SparkPlan): (Int, Int) = {
+ val conditions = collect(plan) { case f: FilterExec => f.condition }
+ val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum
+ val skipped = conditions.map(_.collect {
+ case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain)
+ if a.name == checkedName => 1
+ }.sum).sum
+ (all - skipped, skipped)
+ }
+
+ // The scan offered nothing, and the Filter evaluates the bloom on every row.
+ private def assertNothingOffered(plan: SparkPlan): Unit = {
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.isEmpty, s"the scan must get no storage filter:\n$plan")
+ assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan")
+ }
+
+ // The rows the Filter on the scan returned, counted before the join.
+ private def rowsOutOfFilter(plan: SparkPlan): Long = {
+ val filters = collect(plan) {
+ case f: FilterExec if f.exists(_.isInstanceOf[BatchScanExec]) => f
+ }
+ assert(filters.size == 1, s"expected one Filter on the scan:\n$plan")
+ filters.head.metrics("numOutputRows").value
+ }
+
+ // Runs `df`, checks that it returns `expected`, and returns its plan.
+ private def checkJoin(df: DataFrame, expected: Set[Seq[Long]]): SparkPlan = {
+ val result = df.collect().map(_.toSeq.map(_.asInstanceOf[Long])).toSet
+ assert(result == expected, s"got $result")
+ df.queryExecution.executedPlan
+ }
+
+ // The build side keeps `bv = 5`, so the join returns that one row.
+ private val joinResult = Set(Seq(5L, 50L))
+
+ // Runs the bloom join over the v2 source read with `options`, checks its rows, and returns its
+ // plan.
+ private def runBloomJoin(options: (String, String)*): SparkPlan = {
+ StorageFilterTestSource.reset()
+ checkJoin(bloomJoin(v2Probe(options: _*)), joinResult)
+ }
+
+ for (columnar <- Seq(false, true); aqe <- Seq(false, true)) {
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind scan gets the checked column, and the Filter skips the bloom on checked rows " +
+ s"(AQE $aqe)") {
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val plan = runBloomJoin("columnar" -> columnar.toString)
+ val scan = scanOf(plan)
+ assert(scan.supportsColumnar == columnar, s"the scan must read $kind:\n$plan")
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(scan.appendsCheckedColumn, s"the scan must get the checked column:\n$plan")
+ assert(bloomsAbove(plan) == (0, 1),
+ s"the Filter must skip the bloom on checked rows:\n$plan")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(),
+ s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " +
+ s"${StorageFilterTestSource.rowsReturned.get()}")
+ val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output }
+ assert(probeSide.exists(!_.exists(_.name == checkedName)),
+ s"the Project above the scan must drop the checked column; got $probeSide")
+ assert(StorageFilterTestSource.filteredReaders.get() == 1,
+ s"the one partition's reader must get the filter; " +
+ s"got ${StorageFilterTestSource.filteredReaders.get()}")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 1,
+ "the reader must be asked for the checked column")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ s"the reader must evaluate every row; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(false, true).foreach { columnar =>
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind reader that gives the filter up marks its rows unchecked, and the Filter " +
+ "filters them") {
+ withBloomJoin() {
+ val plan = runBloomJoin("columnar" -> columnar.toString, "giveUp" -> "true")
+ assert(bloomsAbove(plan) == (0, 1), s"the plan is that of a filtering reader:\n$plan")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == 0,
+ s"the reader must evaluate nothing; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() == StorageFilterTestSource.numRows,
+ s"the reader must return every row; got ${StorageFilterTestSource.rowsReturned.get()}")
+ assert(rowsOutOfFilter(plan) < StorageFilterTestSource.numRows,
+ s"the Filter must apply the bloom to unchecked rows; got ${rowsOutOfFilter(plan)}")
+ }
+ }
+
+ test(s"the Filter takes a $kind reader's word on the rows it marks checked") {
+ // This reader breaks the contract on purpose: it applies no filter but marks every row
+ // checked. Only skipping the bloom on those rows keeps all of them.
+ withBloomJoin() {
+ // The join must still filter, which `runBloomJoin` checks.
+ val plan = runBloomJoin("columnar" -> columnar.toString, "markAllChecked" -> "true")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.numRows,
+ s"the Filter must keep every row marked checked; got ${rowsOutOfFilter(plan)}")
+ }
+ }
+
+ Seq(false, true).foreach { aqe =>
+ test(s"a $kind reader that leaves out the checked column it was asked for fails the task, " +
+ s"from any row or batch (AQE $aqe)") {
+ // A reader that gives the filter up returns every row, so it returns several batches.
+ Seq(Seq("dropChecked" -> "true"), Seq("dropChecked" -> "later", "giveUp" -> "true"))
+ .foreach { options =>
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ StorageFilterTestSource.reset()
+ val df = bloomJoin(v2Probe(options :+ ("columnar" -> columnar.toString): _*))
+ // `k`, `v` and the checked column.
+ checkError(
+ exception = intercept[SparkException](df.collect()),
+ condition = "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING",
+ sqlState = Some("KD000"),
+ parameters = Map(
+ "class" -> classOf[StorageFilterTestReaderFactory].getName,
+ "expected" -> "3",
+ "actual" -> "2"))
+ }
+ }
+ }
+ }
+
+ test(s"a $kind reader may return more columns after the checked column") {
+ withBloomJoin() {
+ val plan = runBloomJoin("columnar" -> columnar.toString, "checkedLayout" -> "extra")
+ assert(bloomsAbove(plan) == (0, 1),
+ s"the Filter must skip the bloom on checked rows:\n$plan")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(),
+ s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " +
+ s"${StorageFilterTestSource.rowsReturned.get()}")
+ }
+ }
+
+ test(s"a $kind test reader gives a filter up when it throws, and returns its rows unchecked") {
+ // `1 / k` under ANSI throws at k = 0, the partition's first row, so the reader gives the
+ // filter up before it keeps or drops anything. No runtime bloom throws, so the reader is
+ // driven by hand.
+ val k = AttributeReference("k", LongType)()
+ val inverse = Divide(Literal(1.0), Cast(k, DoubleType), NumericEvalContext(EvalMode.ANSI))
+ val filter = StorageFilterImpl(GreaterThan(inverse, Literal(0.0)), Seq(k))
+ val partition = RangeInputPartition(0, StorageFilterTestSource.numRows)
+ val schema = new StructType().add("k", LongType).add("v", LongType)
+ val rows = mutable.ArrayBuffer.empty[(Long, Boolean)]
+ if (columnar) {
+ val reader = new StorageFilterTestColumnarReader(
+ partition, schema, Some(filter), appendsChecked = true)
+ try {
+ while (reader.next()) {
+ val batch = reader.get()
+ (0 until batch.numRows()).foreach { i =>
+ rows += ((batch.column(0).getLong(i), batch.column(2).getBoolean(i)))
+ }
+ }
+ } finally {
+ reader.close()
+ }
+ } else {
+ val reader =
+ new StorageFilterTestReader(partition, schema, Some(filter), appendsChecked = true)
+ while (reader.next()) {
+ rows += ((reader.get().getLong(0), reader.get().getBoolean(2)))
+ }
+ }
+ assert(rows == (0L until StorageFilterTestSource.numRows.toLong).map(_ -> false),
+ s"expected every row, unchecked; got ${rows.take(5)} and ${rows.size} rows in all")
+ }
+ }
+
+ test("a columnar reader that returns the checked column ahead of the others fails the task") {
+ // The batch is as wide as asked for, but its last column is not a boolean. A row carries no
+ // types, so a row-based reader's mistake goes unnoticed.
+ withBloomJoin() {
+ StorageFilterTestSource.reset()
+ val df = bloomJoin(v2Probe("columnar" -> "true", "checkedLayout" -> "front"))
+ checkError(
+ exception = intercept[SparkException](df.collect()),
+ condition = "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING",
+ sqlState = Some("KD000"),
+ parameters = Map(
+ "class" -> classOf[StorageFilterTestReaderFactory].getName,
+ "expected" -> "3",
+ "actual" -> "3"))
+ }
+ }
+
+ // Under a limit of two fields, the checked column would make three.
+ Seq(
+ ("too wide for whole-stage codegen", SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> "2"),
+ ("without whole-stage codegen", SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false")
+ ).foreach { case (when, conf) =>
+ test(s"a scan $when gets the bloom without the checked column") {
+ withBloomJoin(conf) {
+ val plan = runBloomJoin("columnar" -> "true")
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(!scan.appendsCheckedColumn, s"the scan must not get the checked column:\n$plan")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 0,
+ "the reader must not be asked for the checked column")
+ assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(
+ ("with the conf off", Seq(SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "false"), Nil),
+ ("to a scan that declares no column", Nil, Seq("declared" -> "")),
+ ("to a scan that declares only another column", Nil, Seq("declared" -> "v")),
+ ("to a scan that declares only a nested field", Nil, Seq("declared" -> "k.x"))
+ ).foreach { case (when, conf, options) =>
+ test(s"nothing is offered $when") {
+ withBloomJoin(conf: _*) {
+ val plan = runBloomJoin(options :+ ("columnar" -> "true"): _*)
+ assertNothingOffered(plan)
+ assert(StorageFilterTestSource.filteredReaders.get() == 0,
+ "no reader may get a storage filter")
+ }
+ }
+ }
+
+ test("a scan that reads only the key and a metadata column is offered nothing") {
+ // Only the data columns count as columns left to skip, as on the file-source side.
+ withBloomJoin() {
+ val source = v2Probe("columnar" -> "true", "metadataColumn" -> "true")
+ val pos = StorageFilterTestSource.posName
+ StorageFilterTestSource.reset()
+ assertNothingOffered(
+ checkJoin(bloomJoin(source.select(source("k"), source.metadataColumn(pos))),
+ Set(Seq(5L, 5L))))
+ StorageFilterTestSource.reset()
+ val withValue = checkJoin(
+ bloomJoin(source.select(source("k"), source("v"), source.metadataColumn(pos))),
+ Set(Seq(5L, 50L, 5L)))
+ assert(scanOf(withValue).storageFilters.size == 1,
+ s"with `v` left to skip, the bloom must be offered:\n$withValue")
+ }
+ }
+
+ test("a variant scan that reads only its key column and a metadata column is offered nothing") {
+ // The key is an extraction of the declared variant column, so that column is the scan's only
+ // data column, and the scan's output has fresh exprIds.
+ withBloomJoin() {
+ val pos = StorageFilterTestSource.posName
+ val probe = v2Probe("columnar" -> "true", "variant" -> "true", "metadataColumn" -> "true",
+ "declared" -> "j").selectExpr("variant_get(j, '$.a', 'long') AS a", pos)
+ val build = spark.table("bf_build").where("bv = 5")
+ val df = probe.join(build, probe("a") === build("bk") * 10).select(probe("a"), probe(pos))
+ StorageFilterTestSource.reset()
+ assertNothingOffered(checkJoin(df, Set(Seq(50L, 5L))))
+ }
+ }
+
+ test("a scan that took a pushed variant extraction gets the bloom") {
+ // The extraction keeps the scan's rows.
+ withBloomJoin() {
+ StorageFilterTestSource.reset()
+ val probe = v2Probe("columnar" -> "true", "variant" -> "true")
+ .selectExpr("k", "variant_get(j, '$.a', 'long') AS a")
+ val plan = checkJoin(bloomJoin(probe), joinResult)
+ val scan = scanOf(plan)
+ assert(scan.output.exists(_.dataType.isInstanceOf[StructType]),
+ s"test setup: the scan must read the extraction:\n$plan")
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+
+ for (aqe <- Seq(false, true)) {
+ test("a scan with a column named like the checked column gets the filter, and runtime " +
+ s"filtering still resolves that column (AQE $aqe)") {
+ // Runtime filtering resolves the declared filter attributes by name, which the checked column
+ // the scan appends would make ambiguous. Dynamic partition pruning on `k` would take the
+ // bloom's place.
+ withBloomJoin(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "false",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val plan = runBloomJoin(
+ "columnar" -> "true", "valueName" -> checkedName, "runtimeFiltering" -> "v2")
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.size == 1 && scan.appendsCheckedColumn,
+ s"the scan must get the bloom and the checked column:\n$plan")
+ assert(StorageFilterTestSource.runtimeFilterCalls.get() > 0,
+ "test setup: the bloom must reach runtime filtering")
+ val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output }
+ assert(probeSide.contains(scan.output.init),
+ s"the Project above the scan must keep the table's column; got $probeSide")
+ }
+ }
+
+ test(s"a bloom the scan fully evaluates through runtime filtering is not offered (AQE $aqe)") {
+ // Dynamic partition pruning on `k` would take the bloom's place.
+ withBloomJoin(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "false",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val plan = runBloomJoin("columnar" -> "true", "runtimeFiltering" -> "catalyst")
+ val scan = scanOf(plan)
+ assert(scan.runtimeFilters.exists(_.exists(_.isInstanceOf[BloomFilterMightContain])) &&
+ StorageFilterTestSource.runtimeFilterCalls.get() > 0,
+ s"test setup: the scan must take the bloom as a runtime filter:\n$plan")
+ assert(scan.storageFilters.isEmpty, s"the scan must get no storage filter:\n$plan")
+ assert(bloomsAbove(plan) == (0, 0), s"no Filter may keep the bloom:\n$plan")
+ }
+ }
+
+ test(s"a two-key join offers both blooms in one filter, each bloom with its own bytes " +
+ s"(AQE $aqe)") {
+ // `MergeSubplans` merges the two blooms' aggregates into one subquery returning both blooms.
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ withTable("bf_build2") {
+ spark.range(0, 30).selectExpr("id AS bk", "id * 10 AS bv", "id AS flag")
+ .write.format("parquet").saveAsTable("bf_build2")
+ StorageFilterTestSource.reset()
+ val probe = v2Probe("columnar" -> "true", "wide" -> "true", "declared" -> "k,v")
+ .where("w % 3 != 1")
+ val build = spark.table("bf_build2").where("flag = 5")
+ val df = probe.join(build, probe("k") === build("bk") && probe("v") === build("bv"))
+ .select(probe("k"), probe("v"), probe("w"))
+ val plan = checkJoin(df, Set(Seq(5L, 50L, 500L)))
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.size == 2, s"both blooms must be on the scan:\n$plan")
+ assert(scan.storageFilters.forall(_.exists {
+ case GetStructField(_: ExecSubqueryExpression, _, _) => true
+ case _ => false
+ }), s"test setup: the blooms must come from one merged subquery:\n$plan")
+ assert(bloomsAbove(plan) == (0, 2), s"the Filter must skip both blooms:\n$plan")
+ val filter = collect(plan) {
+ case f: FilterExec if f.child.exists(_.isInstanceOf[BatchScanExec]) => f
+ }.head
+ assert(splitConjunctivePredicates(filter.condition).exists { c =>
+ c.references.map(_.name).toSet == Set("w") && !c.isInstanceOf[If]
+ }, s"the Filter must evaluate the conjunct over `w` on every row:\n$plan")
+ assert(StorageFilterTestSource.filterReferences == Set(Seq("k", "v")),
+ s"one filter over both keys; got ${StorageFilterTestSource.filterReferences}")
+ val bloomBytes = scan.storageFilters.map(_.collectFirst {
+ case m: BloomFilterMightContain =>
+ m.bloomFilterExpression.eval().asInstanceOf[Array[Byte]].length
+ }.get).sum
+ assert(StorageFilterTestSource.filterBytes.get() < bloomBytes * 3 / 2,
+ s"each bloom must ship once; the filter takes ${StorageFilterTestSource.filterBytes} " +
+ s"bytes for $bloomBytes bytes of blooms")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ "the reader must evaluate every row; got " +
+ StorageFilterTestSource.rowsEvaluated.get())
+ }
+ }
+ }
+
+ test(s"the readers of one task share the filter the task deserialized (AQE $aqe)") {
+ // `coalesce(1)` reads both partitions in one task.
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ StorageFilterTestSource.reset()
+ val probe = v2Probe("columnar" -> "true", "numPartitions" -> "2").coalesce(1)
+ val plan = checkJoin(bloomJoin(probe), joinResult)
+ assert(collect(plan) { case c: CoalesceExec => c }.nonEmpty, s"test setup:\n$plan")
+ assert(scanOf(plan).storageFilters.size == 1, s"the bloom must be on the scan:\n$plan")
+ assert(StorageFilterTestSource.filteredReaders.get() == 2,
+ s"both readers must get the filter; got ${StorageFilterTestSource.filteredReaders.get()}")
+ assert(StorageFilterTestSource.filterIdentities.size == 1,
+ "both readers must get the same instance")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ "the readers must evaluate every row; got " + StorageFilterTestSource.rowsEvaluated.get())
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the readers must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ test("the storage filter shows in the scan description and separates canonicalized plans") {
+ withBloomJoin() {
+ val plan = bloomJoin(v2Probe("columnar" -> "true")).queryExecution.executedPlan
+ val scan = scanOf(plan)
+ assert(scan.simpleString(100).contains("StorageFilters:"),
+ s"the scan description must name the filters: ${scan.simpleString(100)}")
+ assert(scan.verboseStringWithOperatorId().contains("StorageFilters:"),
+ s"EXPLAIN FORMATTED must name the filters: ${scan.verboseStringWithOperatorId()}")
+
+ intercept[IllegalArgumentException](scan.copy(storageFilters = Nil))
+ // As the scan would be without the checked column, too wide or without codegen.
+ val unchecked = scan.copy(output = scan.output.init, appendsCheckedColumn = false)
+ assert(scan != unchecked, "equals must see the checked column")
+ val plain = unchecked.copy(storageFilters = Nil)
+ assert(unchecked != plain && unchecked.hashCode != plain.hashCode,
+ "equals and hashCode must see storageFilters")
+ assert(!unchecked.sameResult(plain),
+ s"a filtering scan must not be reusable as a plain one:\n${unchecked.canonicalized}\n" +
+ s"${plain.canonicalized}")
+ assert(!plain.simpleString(100).contains("StorageFilters:"),
+ s"a plain scan's description must not name storage filters: ${plain.simpleString(100)}")
+ assert(!plain.verboseStringWithOperatorId().contains("StorageFilters:"),
+ s"nor its EXPLAIN FORMATTED: ${plain.verboseStringWithOperatorId()}")
+
+ // As another relation instance of the same table gets, with fresh exprIds.
+ val fresh = AttributeMap(scan.output.map(a => a -> a.newInstance()))
+ val twin = scan.copy(
+ output = scan.output.map(fresh(_)),
+ storageFilters = scan.storageFilters.map(_.transform {
+ case a: Attribute => fresh.getOrElse(a, a)
+ }))
+ assert(scan.sameResult(twin),
+ s"the same filtering scan must be reusable:\n${scan.canonicalized}\n${twin.canonicalized}")
+ }
+ }
+
+ // The reader would filter ahead of the operation, which then picks other rows.
+ Seq[(String, DataFrame => DataFrame)](
+ "limit" -> (_.limit(100)),
+ "offset" -> (_.offset(10)),
+ "topN" -> (_.orderBy("v").limit(100)),
+ "sample" -> (_.sample(0.5))
+ ).foreach { case (op, apply) =>
+ test(s"nothing is offered to a scan that took a pushed $op") {
+ withBloomJoin() {
+ val probe = apply(v2Probe("columnar" -> "true", "push" -> op))
+ val plan = bloomJoin(probe).queryExecution.executedPlan
+ assert(collect(plan) {
+ case p @ (_: BaseLimitExec | _: TakeOrderedAndProjectExec | _: SampleExec) => p
+ }.isEmpty, s"the $op must be pushed into the scan whole:\n$plan")
+ assertNothingOffered(plan)
+ }
+ }
+ }
+
+ test("the scan's pushdowns decide on storage filters, not whether it can be merged") {
+ // A pushed variant extraction, for one, keeps the scan's rows but cannot be merged.
+ withBloomJoin() {
+ val optimized = bloomJoin(v2Probe("columnar" -> "true")).queryExecution.optimizedPlan
+ def planned(mergeable: Boolean, rowWise: Boolean): BatchScanExec = scanOf(
+ spark.sessionState.planner.plan(ReturnAnswer(optimized.transform {
Review Comment:
Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb).
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterTestSource.scala:
##########
@@ -0,0 +1,492 @@
+/*
+ * 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.util
+import java.util.OptionalLong
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow}
+import org.apache.spark.sql.connector.catalog.{MetadataColumn, SupportsMetadataColumns, SupportsRead, Table, TableCapability, TableProvider}
+import org.apache.spark.sql.connector.expressions.{Expressions, FieldReference, NamedReference, SortOrder, Transform}
+import org.apache.spark.sql.connector.expressions.filter.Predicate
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.sql.types.{BooleanType, DataType, LongType, StructField, StructType, VariantType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.util.Utils
+
+/**
+ * A read-only v2 source over a fixed table of longs that implements [[SupportsStorageFilters]]:
+ * the key `k`, the value `v` holding `k * 10`, and with `wide` also `w` holding `k * 100`. Its
+ * reader gets the filter in the task and drops the rows the filter rejects. It reads the columns
+ * Spark requires, and builds the filter's row from the filter's `references()`. Its options:
+ * - `columnar`: reads in batches, through `StorageFilter.select`, rather than row by row;
+ * - `giveUp`: reads without the filter, as a reader does for a unit it cannot use the filter on.
+ * Here the unit is the partition;
+ * - `markAllChecked`: reads without the filter but marks every row checked, which breaks the
+ * contract on purpose;
+ * - `dropChecked`: `true` never appends the checked column, and `later` appends it to the first
+ * row or batch only. Either breaks the contract on purpose;
+ * - `checkedLayout`: `front` returns the checked column ahead of the scan's columns, which breaks
+ * the contract on purpose. `extra` returns one more column after it;
+ * - `declared`: the comma-separated storage-filter columns, `k` by default, a dot separating a
+ * nested field;
+ * - `valueName`: the name of the value column, `v` by default;
+ * - `wide`: adds the column `w`;
+ * - `variant`: adds the variant column `j`, which the scan reads only through pushed extractions
+ * of a long. Each one holds `k * 10`;
+ * - `numPartitions`: splits the keys into that many input partitions, 1 by default;
+ * - `metadataColumn`: adds the metadata column `_pos`, which holds the key;
+ * - `runtimeFiltering`: `v2` makes the scan a `SupportsRuntimeV2Filtering` one over every column,
+ * with iterative pushdown on a table partitioned by `k`. `catalyst` makes it a
+ * `SupportsRuntimeCatalystFiltering` one that evaluates the filters on `k` fully. Neither
+ * applies the filters it takes, so they serve planning tests only;
+ * - `push`: the comma-separated pushed operations it takes whole without applying them, of
+ * `limit`, `offset`, `topN` and `sample`, for planning tests only.
+ *
+ * A filter that throws `StorageFilterNotAppliedException` is given up for the rest of the
+ * partition too. The counters below record what the readers got, evaluated and returned.
+ */
+class StorageFilterTestSource extends TableProvider {
+
+ override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+ val schema = new StructType().add("k", LongType).add(options.getOrDefault("valueName", "v"),
+ LongType)
+ val wide = if (options.getBoolean("wide", false)) schema.add("w", LongType) else schema
+ if (options.getBoolean("variant", false)) wide.add("j", VariantType) else wide
+ }
+
+ override def getTable(
+ schema: StructType,
+ partitioning: Array[Transform],
+ properties: util.Map[String, String]): Table = {
+ val options = new CaseInsensitiveStringMap(properties)
+ if (options.getBoolean("metadataColumn", false)) {
+ new StorageFilterTestTable(schema, options) with SupportsMetadataColumns {
+ override def metadataColumns(): Array[MetadataColumn] = Array(new MetadataColumn {
+ override def name(): String = StorageFilterTestSource.posName
+ override def dataType(): DataType = LongType
+ })
+ }
+ } else {
+ new StorageFilterTestTable(schema, options)
+ }
+ }
+}
+
+object StorageFilterTestSource {
+
+ /** Rows the source serves, `k` from 0 until numRows. */
+ val numRows: Int = 200
+
+ /** Rows a batch holds at most in a columnar read. */
+ val batchSize: Int = 64
+
+ /** The name of the metadata column the `metadataColumn` option adds. */
+ val posName: String = "_pos"
+
+ /** How many readers were created with a storage-filter request. */
+ val filteredReaders = new AtomicInteger(0)
+
+ /** How many of those requests asked for the checked column. */
+ val checkedColumnRequests = new AtomicInteger(0)
+
+ /** Rows the reader evaluated the storage filter against. */
+ val rowsEvaluated = new AtomicLong(0)
+
+ /** Rows the reader returned. */
+ val rowsReturned = new AtomicLong(0)
+
+ /** The serialized size of the first filter a reader got, before it evaluated the filter. */
+ val filterBytes = new AtomicLong(0)
+
+ /** The references of the filters the readers got, and the identities of those filters. */
+ val filterReferences = mutable.Set.empty[Seq[String]]
+ val filterIdentities = mutable.Set.empty[Int]
+
+ /** How many times a runtime-filtering scan was handed filters. */
+ val runtimeFilterCalls = new AtomicInteger(0)
+
+ def reset(): Unit = {
+ filteredReaders.set(0)
+ checkedColumnRequests.set(0)
+ rowsEvaluated.set(0)
+ rowsReturned.set(0)
+ filterBytes.set(0)
+ filterReferences.synchronized(filterReferences.clear())
+ filterIdentities.synchronized(filterIdentities.clear())
+ runtimeFilterCalls.set(0)
+ }
+
+ /** The value of column `name` in the row of `key`. Any column but `k`, `w` and `_pos` is `v`. */
+ def valueOf(name: String, key: Long): Long = name match {
+ case "k" | `posName` => key
+ case "w" => key * 100
+ case _ => key * 10
+ }
+
+ /** The value of `field` in the row of `key`, a struct for the extractions of a variant. */
+ def valueOf(field: StructField, key: Long): Any = field.dataType match {
+ case struct: StructType => InternalRow.fromSeq(struct.map(_ => valueOf(field.name, key)))
+ case _ => valueOf(field.name, key)
+ }
+
+ /** The columns of the row the filter evaluates, in that order. */
+ def columnsOf(filter: StorageFilter): Seq[String] = filter.references.map(_.fieldNames.head).toSeq
+
+ // Records the filter a reader got, before the reader evaluates it. Only the first reader
+ // serializes it, since a later reader of the same task may get it evaluated, see `StorageFilter`.
+ private[connector] def record(filter: StorageFilter): Unit = {
+ if (filterBytes.get() == 0) filterBytes.compareAndSet(0, Utils.serialize(filter).length)
+ filterReferences.synchronized(filterReferences += columnsOf(filter))
+ filterIdentities.synchronized(filterIdentities += System.identityHashCode(filter))
+ }
+}
+
+class StorageFilterTestTable(tableSchema: StructType, tableOptions: CaseInsensitiveStringMap)
+ extends Table with SupportsRead {
+
+ override def name(): String = "storage_filter_test"
+
+ override def schema(): StructType = tableSchema
+
+ override def partitioning(): Array[Transform] =
+ if (tableOptions.getOrDefault("runtimeFiltering", "") == "v2") {
+ Array(Expressions.identity("k"))
+ } else {
+ Array.empty
+ }
+
+ override def capabilities(): util.Set[TableCapability] =
+ Set(TableCapability.BATCH_READ).asJava
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder =
+ new StorageFilterTestScanBuilder(tableSchema, options)
+}
+
+/**
+ * Prunes the columns Spark does not read, and takes the pushed operations named in `push` whole.
+ * Its scan does not apply them, so they serve planning tests only. It takes every variant
+ * extraction, and reads each into a struct field of the variant's column named by its position.
+ */
+class StorageFilterTestScanBuilder(tableSchema: StructType, options: CaseInsensitiveStringMap)
+ extends SupportsPushDownRequiredColumns
+ with SupportsPushDownLimit
+ with SupportsPushDownOffset
+ with SupportsPushDownTopN
+ with SupportsPushDownTableSample
+ with SupportsPushDownVariantExtractions {
+
+ private val pushes = Utils.stringToSeq(options.getOrDefault("push", "")).toSet
+ private var required = tableSchema
+ private var extractions = Array.empty[VariantExtraction]
+
+ override def pruneColumns(requiredSchema: StructType): Unit = required = requiredSchema
+
+ override def pushVariantExtractions(pushed: Array[VariantExtraction]): Array[Boolean] = {
+ extractions = pushed
+ pushed.map(_ => true)
+ }
+
+ override def build(): Scan = {
+ val byColumn = extractions.groupBy(_.columnName().head)
+ val schema = StructType(required.map { field =>
+ byColumn.get(field.name).map { columnExtractions =>
+ field.copy(dataType = StructType(columnExtractions.zipWithIndex.map { case (e, i) =>
+ StructField(i.toString, e.expectedDataType(), nullable = true, e.metadata())
+ }))
+ }.getOrElse(field)
+ })
+ val factory = StorageFilterTestReaderFactory(
+ schema,
+ columnar = options.getBoolean("columnar", false),
+ giveUp = options.getBoolean("giveUp", false),
+ marksAllChecked = options.getBoolean("markAllChecked", false),
+ dropChecked = options.getOrDefault("dropChecked", ""),
+ checkedLayout = options.getOrDefault("checkedLayout", ""))
+ val declared = Utils.stringToSeq(options.getOrDefault("declared", "k"))
+ val numPartitions = options.getInt("numPartitions", 1)
+ options.getOrDefault("runtimeFiltering", "") match {
+ case "v2" => new StorageFilterTestScan(schema, factory, declared, numPartitions)
+ with SupportsRuntimeV2Filtering {
+ override def filterAttributes(): Array[NamedReference] =
+ readSchema().fieldNames.map(FieldReference.column).toArray[NamedReference]
+ override def filter(predicates: Array[Predicate]): Unit =
+ StorageFilterTestSource.runtimeFilterCalls.incrementAndGet()
+ override def supportsIterativePushdown(): Boolean = true
+ }
+ case "catalyst" => new StorageFilterTestScan(schema, factory, declared, numPartitions)
+ with SupportsRuntimeCatalystFiltering {
+ override def filterAttributes(): Array[NamedReference] = Array(FieldReference.column("k"))
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+ override def filter(expressions: Array[Expression]): Unit =
+ StorageFilterTestSource.runtimeFilterCalls.incrementAndGet()
+ }
+ case _ => new StorageFilterTestScan(schema, factory, declared, numPartitions)
+ }
+ }
+
+ override def pushLimit(limit: Int): Boolean = pushes("limit")
+ override def pushOffset(offset: Int): Boolean = pushes("offset")
+ override def pushTopN(orders: Array[SortOrder], limit: Int): Boolean = pushes("topN")
+ override def pushTableSample(
+ lowerBound: Double,
+ upperBound: Double,
+ withReplacement: Boolean,
+ seed: Long,
+ sampleMethod: SampleMethod): Boolean = pushes("sample")
+ override def isPartiallyPushed: Boolean = false
+}
+
+class StorageFilterTestScan(
+ schema: StructType,
+ factory: StorageFilterTestReaderFactory,
+ declared: Seq[String],
+ numPartitions: Int)
+ extends Scan with Batch with SupportsReportStatistics with SupportsStorageFilters {
+
+ override def readSchema(): StructType = schema
+
+ override def toBatch: Batch = this
+
+ // Over the suite's application-side scan size threshold, so InjectRuntimeFilter considers this
+ // side worth filtering.
+ override def estimateStatistics(): Statistics = new Statistics {
+ override def sizeInBytes(): OptionalLong = OptionalLong.of(10L * 1024 * 1024)
+ override def numRows(): OptionalLong = OptionalLong.of(StorageFilterTestSource.numRows.toLong)
+ }
+
+ override def storageFilterAttributes(): Array[NamedReference] =
+ declared.map(name => FieldReference(name.split('.').toSeq)).toArray[NamedReference]
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ val n = StorageFilterTestSource.numRows
+ (0 until numPartitions).map { i =>
+ RangeInputPartition(i * n / numPartitions, (i + 1) * n / numPartitions)
+ }.toArray
+ }
+
+ override def createReaderFactory(): PartitionReaderFactory = factory
+}
+
+case class StorageFilterTestReaderFactory(
+ schema: StructType,
+ columnar: Boolean,
+ giveUp: Boolean,
+ marksAllChecked: Boolean,
+ dropChecked: String,
+ checkedLayout: String)
+ extends PartitionReaderFactory {
+
+ override def supportColumnarReads(partition: InputPartition): Boolean = columnar
+
+ // The filter the reader applies, none when it gives the filter up or only marks the rows.
+ private def filterOf(request: StorageFilterRequest): Option[StorageFilter] = {
+ StorageFilterTestSource.filteredReaders.incrementAndGet()
+ StorageFilterTestSource.record(request.filter())
+ if (request.appendsCheckedColumn()) {
+ StorageFilterTestSource.checkedColumnRequests.incrementAndGet()
+ }
+ if (giveUp || marksAllChecked) None else Some(request.filter())
+ }
+
+ override def createReader(partition: InputPartition): PartitionReader[InternalRow] =
+ new StorageFilterTestReader(keys(partition), schema, None, appendsChecked = false)
+
+ override def createReader(
+ partition: InputPartition,
+ request: StorageFilterRequest): PartitionReader[InternalRow] =
+ new StorageFilterTestReader(keys(partition), schema, filterOf(request),
+ request.appendsCheckedColumn() && dropChecked != "true", marksAllChecked,
+ dropsCheckedLater = dropChecked == "later", checkedLayout)
+
+ override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] =
+ new StorageFilterTestColumnarReader(keys(partition), schema, None, appendsChecked = false)
+
+ override def createColumnarReader(
+ partition: InputPartition,
+ request: StorageFilterRequest): PartitionReader[ColumnarBatch] =
+ new StorageFilterTestColumnarReader(keys(partition), schema, filterOf(request),
+ request.appendsCheckedColumn() && dropChecked != "true", marksAllChecked,
+ dropsCheckedLater = dropChecked == "later", checkedLayout)
+
+ // The keys from `start` until `end`.
+ private def keys(partition: InputPartition): RangeInputPartition =
+ partition.asInstanceOf[RangeInputPartition]
+}
+
+/**
+ * Serves the rows of `partition` row by row, holding the columns of `schema`, and drops any row
+ * `filter` rejects, with the checked column last when `appendsChecked`. The filter's row holds
+ * the values of its `references()`. `filter` is cleared when it cannot be applied.
+ * `marksAllChecked` marks every row checked. `dropsCheckedLater` appends the checked column to
+ * the first row only. `checkedLayout` places the checked column as the source's option says.
+ */
+class StorageFilterTestReader(
+ partition: RangeInputPartition,
+ schema: StructType,
+ private var filter: Option[StorageFilter],
+ appendsChecked: Boolean,
+ marksAllChecked: Boolean = false,
+ dropsCheckedLater: Boolean = false,
+ checkedLayout: String = "")
+ extends PartitionReader[InternalRow] {
+
+ private val referenceNames = filter.map(StorageFilterTestSource.columnsOf).getOrElse(Nil)
+ private var current: Long = partition.start - 1L
+ private var row: InternalRow = _
+ private var returnedAny = false
+
+ override def next(): Boolean = {
+ var found = false
+ while (!found && current + 1 < partition.end) {
+ current += 1
+ val key = current
+ val passes = filter.forall { f =>
+ StorageFilterTestSource.rowsEvaluated.incrementAndGet()
+ val keyRow = new GenericInternalRow(
+ referenceNames.map(StorageFilterTestSource.valueOf(_, key)).toArray[Any])
+ try f.eval(keyRow) catch {
+ case _: StorageFilterNotAppliedException =>
+ filter = None
+ true
+ }
+ }
+ if (passes) {
+ val values = schema.fields.map(StorageFilterTestSource.valueOf(_, key))
+ val checked = filter.isDefined || marksAllChecked
+ val withChecked = appendsChecked && !(dropsCheckedLater && returnedAny)
+ row = new GenericInternalRow(
+ if (!withChecked) {
+ values
+ } else {
+ checkedLayout match {
+ case "front" => checked +: values
Review Comment:
Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). The docs say `front` is for columnar reads, and the row reader rejects it.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsStorageFilters.java:
##########
@@ -0,0 +1,64 @@
+/*
+ * 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.read;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link Scan}. A data source implements it if its readers can evaluate a
+ * {@link StorageFilter}, to skip reading the other columns of the rows the filter rejects. The
+ * readers get the filter in the task, through
+ * {@link PartitionReaderFactory#createReader(InputPartition, StorageFilterRequest)} or
+ * {@link PartitionReaderFactory#createColumnarReader(InputPartition, StorageFilterRequest)}, which
+ * the scan's reader factory must implement.
Review Comment:
Thanks, the description's possible follow-ups now name the V2 file sources, with the three gaps you list.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "spark-test (via GitHub)" <gi...@apache.org>.
spark-test commented on PR #59715:
URL: https://github.com/apache/spark/pull/59715#issuecomment-6350384781
The AI review requested by @HyukjinKwon did not finish. A committer can add the `ai-review-requested` label again to retry.
<!-- SPARK_OSS_REVIEW_REQUEST 33986584279 -->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4381899032
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -204,6 +221,13 @@ class DataSourceRDD(
}
}
+ private def checkWidth(expected: Int, actual: Int): Unit = {
+ if (actual != expected) {
Review Comment:
Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). A batch now needs at least the expected number of columns, with a boolean at the checked column's position. A row needs at least that many fields. So a reader may return more columns after the checked one. The message says that the column is missing or not a boolean. This stays a best-effort guard for a reader that breaks the documented layout, which I'd expect to be rare. A row carries no names or types, and a misplaced boolean column still passes. New tests cover a columnar reader that returns the column ahead of the others, and readers that return one more column after it.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,491 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GetStructField, GreaterThan, If, Literal, NumericEvalContext, PredicateHelper, Rand, XxHash64}
+import org.apache.spark.sql.catalyst.plans.logical.ReturnAnswer
+import org.apache.spark.sql.execution.{BaseLimitExec, CoalesceExec, ExecSubqueryExpression, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation}
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{BinaryType, DoubleType, LongType, StructType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper with PredicateHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
Review Comment:
Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). The two-key, two-reader and runtime-filtering tests, and the test of a reader that leaves out the checked column, now run with AQE on and off.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala:
##########
@@ -817,8 +817,8 @@ case class FileSourceScanExec(
protected lazy val preparedStorageFilters: Seq[Expression] = {
val requestedDataAttrs = output.take(requiredSchema.length)
storageFilters.map { expr =>
- val subqueryReplaced = expr.transform { case s: execution.ScalarSubquery => s.toLiteral }
- BindReferences.bindReference(subqueryReplaced, requestedDataAttrs)
+ BindReferences.bindReference(
+ StorageFilterPlanning.withSubqueryResults(expr), requestedDataAttrs)
Review Comment:
Thanks, added in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). `ParquetStorageFilterSuite` has a two-key join whose blooms come from one merged subquery, with AQE on and off. It checks that the scan's RDD, which carries the reader's function, takes less than 1.5 times the bytes of the two blooms. With `s.toLiteral` alone it takes about twice the bytes of the two blooms, and the test fails.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/StorageFilterPlanning.scala:
##########
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, BloomFilterMightContain, Expression, ExpressionSet, If, Literal}
+import org.apache.spark.sql.execution.{ScalarSubquery => ExecScalarSubquery, WholeStageCodegenExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, BooleanType}
+
+/**
+ * The rules file-source and DSv2 scans share for storage filters, see
+ * `FileFormat.buildReaderWithStorageFilters` and `SupportsStorageFilters`. Each strategy decides
+ * on its own whether its scan can take any.
+ */
+object StorageFilterPlanning {
+
+ /**
+ * The conjuncts of `filters` offered to the scan, each once, as an `ExpressionSet` counts them.
+ * Each one:
+ * - is deterministic, since the post-scan `Filter` takes the reader's word on the rows it
+ * checked;
+ * - references at least one column, and only `keyCandidates`;
+ * - is taken by `accepts`.
+ *
+ * Nothing is offered when the offered conjuncts reference every one of `readColumns`, since the
+ * reader then reads the same columns for the same rows either way.
+ */
+ def offeredConjuncts(
+ filters: Seq[Expression],
+ keyCandidates: AttributeSet,
+ readColumns: AttributeSet,
+ accepts: Expression => Boolean): Seq[Expression] = {
+ val offered = ExpressionSet(filters).toSeq.filter { expr =>
+ val refs = expr.references
+ expr.deterministic && refs.nonEmpty && refs.subsetOf(keyCandidates) && accepts(expr)
+ }
+ if (offered.isEmpty) return Nil
+ val storageKeyAttrs = AttributeSet.fromAttributeSets(offered.map(_.references))
+ if (readColumns.subsetOf(storageKeyAttrs)) Nil else offered
+ }
+
+ /**
+ * The column the reader marks the rows it checked in, when a filter is `offered`. The caller
+ * places it in the scan's output. Only where whole-stage codegen is on and `output` with the
+ * column still fits it, so the column keeps the scan in whole-stage codegen. Elsewhere the
+ * Project that drops the column could copy every surviving row. A post-scan conjunct that cannot
+ * be generated still puts that Project in a stage of its own, the trade-off SPARK-60034 accepted.
+ */
+ def checkedColumn(
+ conf: SQLConf,
+ offered: Seq[Expression],
+ output: Seq[Attribute]): Option[AttributeReference] = {
+ if (offered.isEmpty) return None
+ val checked = AttributeReference(FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME, BooleanType)()
+ Option.when(conf.wholeStageEnabled &&
+ !WholeStageCodegenExec.isTooManyFields(conf, (output :+ checked).toStructType))(checked)
+ }
+
+ /**
+ * `filters` with each of `offered` skipped on the rows `checked` marks, or as they are without
+ * the column. Matched by `ExpressionSet`, as `offered` was built. An `If` keeps two things an
+ * `Or` would lose:
+ * - a null mark leaves the conjunct as it was, an `And` around it included;
+ * - subexpression elimination cannot hoist what the conjunct shares with another one ahead of
+ * the skip.
+ */
+ def skipCheckedRows(
+ filters: Seq[Expression],
+ offered: Seq[Expression],
+ checked: Option[Attribute]): Seq[Expression] = checked match {
+ case Some(c) =>
+ val offeredSet = ExpressionSet(offered)
+ filters.map(f => if (offeredSet.contains(f)) If(c, Literal.TrueLiteral, f) else f)
+ case None => filters
+ }
+
+ /**
+ * `filter` with each scalar subquery replaced by its result, for the readers. Each bloom filter
+ * becomes a literal of its own bytes. A subquery `MergeSubplans` merged returns the blooms of
+ * several filters in one struct, so a bloom read out of it would otherwise carry the whole
+ * struct.
+ */
+ def withSubqueryResults(filter: Expression): Expression = filter.transformUp {
+ case s: ExecScalarSubquery => s.toLiteral
+ case m @ BloomFilterMightContain(bloom, _) if bloom.foldable && !bloom.isInstanceOf[Literal] =>
Review Comment:
Thanks, changed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538) to your version. The unit test covers a long field too. One correction on the link: SPARK-60448's key-column conditions hold no subquery. The ones that do come with SPARK-60500, so the description now names both.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -254,6 +279,7 @@ private class PartitionIterator[T](
}
valuePrepared = false
val result = reader.get()
+ check.foreach(_(result))
Review Comment:
Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538).
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/Scan.java:
##########
@@ -70,9 +70,10 @@ default String description() {
* {@link Table} that creates this scan returns {@link TableCapability#BATCH_READ} support in its
* {@link Table#capabilities()}.
* <p>
- * If the scan supports runtime filtering and implements {@link SupportsRuntimeFiltering},
- * this method may be called multiple times. Therefore, implementations can cache some state
- * to avoid planning the job twice.
+ * If the scan supports runtime filtering and implements {@link SupportsRuntimeFiltering}, or
+ * Spark offers it a storage filter through {@link SupportsStorageFilters}, this method may be
Review Comment:
Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538), with your wording.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterTestSource.scala:
##########
@@ -0,0 +1,444 @@
+/*
+ * 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.{ByteArrayOutputStream, ObjectOutputStream}
+import java.util
+import java.util.OptionalLong
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow}
+import org.apache.spark.sql.connector.catalog.{MetadataColumn, SupportsMetadataColumns, SupportsRead, Table, TableCapability, TableProvider}
+import org.apache.spark.sql.connector.expressions.{Expressions, FieldReference, NamedReference, SortOrder, Transform}
+import org.apache.spark.sql.connector.expressions.filter.{Predicate, StorageFilter, StorageFilterNotAppliedException}
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.sql.types.{BooleanType, DataType, LongType, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+/**
+ * A read-only v2 source over a fixed table of longs that implements [[SupportsStorageFilters]]:
+ * the key `k`, the value `v` holding `k * 10`, and with `wide` also `w` holding `k * 100`. Its
+ * reader gets the filter in the task and drops the rows the filter rejects. It reads the columns
+ * Spark requires, and builds the filter's row from the filter's `references()`. Its options:
+ * - `columnar`: reads in batches, through `StorageFilter.select`, rather than row by row;
+ * - `giveUp`: reads without the filter, as a reader does for a unit it cannot use the filter on.
+ * Here the unit is the partition;
+ * - `markAllChecked`: reads without the filter but marks every row checked, which breaks the
+ * contract on purpose;
+ * - `dropChecked`: `true` never appends the checked column, and `later` appends it to the first
+ * row or batch only. Either breaks the contract on purpose;
+ * - `declared`: the comma-separated storage-filter columns, `k` by default, a dot separating a
+ * nested field;
+ * - `valueName`: the name of the value column, `v` by default;
+ * - `wide`: adds the column `w`;
+ * - `numPartitions`: splits the keys into that many input partitions, 1 by default;
+ * - `metadataColumn`: adds the metadata column `_pos`, which holds the key;
+ * - `runtimeFiltering`: `v2` makes the scan a `SupportsRuntimeV2Filtering` one over every column,
+ * with iterative pushdown on a table partitioned by `k`. `catalyst` makes it a
+ * `SupportsRuntimeCatalystFiltering` one that evaluates the filters on `k` fully. Neither
+ * applies the filters it takes, so they serve planning tests only;
+ * - `push`: the comma-separated pushed operations it takes whole without applying them, of
+ * `limit`, `offset`, `topN` and `sample`, for planning tests only.
+ *
+ * A filter that throws `StorageFilterNotAppliedException` is given up for the rest of the
+ * partition too. The counters below record what the readers got, evaluated and returned.
+ */
+class StorageFilterTestSource extends TableProvider {
+
+ override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+ val schema = new StructType().add("k", LongType).add(options.getOrDefault("valueName", "v"),
+ LongType)
+ if (options.getBoolean("wide", false)) schema.add("w", LongType) else schema
+ }
+
+ override def getTable(
+ schema: StructType,
+ partitioning: Array[Transform],
+ properties: util.Map[String, String]): Table = {
+ val options = new CaseInsensitiveStringMap(properties)
+ if (options.getBoolean("metadataColumn", false)) {
+ new StorageFilterTestTable(schema, options) with SupportsMetadataColumns {
+ override def metadataColumns(): Array[MetadataColumn] = Array(new MetadataColumn {
+ override def name(): String = StorageFilterTestSource.posName
+ override def dataType(): DataType = LongType
+ })
+ }
+ } else {
+ new StorageFilterTestTable(schema, options)
+ }
+ }
+}
+
+object StorageFilterTestSource {
+
+ /** Rows the source serves, `k` from 0 until numRows. */
+ val numRows: Int = 200
+
+ /** Rows a batch holds at most in a columnar read. */
+ val batchSize: Int = 64
+
+ /** The name of the metadata column the `metadataColumn` option adds. */
+ val posName: String = "_pos"
+
+ /** How many readers were created with a storage-filter request. */
+ val filteredReaders = new AtomicInteger(0)
+
+ /** How many of those requests asked for the checked column. */
+ val checkedColumnRequests = new AtomicInteger(0)
+
+ /** Rows the reader evaluated the storage filter against. */
+ val rowsEvaluated = new AtomicLong(0)
+
+ /** Rows the reader returned. */
+ val rowsReturned = new AtomicLong(0)
+
+ /** The serialized size of the last filter a reader got. */
+ val filterBytes = new AtomicLong(0)
+
+ /** The references of the filters the readers got, and the identities of those filters. */
+ val filterReferences = mutable.Set.empty[Seq[String]]
+ val filterIdentities = mutable.Set.empty[Int]
+
+ /** How many times a runtime-filtering scan was handed filters. */
+ val runtimeFilterCalls = new AtomicInteger(0)
+
+ def reset(): Unit = {
+ filteredReaders.set(0)
+ checkedColumnRequests.set(0)
+ rowsEvaluated.set(0)
+ rowsReturned.set(0)
+ filterBytes.set(0)
+ filterReferences.synchronized(filterReferences.clear())
+ filterIdentities.synchronized(filterIdentities.clear())
+ runtimeFilterCalls.set(0)
+ }
+
+ /** The value of column `name` in the row of `key`. Any column but `k`, `w` and `_pos` is `v`. */
+ def valueOf(name: String, key: Long): Long = name match {
+ case "k" | `posName` => key
+ case "w" => key * 100
+ case _ => key * 10
+ }
+
+ /** The columns of the row the filter evaluates, in that order. */
+ def columnsOf(filter: StorageFilter): Seq[String] = filter.references.map(_.fieldNames.head).toSeq
+
+ // Records the filter a reader got, before the reader evaluates it.
+ private[connector] def record(filter: StorageFilter): Unit = {
+ val bytes = new ByteArrayOutputStream()
+ val out = new ObjectOutputStream(bytes)
+ try out.writeObject(filter) finally out.close()
Review Comment:
Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). Only the first reader serializes the filter now, before it evaluates it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4380305050
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,491 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GetStructField, GreaterThan, If, Literal, NumericEvalContext, PredicateHelper, Rand, XxHash64}
+import org.apache.spark.sql.catalyst.plans.logical.ReturnAnswer
+import org.apache.spark.sql.execution.{BaseLimitExec, CoalesceExec, ExecSubqueryExpression, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation}
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{BinaryType, DoubleType, LongType, StructType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper with PredicateHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
+ withSQLConf(bloomJoinConf ++ conf: _*) {
+ withTable("bf_build") {
+ spark.range(0, 30).selectExpr("id AS bk", "id AS bv")
+ .write.format("parquet").saveAsTable("bf_build")
+ body
+ }
+ }
+ }
+
+ // A join whose build side is small and carries a selective predicate, which is the shape
+ // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`.
+ private def bloomJoin(probe: DataFrame): DataFrame = {
+ val build = spark.table("bf_build").where("bv = 5")
+ probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*)
+ }
+
+ private def v2Probe(options: (String, String)*): DataFrame =
+ spark.read.format(sourceName).options(options.toMap).load()
+
+ private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME
+
+ private def scanOf(plan: SparkPlan): BatchScanExec =
+ collect(plan) { case b: BatchScanExec => b }.headOption
+ .getOrElse(fail(s"no BatchScanExec in plan:\n$plan"))
+
+ // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows
+ // the reader marks as checked.
+ private def bloomsAbove(plan: SparkPlan): (Int, Int) = {
+ val conditions = collect(plan) { case f: FilterExec => f.condition }
+ val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum
+ val skipped = conditions.map(_.collect {
+ case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain)
+ if a.name == checkedName => 1
+ }.sum).sum
+ (all - skipped, skipped)
+ }
+
+ // The scan offered nothing, and the Filter evaluates the bloom on every row.
+ private def assertNothingOffered(plan: SparkPlan): Unit = {
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.isEmpty && !scan.appendsCheckedColumn,
+ s"the scan must get no storage filter:\n$plan")
+ assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan")
+ }
+
+ // The rows the Filter on the scan returned, counted before the join.
+ private def rowsOutOfFilter(plan: SparkPlan): Long = {
+ val filters = collect(plan) {
+ case f: FilterExec if f.exists(_.isInstanceOf[BatchScanExec]) => f
+ }
+ assert(filters.size == 1, s"expected one Filter on the scan:\n$plan")
+ filters.head.metrics("numOutputRows").value
+ }
+
+ // Runs `df`, checks that it returns `expected`, and returns its plan.
+ private def checkJoin(df: DataFrame, expected: Set[Seq[Long]]): SparkPlan = {
+ val result = df.collect().map(_.toSeq.map(_.asInstanceOf[Long])).toSet
+ assert(result == expected, s"got $result")
+ df.queryExecution.executedPlan
+ }
+
+ // The build side keeps `bv = 5`, so the join returns that one row.
+ private val joinResult = Set(Seq(5L, 50L))
+
+ // Runs the bloom join over the v2 source read with `options`, checks its rows, and returns its
+ // plan.
+ private def runBloomJoin(options: (String, String)*): SparkPlan = {
+ StorageFilterTestSource.reset()
+ checkJoin(bloomJoin(v2Probe(options: _*)), joinResult)
+ }
+
+ for (columnar <- Seq(false, true); aqe <- Seq(false, true)) {
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind scan gets the checked column, and the Filter skips the bloom on checked rows " +
+ s"(AQE $aqe)") {
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val plan = runBloomJoin("columnar" -> columnar.toString)
+ val scan = scanOf(plan)
+ assert(scan.supportsColumnar == columnar, s"the scan must read $kind:\n$plan")
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(scan.appendsCheckedColumn, s"the scan must get the checked column:\n$plan")
+ assert(bloomsAbove(plan) == (0, 1),
+ s"the Filter must skip the bloom on checked rows:\n$plan")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(),
+ s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " +
+ s"${StorageFilterTestSource.rowsReturned.get()}")
+ val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output }
+ assert(probeSide.exists(!_.exists(_.name == checkedName)),
+ s"the Project above the scan must drop the checked column; got $probeSide")
+ assert(StorageFilterTestSource.filteredReaders.get() == 1,
+ s"the one partition's reader must get the filter; " +
+ s"got ${StorageFilterTestSource.filteredReaders.get()}")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 1,
+ "the reader must be asked for the checked column")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ s"the reader must evaluate every row; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(false, true).foreach { columnar =>
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind reader that gives the filter up marks its rows unchecked, and the Filter " +
+ "filters them") {
+ withBloomJoin() {
+ val plan = runBloomJoin("columnar" -> columnar.toString, "giveUp" -> "true")
+ assert(bloomsAbove(plan) == (0, 1), s"the plan is that of a filtering reader:\n$plan")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == 0,
+ s"the reader must evaluate nothing; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() == StorageFilterTestSource.numRows,
+ s"the reader must return every row; got ${StorageFilterTestSource.rowsReturned.get()}")
+ assert(rowsOutOfFilter(plan) < StorageFilterTestSource.numRows,
+ s"the Filter must apply the bloom to unchecked rows; got ${rowsOutOfFilter(plan)}")
+ }
+ }
+
+ test(s"the Filter takes a $kind reader's word on the rows it marks checked") {
+ // This reader breaks the contract on purpose: it applies no filter but marks every row
+ // checked. Only skipping the bloom on those rows keeps all of them.
+ withBloomJoin() {
+ // The join must still filter, which `runBloomJoin` checks.
+ val plan = runBloomJoin("columnar" -> columnar.toString, "markAllChecked" -> "true")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.numRows,
+ s"the Filter must keep every row marked checked; got ${rowsOutOfFilter(plan)}")
+ }
+ }
+
+ test(s"a $kind reader that leaves out the checked column it was asked for fails the task, " +
+ "from any row or batch") {
+ // A reader that gives the filter up returns every row, so it returns several batches.
+ Seq(Seq("dropChecked" -> "true"), Seq("dropChecked" -> "later", "giveUp" -> "true"))
+ .foreach { options =>
+ withBloomJoin() {
+ StorageFilterTestSource.reset()
+ val df = bloomJoin(v2Probe(options :+ ("columnar" -> columnar.toString): _*))
+ val e = intercept[SparkException](df.collect())
+ val missing =
+ Iterator.iterate[Throwable](e)(_.getCause).takeWhile(_ != null).collectFirst {
+ case t: SparkThrowable
+ if t.getCondition == "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING" => t
+ }
+ assert(missing.exists(_.getMessageParameters.get("class") ==
Review Comment:
**10. [Low] The missing-column test checks only the `class` parameter**
This walks the cause chain by hand and compares one message parameter, so nothing checks the new condition's `expected` and `actual` parameters or its SQLSTATE `KD000`. If `checkWidth` (DataSourceRDD.scala:224-228) passed the two widths the other way round, the message would say that a reader returned 3 columns, not 2, for a reader that returned 2 of 3, and all four cases here would still pass.
`DAGScheduler` rethrows a task's error as it is when the error has a condition that is not internal (DAGScheduler.scala:5233-5237), so `checkError` can take the exception directly:
```scala
checkError(
exception = intercept[SparkException](df.collect()),
condition = "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING",
sqlState = Some("KD000"),
parameters = Map(
"class" -> classOf[StorageFilterTestReaderFactory].getName,
"expected" -> "3",
"actual" -> "2"))
```
If item 1 ([thread](https://github.com/apache/spark/pull/59715#discussion_r4379088674)) changes the check or its message, this would pin the new form too.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -212,7 +236,8 @@ class DataSourceRDD(
private class PartitionIterator[T](
private var reader: PartitionReader[T],
rowCount: T => Int,
- taskState: TaskState) extends Iterator[T] {
+ taskState: TaskState,
+ check: Option[T => Unit] = None) extends Iterator[T] {
Review Comment:
**14. [Nit] Nothing uses the `= None` default of `check`**
Both constructions, at L203-204 and L210-211, pass `check`. So this default only lets a future construction compile without the width check from item 2 of my first review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009016)). Since `PartitionIterator` is private, could the default go? This is separate from item 5 ([thread](https://github.com/apache/spark/pull/59715#discussion_r4379088728)), which is about how L282 calls it.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,116 @@
+/*
+ * 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.internal.connector
+
+import org.apache.spark.SparkException
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to
+ * it. The binding is built on first use, so an instance created on the driver can be evaluated on
+ * an executor.
+ */
+class StorageFilterImpl private (
+ private var catalystExpr: CatalystExpression,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes, read on every row.
+ private val numReferences = attrs.length
+
+ // Whether this instance has logged a failure, so a task logs only its first.
+ @transient private var loggedFailure = false
+
+ @transient private lazy val boundPredicate: InternalRow => Boolean = {
+ // The expression is released below, so a copy made after a first evaluation would have nothing
+ // to build from. That is named here rather than left to an NPE.
+ if (catalystExpr == null) {
+ throw SparkException.internalError(
+ s"$this was copied after an evaluation released the expression of its predicate")
+ }
+ val predicate = CatalystPredicate.create(BindReferences.bindReference(catalystExpr, attrs))
+ // A bloom's bytes sit in the expression as a binary literal, while a generated predicate holds
+ // the deserialized filter instead, so dropping the expression lets the task reclaim them. An
+ // interpreted predicate, the fallback, keeps the expression. Every task deserializes its own
+ // copy, and this is the field's only reader.
+ catalystExpr = null
+ predicate.eval
+ }
+
+ /**
+ * Wraps an error in [[StorageFilterNotAppliedException]], unless
+ * [[StorageFilterErrors.rethrowIfMustPropagate]] rethrows it, as the file-source reader does: a
+ * task kill, or a fatal error in its cause chain. Each task logs the first one. A row of another
+ * width than [[references]] is a reader bug, which fails the task, since the predicate would read
+ * the wrong values.
+ */
+ @throws[StorageFilterNotAppliedException]
+ override def eval(row: InternalRow): Boolean = {
+ // Outside the `try`, so a failure to build the predicate, or a row of another width, fails the
+ // task.
+ if (row.numFields != numReferences) {
+ throw SparkException.internalError(
+ s"$this takes rows of $numReferences values, but got ${row.numFields}")
+ }
+ val predicate = boundPredicate
+ try {
+ predicate(row)
+ } catch {
+ case e: Throwable =>
+ StorageFilterErrors.rethrowIfMustPropagate(e, fromRead = false)
+ if (!loggedFailure) {
+ loggedFailure = true
+ logWarning(log"${MDC(LogKeys.EXPR, this)} could not be applied, so the reader reads " +
+ log"without it.", e)
+ }
+ throw new StorageFilterNotAppliedException(e)
+ }
+ }
+
+ @transient override lazy val references: Array[NamedReference] =
Review Comment:
**9. [Low] `references()` hands every reader of a task the same mutable array**
This lazy val returns one array to every caller, and the readers of a task share the filter (StorageFilter.java:51-52). The array's order is the row layout `eval` and `select` expect (StorageFilter.java:61-62), while the predicate stays bound to `attrs`. So a reader that sorts the array in place, e.g. into its file's column order, changes the layout that it and every later reader of the task read from it. The `references()` doc does not say that the array must not be modified.
Take a two-key bloom join on BIGINT `k` and `v`, so one filter over `[k, v]`. A reader that sorts the array to `[v, k]` builds `(v, k)` rows from then on, and so does every later reader of its task, e.g. under `coalesce(1)`. The width check at L74 passes, and the `k` bloom is probed with `v` values. The rows it rejects are dropped inside the reader, so the post-scan `Filter` never sees them, and the join loses them without an error. `PartitionPredicateImpl` caches its array too, but each of its references carries its own ordinal, so a reordered array changes nothing there.
Could `references()` return a copy, e.g. `references.clone()`, or could its doc say that the array must not be modified?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -187,11 +196,19 @@ class DataSourceRDD(
castPartition(split).inputPartition.iterator.flatMap { inputPartition =>
val iter = if (columnarReads) {
- val batchReader = partitionReaderFactory.createColumnarReader(inputPartition)
- new PartitionIterator[ColumnarBatch](batchReader, _.numRows, taskState)
+ val batchReader = storageFilterRequest match {
Review Comment:
**17. [Nit] Both read modes repeat the same reader routing**
L199-202 and L206-209 match `storageFilterRequest` the same way, and L204 and L211 build the same width check, differing only in how they read the width. A small helper that takes the two factory methods would keep the routing in one place. It saves only a few lines, so keeping it as is is fine too.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,116 @@
+/*
+ * 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.internal.connector
+
+import org.apache.spark.SparkException
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to
+ * it. The binding is built on first use, so an instance created on the driver can be evaluated on
+ * an executor.
+ */
+class StorageFilterImpl private (
+ private var catalystExpr: CatalystExpression,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes, read on every row.
+ private val numReferences = attrs.length
+
+ // Whether this instance has logged a failure, so a task logs only its first.
+ @transient private var loggedFailure = false
+
+ @transient private lazy val boundPredicate: InternalRow => Boolean = {
Review Comment:
**12. [Nit] `boundPredicate` repeats `ParquetStorageFilter.preparedPredicate`**
Since the fix for item 9 of my first review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009048)), this block follows ParquetStorageFilter.scala:138-153 almost line by line: the null check with its internal error, the build, the comment on the bloom's bytes, and the release. The two copies differ only in their internal-error messages. The fix for item 1 of that review moved the error rule into `StorageFilterErrors`, so both paths share one rule. Could this step be shared the same way, e.g. through a small serializable holder of the bound expression next to `StorageFilterErrors`, with `apply` binding the expression?
Also, `CatalystPredicate.create(BindReferences.bindReference(catalystExpr, attrs))` at L54 is the existing `CatalystPredicate.create(catalystExpr, attrs)` (predicates.scala:92).
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterTestSource.scala:
##########
@@ -0,0 +1,444 @@
+/*
+ * 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.{ByteArrayOutputStream, ObjectOutputStream}
+import java.util
+import java.util.OptionalLong
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow}
+import org.apache.spark.sql.connector.catalog.{MetadataColumn, SupportsMetadataColumns, SupportsRead, Table, TableCapability, TableProvider}
+import org.apache.spark.sql.connector.expressions.{Expressions, FieldReference, NamedReference, SortOrder, Transform}
+import org.apache.spark.sql.connector.expressions.filter.{Predicate, StorageFilter, StorageFilterNotAppliedException}
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.sql.types.{BooleanType, DataType, LongType, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+/**
+ * A read-only v2 source over a fixed table of longs that implements [[SupportsStorageFilters]]:
+ * the key `k`, the value `v` holding `k * 10`, and with `wide` also `w` holding `k * 100`. Its
+ * reader gets the filter in the task and drops the rows the filter rejects. It reads the columns
+ * Spark requires, and builds the filter's row from the filter's `references()`. Its options:
+ * - `columnar`: reads in batches, through `StorageFilter.select`, rather than row by row;
+ * - `giveUp`: reads without the filter, as a reader does for a unit it cannot use the filter on.
+ * Here the unit is the partition;
+ * - `markAllChecked`: reads without the filter but marks every row checked, which breaks the
+ * contract on purpose;
+ * - `dropChecked`: `true` never appends the checked column, and `later` appends it to the first
+ * row or batch only. Either breaks the contract on purpose;
+ * - `declared`: the comma-separated storage-filter columns, `k` by default, a dot separating a
+ * nested field;
+ * - `valueName`: the name of the value column, `v` by default;
+ * - `wide`: adds the column `w`;
+ * - `numPartitions`: splits the keys into that many input partitions, 1 by default;
+ * - `metadataColumn`: adds the metadata column `_pos`, which holds the key;
+ * - `runtimeFiltering`: `v2` makes the scan a `SupportsRuntimeV2Filtering` one over every column,
+ * with iterative pushdown on a table partitioned by `k`. `catalyst` makes it a
+ * `SupportsRuntimeCatalystFiltering` one that evaluates the filters on `k` fully. Neither
+ * applies the filters it takes, so they serve planning tests only;
+ * - `push`: the comma-separated pushed operations it takes whole without applying them, of
+ * `limit`, `offset`, `topN` and `sample`, for planning tests only.
+ *
+ * A filter that throws `StorageFilterNotAppliedException` is given up for the rest of the
+ * partition too. The counters below record what the readers got, evaluated and returned.
+ */
+class StorageFilterTestSource extends TableProvider {
+
+ override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+ val schema = new StructType().add("k", LongType).add(options.getOrDefault("valueName", "v"),
+ LongType)
+ if (options.getBoolean("wide", false)) schema.add("w", LongType) else schema
+ }
+
+ override def getTable(
+ schema: StructType,
+ partitioning: Array[Transform],
+ properties: util.Map[String, String]): Table = {
+ val options = new CaseInsensitiveStringMap(properties)
+ if (options.getBoolean("metadataColumn", false)) {
+ new StorageFilterTestTable(schema, options) with SupportsMetadataColumns {
+ override def metadataColumns(): Array[MetadataColumn] = Array(new MetadataColumn {
+ override def name(): String = StorageFilterTestSource.posName
+ override def dataType(): DataType = LongType
+ })
+ }
+ } else {
+ new StorageFilterTestTable(schema, options)
+ }
+ }
+}
+
+object StorageFilterTestSource {
+
+ /** Rows the source serves, `k` from 0 until numRows. */
+ val numRows: Int = 200
+
+ /** Rows a batch holds at most in a columnar read. */
+ val batchSize: Int = 64
+
+ /** The name of the metadata column the `metadataColumn` option adds. */
+ val posName: String = "_pos"
+
+ /** How many readers were created with a storage-filter request. */
+ val filteredReaders = new AtomicInteger(0)
+
+ /** How many of those requests asked for the checked column. */
+ val checkedColumnRequests = new AtomicInteger(0)
+
+ /** Rows the reader evaluated the storage filter against. */
+ val rowsEvaluated = new AtomicLong(0)
+
+ /** Rows the reader returned. */
+ val rowsReturned = new AtomicLong(0)
+
+ /** The serialized size of the last filter a reader got. */
+ val filterBytes = new AtomicLong(0)
+
+ /** The references of the filters the readers got, and the identities of those filters. */
+ val filterReferences = mutable.Set.empty[Seq[String]]
+ val filterIdentities = mutable.Set.empty[Int]
+
+ /** How many times a runtime-filtering scan was handed filters. */
+ val runtimeFilterCalls = new AtomicInteger(0)
+
+ def reset(): Unit = {
+ filteredReaders.set(0)
+ checkedColumnRequests.set(0)
+ rowsEvaluated.set(0)
+ rowsReturned.set(0)
+ filterBytes.set(0)
+ filterReferences.synchronized(filterReferences.clear())
+ filterIdentities.synchronized(filterIdentities.clear())
+ runtimeFilterCalls.set(0)
+ }
+
+ /** The value of column `name` in the row of `key`. Any column but `k`, `w` and `_pos` is `v`. */
+ def valueOf(name: String, key: Long): Long = name match {
+ case "k" | `posName` => key
+ case "w" => key * 100
+ case _ => key * 10
+ }
+
+ /** The columns of the row the filter evaluates, in that order. */
+ def columnsOf(filter: StorageFilter): Seq[String] = filter.references.map(_.fieldNames.head).toSeq
+
+ // Records the filter a reader got, before the reader evaluates it.
+ private[connector] def record(filter: StorageFilter): Unit = {
+ val bytes = new ByteArrayOutputStream()
Review Comment:
**16. [Nit] The test source repeats `Utils.serialize` and `RangeInputPartition`**
L150-153 compute `Utils.serialize(filter).length`, which would also drop the `java.io` import at L20, with or without item 7 ([thread](https://github.com/apache/spark/pull/59715#discussion_r4379088759)). `StorageFilterTestPartition` (L269) is the same as `RangeInputPartition` (DataSourceV2Suite.scala:2162), which `FakeV2Provider` in this package already reuses. And `Utils.stringToSeq` could split the `push` and `declared` options at L191 and L203.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,491 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GetStructField, GreaterThan, If, Literal, NumericEvalContext, PredicateHelper, Rand, XxHash64}
+import org.apache.spark.sql.catalyst.plans.logical.ReturnAnswer
+import org.apache.spark.sql.execution.{BaseLimitExec, CoalesceExec, ExecSubqueryExpression, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation}
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{BinaryType, DoubleType, LongType, StructType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper with PredicateHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
+ withSQLConf(bloomJoinConf ++ conf: _*) {
+ withTable("bf_build") {
+ spark.range(0, 30).selectExpr("id AS bk", "id AS bv")
+ .write.format("parquet").saveAsTable("bf_build")
+ body
+ }
+ }
+ }
+
+ // A join whose build side is small and carries a selective predicate, which is the shape
+ // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`.
+ private def bloomJoin(probe: DataFrame): DataFrame = {
+ val build = spark.table("bf_build").where("bv = 5")
+ probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*)
+ }
+
+ private def v2Probe(options: (String, String)*): DataFrame =
+ spark.read.format(sourceName).options(options.toMap).load()
+
+ private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME
+
+ private def scanOf(plan: SparkPlan): BatchScanExec =
+ collect(plan) { case b: BatchScanExec => b }.headOption
+ .getOrElse(fail(s"no BatchScanExec in plan:\n$plan"))
+
+ // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows
+ // the reader marks as checked.
+ private def bloomsAbove(plan: SparkPlan): (Int, Int) = {
+ val conditions = collect(plan) { case f: FilterExec => f.condition }
+ val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum
+ val skipped = conditions.map(_.collect {
+ case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain)
+ if a.name == checkedName => 1
+ }.sum).sum
+ (all - skipped, skipped)
+ }
+
+ // The scan offered nothing, and the Filter evaluates the bloom on every row.
+ private def assertNothingOffered(plan: SparkPlan): Unit = {
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.isEmpty && !scan.appendsCheckedColumn,
Review Comment:
**15. [Nit] The second half of this assertion cannot fail**
`BatchScanExec` requires storage filters for a checked column (BatchScanExec.scala:62-64), so `!scan.appendsCheckedColumn` holds whenever `scan.storageFilters.isEmpty` does. L398 already pins that `require`. Could this assert only `scan.storageFilters.isEmpty`, here and at L323?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -957,7 +972,8 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
// changes with variant extraction. When validConstraints is wired up, this needs revisiting.
// Pushed variant extraction cannot be reproduced by rebuilding the scan, so this scan stays
// not-mergeable (the default `mergeableScan = false`).
- val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput)
+ val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput,
+ onlyRowWisePushdowns = !hasRowSelectingPushdown(holder))
Review Comment:
**11. [Low] No test covers the storage-filter path of a variant-extraction scan**
Two changes in this PR exist only for this kind of scan, and no test reaches either:
- This argument, from the fix for item 18 of my first review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009081)). `StorageFilterTestScanBuilder` does not implement `SupportsPushDownVariantExtractions`, and the planning test sets the flag by hand with `r.copy(mergeableScan = ..., onlyRowWisePushdowns = ...)` (StorageFilterPushDownSuite.scala:451). So if this argument were dropped, no variant-extraction scan would be offered a storage filter, as in that item's example, and every test would still pass.
- The name match for metadata columns from `cb01a6f` (DataSourceV2Strategy.scala:144-148), which is there because this scan's output has fresh exprIds from `toAttributes(scan.readSchema())` (L965). The metadata-column test (StorageFilterPushDownSuite.scala:278-290) reads through `pruneColumns`, whose output keeps the relation's exprIds (PushDownUtils.scala:715-722). There the old exprId match and the new name match agree, so going back to the exprId match would pass every test too.
Could the test source support variant extraction, with a test that runs that item's example query over it?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -127,6 +127,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
tableSpec.withNewLocation(newLoc)
}
+ /**
+ * The conjuncts of `filters` offered to the scan as storage filters, see
+ * `SupportsStorageFilters` and `StorageFilterPlanning.offeredConjuncts`. Only runtime blooms for
+ * now, as on the file-source side. The columns left to skip are the data columns, as there, so a
+ * scan that reads only the key and metadata columns is offered none.
+ */
+ private def storageFiltersFor(
+ relation: DataSourceV2ScanRelation,
+ filters: Seq[Expression]): Seq[Expression] = {
+ if (!session.sessionState.conf.v2StorageFilterPushdownEnabled) return Nil
+ if (relation.storageFilterAttrs.isEmpty) return Nil
Review Comment:
**13. [Nit] `storageFiltersFor` matches the declared columns before its cheaper checks**
With the conf on, every scan that implements `SupportsStorageFilters` reaches `storageFilterAttrs` here. That calls `storageFilterAttributes()` and matches every declared name against every output column (DataSourceV2Relation.scala:285-290), before L143 checks `onlyRowWisePushdowns`, and although only a bloom can be offered. Most scans get no runtime bloom. For a wide table that declares all of its 5,000 columns, that is up to 25 million name comparisons each time such a scan is planned. `eaf3758935e` removed the early return that skipped this. Could it come back, after the flag check?
```scala
if (!relation.onlyRowWisePushdowns) return Nil
if (!filters.exists(_.isInstanceOf[BloomFilterMightContain])) return Nil
if (relation.storageFilterAttrs.isEmpty) return Nil
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4386746505
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/StorageFilterPlanning.scala:
##########
@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, Expression, ExpressionSet, GetStructField, If, Literal}
+import org.apache.spark.sql.execution.{ScalarSubquery => ExecScalarSubquery, WholeStageCodegenExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.BooleanType
+
+/**
+ * The rules file-source and DSv2 scans share for storage filters, see
+ * `FileFormat.buildReaderWithStorageFilters` and `SupportsStorageFilters`. Each strategy decides
+ * on its own whether its scan can take any.
+ */
+object StorageFilterPlanning {
+
+ /**
+ * The conjuncts of `filters` offered to the scan, each once, as an `ExpressionSet` counts them.
+ * Each one:
+ * - is deterministic, since the post-scan `Filter` takes the reader's word on the rows it
+ * checked;
+ * - references at least one column, and only `keyCandidates`;
+ * - is taken by `accepts`.
+ *
+ * Nothing is offered when the offered conjuncts reference every one of `readColumns`, since the
+ * reader then reads the same columns for the same rows either way.
+ */
+ def offeredConjuncts(
+ filters: Seq[Expression],
+ keyCandidates: AttributeSet,
+ readColumns: AttributeSet,
+ accepts: Expression => Boolean): Seq[Expression] = {
+ val offered = ExpressionSet(filters).toSeq.filter { expr =>
+ val refs = expr.references
+ expr.deterministic && refs.nonEmpty && refs.subsetOf(keyCandidates) && accepts(expr)
+ }
+ if (offered.isEmpty) return Nil
+ val storageKeyAttrs = AttributeSet.fromAttributeSets(offered.map(_.references))
+ if (readColumns.subsetOf(storageKeyAttrs)) Nil else offered
+ }
+
+ /**
+ * The column the reader marks the rows it checked in, when a filter is `offered`. The caller
+ * places it in the scan's output. Only where whole-stage codegen is on and `output` with the
+ * column still fits it, so the column keeps the scan in whole-stage codegen. Elsewhere the
+ * Project that drops the column could copy every surviving row. A post-scan conjunct that cannot
+ * be generated still puts that Project in a stage of its own, the trade-off SPARK-60034 accepted.
+ */
+ def checkedColumn(
+ conf: SQLConf,
+ offered: Seq[Expression],
+ output: Seq[Attribute]): Option[AttributeReference] = {
+ if (offered.isEmpty) return None
+ val checked = AttributeReference(FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME, BooleanType)()
+ Option.when(conf.wholeStageEnabled &&
+ !WholeStageCodegenExec.isTooManyFields(conf, (output :+ checked).toStructType))(checked)
+ }
+
+ /**
+ * `filters` with each of `offered` skipped on the rows `checked` marks, or as they are without
+ * the column. Matched by `ExpressionSet`, as `offered` was built. An `If` keeps two things an
+ * `Or` would lose:
+ * - a null mark leaves the conjunct as it was, an `And` around it included;
+ * - subexpression elimination cannot hoist what the conjunct shares with another one ahead of
+ * the skip.
+ */
+ def skipCheckedRows(
+ filters: Seq[Expression],
+ offered: Seq[Expression],
+ checked: Option[Attribute]): Seq[Expression] = checked match {
+ case Some(c) =>
+ val offeredSet = ExpressionSet(offered)
+ filters.map(f => if (offeredSet.contains(f)) If(c, Literal.TrueLiteral, f) else f)
+ case None => filters
+ }
+
+ /**
+ * `filter` with each scalar subquery replaced by its result, for the readers. A field read out of
+ * a result becomes a literal of its own. A subquery `MergeSubplans` merged returns the results of
+ * several subqueries in one struct, such as the blooms of several filters, so a field read out of
+ * it would otherwise carry the whole struct.
+ */
+ def withSubqueryResults(filter: Expression): Expression = filter.transformUp {
+ case s: ExecScalarSubquery => s.toLiteral
Review Comment:
Thanks, changed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb) to your version. The field is read straight from the subquery's result, so only that field is copied.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/StorageFilterRequest.java:
##########
@@ -0,0 +1,51 @@
+/*
+ * 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.read;
+
+import java.io.Serializable;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * What Spark hands a reader of a {@link SupportsStorageFilters} scan, through
+ * {@link PartitionReaderFactory#createReader(InputPartition, StorageFilterRequest)} or
+ * {@link PartitionReaderFactory#createColumnarReader(InputPartition, StorageFilterRequest)}.
+ *
+ * @since 5.0.0
+ */
+@Evolving
+public interface StorageFilterRequest extends Serializable {
+
+ /** The filter the reader may apply. */
+ StorageFilter filter();
+
+ /**
+ * Whether the reader must append a boolean column after the columns of
+ * {@link Scan#readSchema()}. It is true only on a row the filter was evaluated on and kept. It is
+ * false on every other row the reader returns, such as a row read without the filter, or one
+ * kept although the filter rejected it. Spark skips the filter's predicate above the scan on the
+ * rows marked true. A reader asked for the column that returns a row or a batch without it fails
+ * the task. The reader may return more columns after the checked column, which Spark does not
+ * read.
Review Comment:
Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). I went back to the exact width check, and kept the boolean check for a batch. So your example fails the task again, and so does a reader that returns more columns after the checked one. The doc and the description now say what Spark checks. It fails the task for a row or a batch of another width, or a batch whose last column is not a boolean.
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -2397,6 +2397,18 @@
],
"sqlState" : "0A000"
},
+ "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING" : {
+ "message" : [
+ "A reader of <class> was asked for the column that marks the rows a storage filter checked, at position <expected>. That column is missing or not a boolean. The reader returned <actual> columns."
Review Comment:
Thanks, changed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb) to your message and parameter names. It also says the column may not be the last one, since the check is exact again, see item 2.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -1060,6 +1061,46 @@ class ParquetStorageFilterSuite extends QueryTest with SharedSparkSession
}
}
+ Seq(false, true).foreach { aqe =>
+ test(s"a two-key join ships each bloom to the reader once (AQE $aqe)") {
+ // `MergeSubplans` merges the two blooms' aggregates into one subquery returning both blooms.
+ withTable("bf3", "bf4") {
+ spark.range(600).selectExpr("id AS k", "id * 10 AS v", "id * 100 AS w")
+ .write.format("parquet").saveAsTable("bf3")
+ spark.range(30).selectExpr("id AS k", "id * 10 AS v", "id AS flag")
+ .write.format("parquet").saveAsTable("bf4")
+ withSQLConf(
+ SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val df = spark.sql(
+ "SELECT bf3.k, bf3.v, bf3.w FROM bf3 JOIN bf4 ON bf3.k = bf4.k AND bf3.v = bf4.v " +
+ "WHERE bf4.flag = 5")
+ val rows = df.collect().map(r => (r.getLong(0), r.getLong(1), r.getLong(2))).toSeq
+ assert(rows == Seq((5L, 50L, 500L)), s"got $rows")
+ val plan = df.queryExecution.executedPlan
+ val scan = collectFirst(plan) {
+ case s: FileSourceScanExec if s.storageFilters.nonEmpty => s
+ }.getOrElse(fail(s"no scan with storage filters:\n$plan"))
+ assert(scan.storageFilters.size == 2 && scan.storageFilters.forall(_.exists {
+ case GetStructField(_: ExecSubqueryExpression, _, _) => true
+ case _ => false
+ }), s"test setup: both blooms must come from one merged subquery:\n$plan")
+ val bloomBytes = scan.storageFilters.map(_.collectFirst {
+ case m: BloomFilterMightContain =>
+ m.bloomFilterExpression.eval().asInstanceOf[Array[Byte]].length
+ }.get).sum
+ // The reader's function ships with the scan's RDD, and the blooms with it.
+ val shipped = Utils.serialize(scan.inputRDD).length
+ assert(shipped < bloomBytes * 3 / 2,
Review Comment:
Thanks, added in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). The test now also checks that the reader excluded rows. It fails when the reader gets no filter.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterImplSuite.scala:
##########
@@ -0,0 +1,172 @@
+/*
+ * 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 org.apache.spark.{SparkArithmeticException, SparkConf, SparkException, SparkFunSuite, SparkUnsupportedOperationException, TaskContext}
+import org.apache.spark.serializer.JavaSerializer
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, Cast, Divide, EqualTo, EvalMode, Expression, GreaterThan, Literal, NumericEvalContext, Remainder, UnaryExpression}
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory, StorageFilterNotAppliedException}
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.internal.connector.{StorageFilterImpl, StorageFilterRequestImpl}
+import org.apache.spark.sql.types.{BooleanType, DataType, DoubleType, LongType}
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+/** Tests the Catalyst-backed [[StorageFilterImpl]] a connector's reader evaluates. */
+class StorageFilterImplSuite extends SparkFunSuite {
+
+ private val k = AttributeReference("k", LongType)()
+ private val v = AttributeReference("v", LongType)()
+
+ // `k % 2 = 0` over a scan whose output has `v` before `k`.
+ private def evenKeys: StorageFilterImpl =
+ StorageFilterImpl(EqualTo(Remainder(k, Literal(2L)), Literal(0L)), Seq(v, k))
+
+ // `1 / k > 0`, which raises an error at `k = 0` under ANSI.
+ private def positiveInverse: StorageFilterImpl = StorageFilterImpl(
+ GreaterThan(
+ Divide(Literal(1.0), Cast(k, DoubleType), NumericEvalContext(EvalMode.ANSI)),
+ Literal(0.0)),
+ Seq(k))
+
+ private def withBatchOf(keys: Long*)(body: ColumnarBatch => Unit): Unit = {
+ val vector = new OnHeapColumnVector(keys.length, LongType)
+ keys.zipWithIndex.foreach { case (key, i) => vector.putLong(i, key) }
+ val batch = new ColumnarBatch(Array(vector), keys.length)
+ try body(batch) finally batch.close()
+ }
+
+ test("the references are the columns the filter reads") {
+ assert(evenKeys.references.map(_.toString).toSeq == Seq("k"))
+ }
+
+ test("a filter over two columns takes them in the scan output's order") {
+ val filter =
+ StorageFilterImpl(And(EqualTo(k, Literal(1L)), EqualTo(v, Literal(10L))), Seq(v, k))
+ assert(filter.references.map(_.toString).toSeq == Seq("v", "k"))
+ assert(filter.eval(InternalRow(10L, 1L)), "the row holds v, then k")
+ assert(!filter.eval(InternalRow(1L, 10L)))
+ }
+
+ test("a reader that reorders the references changes the layout of no other reader") {
+ val filter =
+ StorageFilterImpl(And(EqualTo(k, Literal(1L)), EqualTo(v, Literal(10L))), Seq(v, k))
+ val references = filter.references()
+ references.sortInPlaceBy(_.toString)
+ assert(filter.references().map(_.toString).toSeq == Seq("v", "k"))
+ }
+
+ test("a row of another width than the references fails, rather than being misread") {
+ // `evenKeys` reads `k` alone, so a row holding `v, k` would put `v` where `k` is read.
+ val e = intercept[SparkException](evenKeys.eval(InternalRow(3L, 4L)))
+ assert(e.getCondition == "INTERNAL_ERROR", s"got ${e.getCondition}")
+ }
+
+ test("a filter over a column outside the scan output is an internal error") {
+ val e = intercept[SparkException](StorageFilterImpl(EqualTo(k, Literal(1L)), Seq(v)))
+ assert(e.getCondition == "INTERNAL_ERROR", s"got ${e.getCondition}")
+ }
+
+ test("select writes the kept rows' positions in ascending order") {
+ withBatchOf(3L, 4L, 5L, 6L, 8L) { batch =>
+ val selected = new Array[Int](batch.numRows)
+ val count = evenKeys.select(batch, selected)
+ assert(selected.take(count).toSeq == Seq(1, 3, 4))
+ }
+ }
+
+ test("an error evaluating the filter answers that it is not applied") {
+ val filter = positiveInverse
+ assert(filter.eval(InternalRow(2L)))
+ val e = intercept[StorageFilterNotAppliedException](filter.eval(InternalRow(0L)))
+ assert(e.getCondition == "STORAGE_FILTER_NOT_APPLIED", s"got ${e.getCondition}")
+ assert(e.getCause.isInstanceOf[SparkArithmeticException], s"got ${e.getCause}")
+ withBatchOf(1L, 0L) { batch =>
+ intercept[StorageFilterNotAppliedException](filter.select(batch, new Array[Int](2)))
+ }
+ }
+
+ test("a reader factory without the storage-filter methods cannot create a filtering reader") {
+ val factory = new PartitionReaderFactory {
+ override def createReader(partition: InputPartition): PartitionReader[InternalRow] = null
+ }
+ val request = StorageFilterRequestImpl(evenKeys, appendsCheckedColumn = false)
+ Seq(
+ () => factory.createReader(RangeInputPartition(0, 1), request),
+ () => factory.createColumnarReader(RangeInputPartition(0, 1), request)
+ ).foreach { create =>
+ val e = intercept[SparkUnsupportedOperationException](create())
+ assert(e.getCondition == "DATA_SOURCE_STORAGE_FILTER_READER_NOT_SUPPORTED",
Review Comment:
Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). Both tests use `checkError` with the SQLSTATE. The one for `DATA_SOURCE_STORAGE_FILTER_READER_NOT_SUPPORTED` checks the factory's class too.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -124,9 +124,10 @@ case class GroupPartitionsExec(
* where this check passes, and `UnknownPartitioning`, where re-deriving reports the same thing.
*
* The comparison is by reference for a child that reports the same object, which the columnar and
- * codegen wrappers do. `PlanAdaptiveDynamicPruningFilters` rebuilds `BatchScanExec`, so there it
- * compares the partition keys by value, and reads equal only while the connector returns the same
- * input partitions for the same scan.
+ * codegen wrappers do. `PlanAdaptiveDynamicPruningFilters` and `ReuseAdaptiveSubquery` rebuild a
+ * `BatchScanExec` whose runtime or storage filters hold a subquery, so there it compares the
+ * partition keys by value, and reads equal only while the connector returns the same input
+ * partitions for the same scan.
Review Comment:
Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb), with your wording.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #59715:
URL: https://github.com/apache/spark/pull/59715#issuecomment-6442788106
Thank you so much, @peter-toth !
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4383481085
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/StorageFilterPlanning.scala:
##########
@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, Expression, ExpressionSet, GetStructField, If, Literal}
+import org.apache.spark.sql.execution.{ScalarSubquery => ExecScalarSubquery, WholeStageCodegenExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.BooleanType
+
+/**
+ * The rules file-source and DSv2 scans share for storage filters, see
+ * `FileFormat.buildReaderWithStorageFilters` and `SupportsStorageFilters`. Each strategy decides
+ * on its own whether its scan can take any.
+ */
+object StorageFilterPlanning {
+
+ /**
+ * The conjuncts of `filters` offered to the scan, each once, as an `ExpressionSet` counts them.
+ * Each one:
+ * - is deterministic, since the post-scan `Filter` takes the reader's word on the rows it
+ * checked;
+ * - references at least one column, and only `keyCandidates`;
+ * - is taken by `accepts`.
+ *
+ * Nothing is offered when the offered conjuncts reference every one of `readColumns`, since the
+ * reader then reads the same columns for the same rows either way.
+ */
+ def offeredConjuncts(
+ filters: Seq[Expression],
+ keyCandidates: AttributeSet,
+ readColumns: AttributeSet,
+ accepts: Expression => Boolean): Seq[Expression] = {
+ val offered = ExpressionSet(filters).toSeq.filter { expr =>
+ val refs = expr.references
+ expr.deterministic && refs.nonEmpty && refs.subsetOf(keyCandidates) && accepts(expr)
+ }
+ if (offered.isEmpty) return Nil
+ val storageKeyAttrs = AttributeSet.fromAttributeSets(offered.map(_.references))
+ if (readColumns.subsetOf(storageKeyAttrs)) Nil else offered
+ }
+
+ /**
+ * The column the reader marks the rows it checked in, when a filter is `offered`. The caller
+ * places it in the scan's output. Only where whole-stage codegen is on and `output` with the
+ * column still fits it, so the column keeps the scan in whole-stage codegen. Elsewhere the
+ * Project that drops the column could copy every surviving row. A post-scan conjunct that cannot
+ * be generated still puts that Project in a stage of its own, the trade-off SPARK-60034 accepted.
+ */
+ def checkedColumn(
+ conf: SQLConf,
+ offered: Seq[Expression],
+ output: Seq[Attribute]): Option[AttributeReference] = {
+ if (offered.isEmpty) return None
+ val checked = AttributeReference(FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME, BooleanType)()
+ Option.when(conf.wholeStageEnabled &&
+ !WholeStageCodegenExec.isTooManyFields(conf, (output :+ checked).toStructType))(checked)
+ }
+
+ /**
+ * `filters` with each of `offered` skipped on the rows `checked` marks, or as they are without
+ * the column. Matched by `ExpressionSet`, as `offered` was built. An `If` keeps two things an
+ * `Or` would lose:
+ * - a null mark leaves the conjunct as it was, an `And` around it included;
+ * - subexpression elimination cannot hoist what the conjunct shares with another one ahead of
+ * the skip.
+ */
+ def skipCheckedRows(
+ filters: Seq[Expression],
+ offered: Seq[Expression],
+ checked: Option[Attribute]): Seq[Expression] = checked match {
+ case Some(c) =>
+ val offeredSet = ExpressionSet(offered)
+ filters.map(f => if (offeredSet.contains(f)) If(c, Literal.TrueLiteral, f) else f)
+ case None => filters
+ }
+
+ /**
+ * `filter` with each scalar subquery replaced by its result, for the readers. A field read out of
+ * a result becomes a literal of its own. A subquery `MergeSubplans` merged returns the results of
+ * several subqueries in one struct, such as the blooms of several filters, so a field read out of
+ * it would otherwise carry the whole struct.
+ */
+ def withSubqueryResults(filter: Expression): Expression = filter.transformUp {
+ case s: ExecScalarSubquery => s.toLiteral
Review Comment:
**1. [Low] Each conjunct still copies every bloom of a merged subquery on the driver**
This is the shape I suggested in item 4 of my second review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4379088708)), and it has a cost I missed. `transformUp` turns the `ScalarSubquery` leaf into `s.toLiteral` before it reaches the `GetStructField` above it, and `toLiteral` is `Literal.create(result, dataType)` over the whole merged struct (subquery.scala:112-115). `Literal`'s constructor always validates its value (literals.scala:499), and for a struct it reads every field with `row.get(i, fieldDataType)` (literals.scala:286-292). `result` is a view into the `UnsafeRow` the subquery collected (subquery.scala:95), so each of those reads is `UnsafeRow.getBinary`, which copies the whole bloom (UnsafeRow.java:454-455), only for the copy to be dropped after a type check.
So a conjunct over a merge of N blooms makes N throwaway copies before the one it keeps, and a scan that takes all N blooms allocates N x (N + 1) bloom-sized arrays on the driver each time it runs. In the two-key tests that is 4 MiB of garbage for two 1 MiB blooms. `BatchScanExec.storageFilterRequest` (BatchScanExec.scala:130-131) adds this, and `FileSourceScanExec.preparedStorageFilters` already paid the N x N part on master through its own `s.toLiteral`. Reading the field straight from the subquery never builds the struct literal:
```scala
def withSubqueryResults(filter: Expression): Expression = filter.transformDown {
case g @ GetStructField(_: ExecScalarSubquery | _: Literal, _, _) =>
Literal(g.eval(), g.dataType)
case s: ExecScalarSubquery => s.toLiteral
}
```
`g.eval()` reaches `ScalarSubquery.eval`, which returns `result` without building a `Literal`, so only `GetStructField` copies its one field. Unlike `transformUp`, this folds only the inner read of a nested `GetStructField(GetStructField(...))`, but `MergeSubplans` builds single-level reads only (MergeSubplans.scala:331-335).
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterImplSuite.scala:
##########
@@ -0,0 +1,172 @@
+/*
+ * 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 org.apache.spark.{SparkArithmeticException, SparkConf, SparkException, SparkFunSuite, SparkUnsupportedOperationException, TaskContext}
+import org.apache.spark.serializer.JavaSerializer
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, Cast, Divide, EqualTo, EvalMode, Expression, GreaterThan, Literal, NumericEvalContext, Remainder, UnaryExpression}
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory, StorageFilterNotAppliedException}
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.internal.connector.{StorageFilterImpl, StorageFilterRequestImpl}
+import org.apache.spark.sql.types.{BooleanType, DataType, DoubleType, LongType}
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+/** Tests the Catalyst-backed [[StorageFilterImpl]] a connector's reader evaluates. */
+class StorageFilterImplSuite extends SparkFunSuite {
+
+ private val k = AttributeReference("k", LongType)()
+ private val v = AttributeReference("v", LongType)()
+
+ // `k % 2 = 0` over a scan whose output has `v` before `k`.
+ private def evenKeys: StorageFilterImpl =
+ StorageFilterImpl(EqualTo(Remainder(k, Literal(2L)), Literal(0L)), Seq(v, k))
+
+ // `1 / k > 0`, which raises an error at `k = 0` under ANSI.
+ private def positiveInverse: StorageFilterImpl = StorageFilterImpl(
+ GreaterThan(
+ Divide(Literal(1.0), Cast(k, DoubleType), NumericEvalContext(EvalMode.ANSI)),
+ Literal(0.0)),
+ Seq(k))
+
+ private def withBatchOf(keys: Long*)(body: ColumnarBatch => Unit): Unit = {
+ val vector = new OnHeapColumnVector(keys.length, LongType)
+ keys.zipWithIndex.foreach { case (key, i) => vector.putLong(i, key) }
+ val batch = new ColumnarBatch(Array(vector), keys.length)
+ try body(batch) finally batch.close()
+ }
+
+ test("the references are the columns the filter reads") {
+ assert(evenKeys.references.map(_.toString).toSeq == Seq("k"))
+ }
+
+ test("a filter over two columns takes them in the scan output's order") {
+ val filter =
+ StorageFilterImpl(And(EqualTo(k, Literal(1L)), EqualTo(v, Literal(10L))), Seq(v, k))
+ assert(filter.references.map(_.toString).toSeq == Seq("v", "k"))
+ assert(filter.eval(InternalRow(10L, 1L)), "the row holds v, then k")
+ assert(!filter.eval(InternalRow(1L, 10L)))
+ }
+
+ test("a reader that reorders the references changes the layout of no other reader") {
+ val filter =
+ StorageFilterImpl(And(EqualTo(k, Literal(1L)), EqualTo(v, Literal(10L))), Seq(v, k))
+ val references = filter.references()
+ references.sortInPlaceBy(_.toString)
+ assert(filter.references().map(_.toString).toSeq == Seq("v", "k"))
+ }
+
+ test("a row of another width than the references fails, rather than being misread") {
+ // `evenKeys` reads `k` alone, so a row holding `v, k` would put `v` where `k` is read.
+ val e = intercept[SparkException](evenKeys.eval(InternalRow(3L, 4L)))
+ assert(e.getCondition == "INTERNAL_ERROR", s"got ${e.getCondition}")
+ }
+
+ test("a filter over a column outside the scan output is an internal error") {
+ val e = intercept[SparkException](StorageFilterImpl(EqualTo(k, Literal(1L)), Seq(v)))
+ assert(e.getCondition == "INTERNAL_ERROR", s"got ${e.getCondition}")
+ }
+
+ test("select writes the kept rows' positions in ascending order") {
+ withBatchOf(3L, 4L, 5L, 6L, 8L) { batch =>
+ val selected = new Array[Int](batch.numRows)
+ val count = evenKeys.select(batch, selected)
+ assert(selected.take(count).toSeq == Seq(1, 3, 4))
+ }
+ }
+
+ test("an error evaluating the filter answers that it is not applied") {
+ val filter = positiveInverse
+ assert(filter.eval(InternalRow(2L)))
+ val e = intercept[StorageFilterNotAppliedException](filter.eval(InternalRow(0L)))
+ assert(e.getCondition == "STORAGE_FILTER_NOT_APPLIED", s"got ${e.getCondition}")
+ assert(e.getCause.isInstanceOf[SparkArithmeticException], s"got ${e.getCause}")
+ withBatchOf(1L, 0L) { batch =>
+ intercept[StorageFilterNotAppliedException](filter.select(batch, new Array[Int](2)))
+ }
+ }
+
+ test("a reader factory without the storage-filter methods cannot create a filtering reader") {
+ val factory = new PartitionReaderFactory {
+ override def createReader(partition: InputPartition): PartitionReader[InternalRow] = null
+ }
+ val request = StorageFilterRequestImpl(evenKeys, appendsCheckedColumn = false)
+ Seq(
+ () => factory.createReader(RangeInputPartition(0, 1), request),
+ () => factory.createColumnarReader(RangeInputPartition(0, 1), request)
+ ).foreach { create =>
+ val e = intercept[SparkUnsupportedOperationException](create())
+ assert(e.getCondition == "DATA_SOURCE_STORAGE_FILTER_READER_NOT_SUPPORTED",
Review Comment:
**5. [Nit] The SQLSTATEs of the two new conditions, and the `class` parameter, are checked nowhere**
This line and L98 compare only `e.getCondition`, and no other test raises either condition, since the test readers catch `StorageFilterNotAppliedException` (StorageFilterTestSource.scala:372 and 457). So the defaults at PartitionReaderFactory.java:70 and 84 could pass another class as `class`, e.g. `partition.getClass().getName()`, or this condition's `0A000` could become the `KD000` of the entry above it, and every test would still pass. The same holds for the `KD000` the description gives `STORAGE_FILTER_NOT_APPLIED`. As in item 10 of my second review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4380305050)), could both use `checkError`?
```scala
checkError(
exception = intercept[SparkUnsupportedOperationException](create()),
condition = "DATA_SOURCE_STORAGE_FILTER_READER_NOT_SUPPORTED",
sqlState = Some("0A000"),
parameters = Map("class" -> factory.getClass.getName))
```
and `checkError(exception = e, condition = "STORAGE_FILTER_NOT_APPLIED", sqlState = Some("KD000"))` at L98.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -124,9 +124,10 @@ case class GroupPartitionsExec(
* where this check passes, and `UnknownPartitioning`, where re-deriving reports the same thing.
*
* The comparison is by reference for a child that reports the same object, which the columnar and
- * codegen wrappers do. `PlanAdaptiveDynamicPruningFilters` rebuilds `BatchScanExec`, so there it
- * compares the partition keys by value, and reads equal only while the connector returns the same
- * input partitions for the same scan.
+ * codegen wrappers do. `PlanAdaptiveDynamicPruningFilters` and `ReuseAdaptiveSubquery` rebuild a
+ * `BatchScanExec` whose runtime or storage filters hold a subquery, so there it compares the
+ * partition keys by value, and reads equal only while the connector returns the same input
+ * partitions for the same scan.
Review Comment:
**6. [Nit] Without AQE, `ReuseExchangeAndSubquery` rebuilds the scan too**
This sentence, from the answer to item 7 of my first review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009038)), names only AQE rules. Without AQE, `ReuseExchangeAndSubquery` runs after `EnsureRequirements` (QueryExecution.scala:848 and 886). It walks the plan bottom-up and keeps the first of two equal subqueries, so a scan that holds one keeps it, but a second one in the same scan becomes a `ReusedSubqueryExec` (ReuseExchangeAndSubquery.scala:53-67), and the scan is copied under the `GroupPartitionsExec` that is already planned. A storage-filtered scan gets there in two ways: two blooms that `MergeSubplans` merged both sit in `storageFilters`, as in the two-key test, or, with dynamic partition pruning off, one bloom on a key the scan declares for both `SupportsRuntimeV2Filtering` and `SupportsStorageFilters` sits in `runtimeFilters` and `storageFilters`. Then `checkChildStillMatches` compares the partition keys by value too. Maybe:
```suggestion
* codegen wrappers do. A rule that rewrites a subquery in the runtime or storage filters of a
* `BatchScanExec` after `EnsureRequirements` rebuilds it, so there it compares the partition
* keys by value, and reads equal only while the connector returns the same input partitions for
* the same scan. Such rules are `PlanAdaptiveDynamicPruningFilters` and `ReuseAdaptiveSubquery`
* with AQE, and `ReuseExchangeAndSubquery` without it.
```
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -1060,6 +1061,46 @@ class ParquetStorageFilterSuite extends QueryTest with SharedSparkSession
}
}
+ Seq(false, true).foreach { aqe =>
+ test(s"a two-key join ships each bloom to the reader once (AQE $aqe)") {
+ // `MergeSubplans` merges the two blooms' aggregates into one subquery returning both blooms.
+ withTable("bf3", "bf4") {
+ spark.range(600).selectExpr("id AS k", "id * 10 AS v", "id * 100 AS w")
+ .write.format("parquet").saveAsTable("bf3")
+ spark.range(30).selectExpr("id AS k", "id * 10 AS v", "id AS flag")
+ .write.format("parquet").saveAsTable("bf4")
+ withSQLConf(
+ SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val df = spark.sql(
+ "SELECT bf3.k, bf3.v, bf3.w FROM bf3 JOIN bf4 ON bf3.k = bf4.k AND bf3.v = bf4.v " +
+ "WHERE bf4.flag = 5")
+ val rows = df.collect().map(r => (r.getLong(0), r.getLong(1), r.getLong(2))).toSeq
+ assert(rows == Seq((5L, 50L, 500L)), s"got $rows")
+ val plan = df.queryExecution.executedPlan
+ val scan = collectFirst(plan) {
+ case s: FileSourceScanExec if s.storageFilters.nonEmpty => s
+ }.getOrElse(fail(s"no scan with storage filters:\n$plan"))
+ assert(scan.storageFilters.size == 2 && scan.storageFilters.forall(_.exists {
+ case GetStructField(_: ExecSubqueryExpression, _, _) => true
+ case _ => false
+ }), s"test setup: both blooms must come from one merged subquery:\n$plan")
+ val bloomBytes = scan.storageFilters.map(_.collectFirst {
+ case m: BloomFilterMightContain =>
+ m.bloomFilterExpression.eval().asInstanceOf[Array[Byte]].length
+ }.get).sum
+ // The reader's function ships with the scan's RDD, and the blooms with it.
+ val shipped = Utils.serialize(scan.inputRDD).length
+ assert(shipped < bloomBytes * 3 / 2,
Review Comment:
**4. [Nit] The new V1 two-key test passes when no bloom reaches the reader**
Thank you for adding it for item 3 of my second review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4379088695)). It bounds the shipped bytes only from above, so it also passes when the reader gets no filter at all: if `storageFilterOpt` were `None` on this path (ParquetFileFormat.scala:294-304), `shipped` would drop to a few KB. The answer (L1081) and the planning check (L1086-1089) hold either way, and so does this bound when the reader gets the blooms but gives them up. As the test above puts it, "The metrics are what say it did" (L1055-1058). Could this one also assert `rowsExcludedByStorageFilters(plan) > 0` (L988-994)?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsStorageFilters.java:
##########
@@ -0,0 +1,64 @@
+/*
+ * 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.read;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link Scan}. A data source implements it if its readers can evaluate a
+ * {@link StorageFilter}, to skip reading the other columns of the rows the filter rejects. The
+ * readers get the filter in the task, through
+ * {@link PartitionReaderFactory#createReader(InputPartition, StorageFilterRequest)} or
+ * {@link PartitionReaderFactory#createColumnarReader(InputPartition, StorageFilterRequest)}, which
+ * the scan's reader factory must implement.
Review Comment:
**12. [Nit] `FilePartitionReaderFactory` has no route for a request**
Spark's own base for file-based scans, `FilePartitionReaderFactory`, implements only the plain `createReader` and `createColumnarReader` (FilePartitionReaderFactory.scala:29-39). So a `FileScan` that adds `SupportsStorageFilters` fails every filtered task with `DATA_SOURCE_STORAGE_FILTER_READER_NOT_SUPPORTED` unless its factory wraps `FilePartitionReader` itself. That is what this sentence asks for, and no built-in `FileScan` implements the mix-in, so nothing fails today. But an adopter such as the V2 Parquet scan would also have to place the checked column after the partition values, which `PartitionReaderWithPartitionValues` appends last, and the metrics that L47-48 point to would not reach Spark, since `FilePartitionReader`, `PartitionedFileReader` and `PartitionReaderWithPartitionValues` forward only `next`, `get` and `close`. That last part is the same on master. Could the description's list of possible next steps name the V2 file sources?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/StorageFilterRequest.java:
##########
@@ -0,0 +1,51 @@
+/*
+ * 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.read;
+
+import java.io.Serializable;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * What Spark hands a reader of a {@link SupportsStorageFilters} scan, through
+ * {@link PartitionReaderFactory#createReader(InputPartition, StorageFilterRequest)} or
+ * {@link PartitionReaderFactory#createColumnarReader(InputPartition, StorageFilterRequest)}.
+ *
+ * @since 5.0.0
+ */
+@Evolving
+public interface StorageFilterRequest extends Serializable {
+
+ /** The filter the reader may apply. */
+ StorageFilter filter();
+
+ /**
+ * Whether the reader must append a boolean column after the columns of
+ * {@link Scan#readSchema()}. It is true only on a row the filter was evaluated on and kept. It is
+ * false on every other row the reader returns, such as a row read without the filter, or one
+ * kept although the filter rejected it. Spark skips the filter's predicate above the scan on the
+ * rows marked true. A reader asked for the column that returns a row or a batch without it fails
+ * the task. The reader may return more columns after the checked column, which Spark does not
+ * read.
Review Comment:
**2. [Low] The `appendsCheckedColumn` doc still promises a failure the new check no longer gives**
L42-43 say "A reader asked for the column that returns a row or a batch without it fails the task", and the new sentence after it lets the reader return more columns. Since the fix for item 1 of my second review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4379088674)), `DataSourceRDD` checks only that a row is wide enough (DataSourceRDD.scala:235), and that a batch is wide enough and holds a boolean at `width - 1` (DataSourceRDD.scala:229). A width check cannot tell the checked column from the reader's own first extra column, so a reader that leaves the column out but returns more columns passes it.
Take a scan that reads `k, v` (width 3), and a reader whose new overload forwards to its plain reader, the easy mistake of item 2 of my first review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009016)), returning `k, v, _pos: long, _deleted: boolean`. At `cb01a6f`, `4 != 3` failed the task with `DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING`, as documented. Now the row passes, and `If(checked, true, bloom)` reads ordinal 2. From an `UnsafeRow` that is a byte of `_pos`, so the `Filter` skips the bloom on arbitrary rows with no error. From a `GenericInternalRow` it is a `ClassCastException` that names neither the connector nor the request. A columnar `k, v, deleted: boolean, pos` passes too, with `deleted` as the mark. The join still compares the keys, so a bloom's answer stays right. Could the doc say what Spark checks, e.g.
```suggestion
* rows marked true. The reader may return more columns after the checked column, which Spark
* does not read. Spark checks only that a row or a batch is wide enough for the column, and that
* a batch holds a boolean there, and fails the task otherwise. A reader that leaves the column
* out but returns more columns can pass that check, and Spark then reads the column at that
* position as the checked column.
```
The description's "fails the task with `DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING`" could say the same.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterTestSource.scala:
##########
@@ -0,0 +1,492 @@
+/*
+ * 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.util
+import java.util.OptionalLong
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow}
+import org.apache.spark.sql.connector.catalog.{MetadataColumn, SupportsMetadataColumns, SupportsRead, Table, TableCapability, TableProvider}
+import org.apache.spark.sql.connector.expressions.{Expressions, FieldReference, NamedReference, SortOrder, Transform}
+import org.apache.spark.sql.connector.expressions.filter.Predicate
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.sql.types.{BooleanType, DataType, LongType, StructField, StructType, VariantType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.util.Utils
+
+/**
+ * A read-only v2 source over a fixed table of longs that implements [[SupportsStorageFilters]]:
+ * the key `k`, the value `v` holding `k * 10`, and with `wide` also `w` holding `k * 100`. Its
+ * reader gets the filter in the task and drops the rows the filter rejects. It reads the columns
+ * Spark requires, and builds the filter's row from the filter's `references()`. Its options:
+ * - `columnar`: reads in batches, through `StorageFilter.select`, rather than row by row;
+ * - `giveUp`: reads without the filter, as a reader does for a unit it cannot use the filter on.
+ * Here the unit is the partition;
+ * - `markAllChecked`: reads without the filter but marks every row checked, which breaks the
+ * contract on purpose;
+ * - `dropChecked`: `true` never appends the checked column, and `later` appends it to the first
+ * row or batch only. Either breaks the contract on purpose;
+ * - `checkedLayout`: `front` returns the checked column ahead of the scan's columns, which breaks
+ * the contract on purpose. `extra` returns one more column after it;
+ * - `declared`: the comma-separated storage-filter columns, `k` by default, a dot separating a
+ * nested field;
+ * - `valueName`: the name of the value column, `v` by default;
+ * - `wide`: adds the column `w`;
+ * - `variant`: adds the variant column `j`, which the scan reads only through pushed extractions
+ * of a long. Each one holds `k * 10`;
+ * - `numPartitions`: splits the keys into that many input partitions, 1 by default;
+ * - `metadataColumn`: adds the metadata column `_pos`, which holds the key;
+ * - `runtimeFiltering`: `v2` makes the scan a `SupportsRuntimeV2Filtering` one over every column,
+ * with iterative pushdown on a table partitioned by `k`. `catalyst` makes it a
+ * `SupportsRuntimeCatalystFiltering` one that evaluates the filters on `k` fully. Neither
+ * applies the filters it takes, so they serve planning tests only;
+ * - `push`: the comma-separated pushed operations it takes whole without applying them, of
+ * `limit`, `offset`, `topN` and `sample`, for planning tests only.
+ *
+ * A filter that throws `StorageFilterNotAppliedException` is given up for the rest of the
+ * partition too. The counters below record what the readers got, evaluated and returned.
+ */
+class StorageFilterTestSource extends TableProvider {
+
+ override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+ val schema = new StructType().add("k", LongType).add(options.getOrDefault("valueName", "v"),
+ LongType)
+ val wide = if (options.getBoolean("wide", false)) schema.add("w", LongType) else schema
+ if (options.getBoolean("variant", false)) wide.add("j", VariantType) else wide
+ }
+
+ override def getTable(
+ schema: StructType,
+ partitioning: Array[Transform],
+ properties: util.Map[String, String]): Table = {
+ val options = new CaseInsensitiveStringMap(properties)
+ if (options.getBoolean("metadataColumn", false)) {
+ new StorageFilterTestTable(schema, options) with SupportsMetadataColumns {
+ override def metadataColumns(): Array[MetadataColumn] = Array(new MetadataColumn {
+ override def name(): String = StorageFilterTestSource.posName
+ override def dataType(): DataType = LongType
+ })
+ }
+ } else {
+ new StorageFilterTestTable(schema, options)
+ }
+ }
+}
+
+object StorageFilterTestSource {
+
+ /** Rows the source serves, `k` from 0 until numRows. */
+ val numRows: Int = 200
+
+ /** Rows a batch holds at most in a columnar read. */
+ val batchSize: Int = 64
+
+ /** The name of the metadata column the `metadataColumn` option adds. */
+ val posName: String = "_pos"
+
+ /** How many readers were created with a storage-filter request. */
+ val filteredReaders = new AtomicInteger(0)
+
+ /** How many of those requests asked for the checked column. */
+ val checkedColumnRequests = new AtomicInteger(0)
+
+ /** Rows the reader evaluated the storage filter against. */
+ val rowsEvaluated = new AtomicLong(0)
+
+ /** Rows the reader returned. */
+ val rowsReturned = new AtomicLong(0)
+
+ /** The serialized size of the first filter a reader got, before it evaluated the filter. */
+ val filterBytes = new AtomicLong(0)
+
+ /** The references of the filters the readers got, and the identities of those filters. */
+ val filterReferences = mutable.Set.empty[Seq[String]]
+ val filterIdentities = mutable.Set.empty[Int]
+
+ /** How many times a runtime-filtering scan was handed filters. */
+ val runtimeFilterCalls = new AtomicInteger(0)
+
+ def reset(): Unit = {
+ filteredReaders.set(0)
+ checkedColumnRequests.set(0)
+ rowsEvaluated.set(0)
+ rowsReturned.set(0)
+ filterBytes.set(0)
+ filterReferences.synchronized(filterReferences.clear())
+ filterIdentities.synchronized(filterIdentities.clear())
+ runtimeFilterCalls.set(0)
+ }
+
+ /** The value of column `name` in the row of `key`. Any column but `k`, `w` and `_pos` is `v`. */
+ def valueOf(name: String, key: Long): Long = name match {
+ case "k" | `posName` => key
+ case "w" => key * 100
+ case _ => key * 10
+ }
+
+ /** The value of `field` in the row of `key`, a struct for the extractions of a variant. */
+ def valueOf(field: StructField, key: Long): Any = field.dataType match {
+ case struct: StructType => InternalRow.fromSeq(struct.map(_ => valueOf(field.name, key)))
+ case _ => valueOf(field.name, key)
+ }
+
+ /** The columns of the row the filter evaluates, in that order. */
+ def columnsOf(filter: StorageFilter): Seq[String] = filter.references.map(_.fieldNames.head).toSeq
+
+ // Records the filter a reader got, before the reader evaluates it. Only the first reader
+ // serializes it, since a later reader of the same task may get it evaluated, see `StorageFilter`.
+ private[connector] def record(filter: StorageFilter): Unit = {
+ if (filterBytes.get() == 0) filterBytes.compareAndSet(0, Utils.serialize(filter).length)
+ filterReferences.synchronized(filterReferences += columnsOf(filter))
+ filterIdentities.synchronized(filterIdentities += System.identityHashCode(filter))
+ }
+}
+
+class StorageFilterTestTable(tableSchema: StructType, tableOptions: CaseInsensitiveStringMap)
+ extends Table with SupportsRead {
+
+ override def name(): String = "storage_filter_test"
+
+ override def schema(): StructType = tableSchema
+
+ override def partitioning(): Array[Transform] =
+ if (tableOptions.getOrDefault("runtimeFiltering", "") == "v2") {
+ Array(Expressions.identity("k"))
+ } else {
+ Array.empty
+ }
+
+ override def capabilities(): util.Set[TableCapability] =
+ Set(TableCapability.BATCH_READ).asJava
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder =
+ new StorageFilterTestScanBuilder(tableSchema, options)
+}
+
+/**
+ * Prunes the columns Spark does not read, and takes the pushed operations named in `push` whole.
+ * Its scan does not apply them, so they serve planning tests only. It takes every variant
+ * extraction, and reads each into a struct field of the variant's column named by its position.
+ */
+class StorageFilterTestScanBuilder(tableSchema: StructType, options: CaseInsensitiveStringMap)
+ extends SupportsPushDownRequiredColumns
+ with SupportsPushDownLimit
+ with SupportsPushDownOffset
+ with SupportsPushDownTopN
+ with SupportsPushDownTableSample
+ with SupportsPushDownVariantExtractions {
+
+ private val pushes = Utils.stringToSeq(options.getOrDefault("push", "")).toSet
+ private var required = tableSchema
+ private var extractions = Array.empty[VariantExtraction]
+
+ override def pruneColumns(requiredSchema: StructType): Unit = required = requiredSchema
+
+ override def pushVariantExtractions(pushed: Array[VariantExtraction]): Array[Boolean] = {
+ extractions = pushed
+ pushed.map(_ => true)
+ }
+
+ override def build(): Scan = {
+ val byColumn = extractions.groupBy(_.columnName().head)
+ val schema = StructType(required.map { field =>
+ byColumn.get(field.name).map { columnExtractions =>
+ field.copy(dataType = StructType(columnExtractions.zipWithIndex.map { case (e, i) =>
+ StructField(i.toString, e.expectedDataType(), nullable = true, e.metadata())
+ }))
+ }.getOrElse(field)
+ })
+ val factory = StorageFilterTestReaderFactory(
+ schema,
+ columnar = options.getBoolean("columnar", false),
+ giveUp = options.getBoolean("giveUp", false),
+ marksAllChecked = options.getBoolean("markAllChecked", false),
+ dropChecked = options.getOrDefault("dropChecked", ""),
+ checkedLayout = options.getOrDefault("checkedLayout", ""))
+ val declared = Utils.stringToSeq(options.getOrDefault("declared", "k"))
+ val numPartitions = options.getInt("numPartitions", 1)
+ options.getOrDefault("runtimeFiltering", "") match {
+ case "v2" => new StorageFilterTestScan(schema, factory, declared, numPartitions)
+ with SupportsRuntimeV2Filtering {
+ override def filterAttributes(): Array[NamedReference] =
+ readSchema().fieldNames.map(FieldReference.column).toArray[NamedReference]
+ override def filter(predicates: Array[Predicate]): Unit =
+ StorageFilterTestSource.runtimeFilterCalls.incrementAndGet()
+ override def supportsIterativePushdown(): Boolean = true
+ }
+ case "catalyst" => new StorageFilterTestScan(schema, factory, declared, numPartitions)
+ with SupportsRuntimeCatalystFiltering {
+ override def filterAttributes(): Array[NamedReference] = Array(FieldReference.column("k"))
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes()
+ override def filter(expressions: Array[Expression]): Unit =
+ StorageFilterTestSource.runtimeFilterCalls.incrementAndGet()
+ }
+ case _ => new StorageFilterTestScan(schema, factory, declared, numPartitions)
+ }
+ }
+
+ override def pushLimit(limit: Int): Boolean = pushes("limit")
+ override def pushOffset(offset: Int): Boolean = pushes("offset")
+ override def pushTopN(orders: Array[SortOrder], limit: Int): Boolean = pushes("topN")
+ override def pushTableSample(
+ lowerBound: Double,
+ upperBound: Double,
+ withReplacement: Boolean,
+ seed: Long,
+ sampleMethod: SampleMethod): Boolean = pushes("sample")
+ override def isPartiallyPushed: Boolean = false
+}
+
+class StorageFilterTestScan(
+ schema: StructType,
+ factory: StorageFilterTestReaderFactory,
+ declared: Seq[String],
+ numPartitions: Int)
+ extends Scan with Batch with SupportsReportStatistics with SupportsStorageFilters {
+
+ override def readSchema(): StructType = schema
+
+ override def toBatch: Batch = this
+
+ // Over the suite's application-side scan size threshold, so InjectRuntimeFilter considers this
+ // side worth filtering.
+ override def estimateStatistics(): Statistics = new Statistics {
+ override def sizeInBytes(): OptionalLong = OptionalLong.of(10L * 1024 * 1024)
+ override def numRows(): OptionalLong = OptionalLong.of(StorageFilterTestSource.numRows.toLong)
+ }
+
+ override def storageFilterAttributes(): Array[NamedReference] =
+ declared.map(name => FieldReference(name.split('.').toSeq)).toArray[NamedReference]
+
+ override def planInputPartitions(): Array[InputPartition] = {
+ val n = StorageFilterTestSource.numRows
+ (0 until numPartitions).map { i =>
+ RangeInputPartition(i * n / numPartitions, (i + 1) * n / numPartitions)
+ }.toArray
+ }
+
+ override def createReaderFactory(): PartitionReaderFactory = factory
+}
+
+case class StorageFilterTestReaderFactory(
+ schema: StructType,
+ columnar: Boolean,
+ giveUp: Boolean,
+ marksAllChecked: Boolean,
+ dropChecked: String,
+ checkedLayout: String)
+ extends PartitionReaderFactory {
+
+ override def supportColumnarReads(partition: InputPartition): Boolean = columnar
+
+ // The filter the reader applies, none when it gives the filter up or only marks the rows.
+ private def filterOf(request: StorageFilterRequest): Option[StorageFilter] = {
+ StorageFilterTestSource.filteredReaders.incrementAndGet()
+ StorageFilterTestSource.record(request.filter())
+ if (request.appendsCheckedColumn()) {
+ StorageFilterTestSource.checkedColumnRequests.incrementAndGet()
+ }
+ if (giveUp || marksAllChecked) None else Some(request.filter())
+ }
+
+ override def createReader(partition: InputPartition): PartitionReader[InternalRow] =
+ new StorageFilterTestReader(keys(partition), schema, None, appendsChecked = false)
+
+ override def createReader(
+ partition: InputPartition,
+ request: StorageFilterRequest): PartitionReader[InternalRow] =
+ new StorageFilterTestReader(keys(partition), schema, filterOf(request),
+ request.appendsCheckedColumn() && dropChecked != "true", marksAllChecked,
+ dropsCheckedLater = dropChecked == "later", checkedLayout)
+
+ override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] =
+ new StorageFilterTestColumnarReader(keys(partition), schema, None, appendsChecked = false)
+
+ override def createColumnarReader(
+ partition: InputPartition,
+ request: StorageFilterRequest): PartitionReader[ColumnarBatch] =
+ new StorageFilterTestColumnarReader(keys(partition), schema, filterOf(request),
+ request.appendsCheckedColumn() && dropChecked != "true", marksAllChecked,
+ dropsCheckedLater = dropChecked == "later", checkedLayout)
+
+ // The keys from `start` until `end`.
+ private def keys(partition: InputPartition): RangeInputPartition =
+ partition.asInstanceOf[RangeInputPartition]
+}
+
+/**
+ * Serves the rows of `partition` row by row, holding the columns of `schema`, and drops any row
+ * `filter` rejects, with the checked column last when `appendsChecked`. The filter's row holds
+ * the values of its `references()`. `filter` is cleared when it cannot be applied.
+ * `marksAllChecked` marks every row checked. `dropsCheckedLater` appends the checked column to
+ * the first row only. `checkedLayout` places the checked column as the source's option says.
+ */
+class StorageFilterTestReader(
+ partition: RangeInputPartition,
+ schema: StructType,
+ private var filter: Option[StorageFilter],
+ appendsChecked: Boolean,
+ marksAllChecked: Boolean = false,
+ dropsCheckedLater: Boolean = false,
+ checkedLayout: String = "")
+ extends PartitionReader[InternalRow] {
+
+ private val referenceNames = filter.map(StorageFilterTestSource.columnsOf).getOrElse(Nil)
+ private var current: Long = partition.start - 1L
+ private var row: InternalRow = _
+ private var returnedAny = false
+
+ override def next(): Boolean = {
+ var found = false
+ while (!found && current + 1 < partition.end) {
+ current += 1
+ val key = current
+ val passes = filter.forall { f =>
+ StorageFilterTestSource.rowsEvaluated.incrementAndGet()
+ val keyRow = new GenericInternalRow(
+ referenceNames.map(StorageFilterTestSource.valueOf(_, key)).toArray[Any])
+ try f.eval(keyRow) catch {
+ case _: StorageFilterNotAppliedException =>
+ filter = None
+ true
+ }
+ }
+ if (passes) {
+ val values = schema.fields.map(StorageFilterTestSource.valueOf(_, key))
+ val checked = filter.isDefined || marksAllChecked
+ val withChecked = appendsChecked && !(dropsCheckedLater && returnedAny)
+ row = new GenericInternalRow(
+ if (!withChecked) {
+ values
+ } else {
+ checkedLayout match {
+ case "front" => checked +: values
Review Comment:
**11. [Nit] The row reader's `front` layout runs in no test**
Only the columnar test sets `checkedLayout -> front` (StorageFilterPushDownSuite.scala:259), while the option's doc (L52-53) and this reader's doc (L345) offer it to both readers. A row-based `front` read would not end in the contract error either: `checkRow` checks only the width, and the `Filter` above would fail with a `ClassCastException`, as the comment at StorageFilterPushDownSuite.scala:255-256 suggests. Could the docs say `front` is for columnar reads, or this reader reject it?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterPredicate.scala:
##########
@@ -0,0 +1,52 @@
+/*
+ * 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.internal.connector
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BasePredicate, BindReferences, Expression, Predicate}
+
+/**
+ * The predicate of a storage filter, built on first use from `boundExpression`, since a generated
+ * one can be awkward to serialize. The file-source Parquet reader and [[StorageFilterImpl]] share
+ * it.
+ */
+class StorageFilterPredicate(private var boundExpression: Expression) extends Serializable {
+
+ @transient lazy val prepared: BasePredicate = {
+ // The expression is released below, so a copy made after a first evaluation would have nothing
+ // to build from. That is named here rather than left to an NPE.
+ if (boundExpression == null) {
+ throw SparkException.internalError("A storage filter was copied after an evaluation " +
Review Comment:
**7. [Nit] The copied-filter error no longer names the filter**
Before the move for item 12 of my second review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4380305098)), this was `s"$this was copied after an evaluation released the expression of its predicate"`, which names the columns, as item 20 of my first review suggested ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009086)). `StorageFilterPredicate` does not know its filter, so the message is now a constant, while the width error just before it in `StorageFilterImpl.eval` (StorageFilterImpl.scala:57-60) still names the filter. With two storage-filtered scans in a stage, a connector author who hits it cannot tell which filter a reader copied. The V1 message also lost its hint, "a storage filter must only be evaluated on an executor". Could `StorageFilterPredicate` take a short description for this message, e.g. `StorageFilterImpl.toString` and the V1 filter's key columns?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,103 @@
+/*
+ * 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.internal.connector
+
+import org.apache.spark.SparkException
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression => CatalystExpression}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the predicate is bound to
+ * it. The predicate is built on first use, see [[StorageFilterPredicate]].
+ */
+class StorageFilterImpl private (
+ private val predicate: StorageFilterPredicate,
+ private val attrs: Seq[Attribute])
Review Comment:
**9. [Nit] `attrs` is now kept only for its length and the column names**
Since `apply` binds the expression on the driver (L101), `attrs` is read only by `numReferences` (L41) and to rebuild the column names in the transient lazy `referenceArray` (L76-77). So each task deserializes full attributes for their names, and the class holds the layout twice. Could the constructor take the references instead, e.g. `private val referenceArray: Array[NamedReference]` with `numReferences = referenceArray.length`, and `apply` pass `referenced.map(a => FieldReference.column(a.name)).toArray`? `references()` keeps its clone. The one-caller `StorageFilterPredicate.apply` could then go too, as `ParquetStorageFilter` builds `new StorageFilterPredicate(remapped)` itself (ParquetStorageFilter.scala:185).
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,560 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GetStructField, GreaterThan, If, LessThanOrEqual, Literal, NumericEvalContext, PredicateHelper, Rand, XxHash64}
+import org.apache.spark.sql.catalyst.plans.logical.ReturnAnswer
+import org.apache.spark.sql.execution.{BaseLimitExec, CoalesceExec, ExecSubqueryExpression, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation}
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{BinaryType, DoubleType, LongType, StructType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper with PredicateHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
+ withSQLConf(bloomJoinConf ++ conf: _*) {
+ withTable("bf_build") {
+ spark.range(0, 30).selectExpr("id AS bk", "id AS bv")
+ .write.format("parquet").saveAsTable("bf_build")
+ body
+ }
+ }
+ }
+
+ // A join whose build side is small and carries a selective predicate, which is the shape
+ // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`.
+ private def bloomJoin(probe: DataFrame): DataFrame = {
+ val build = spark.table("bf_build").where("bv = 5")
+ probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*)
+ }
+
+ private def v2Probe(options: (String, String)*): DataFrame =
+ spark.read.format(sourceName).options(options.toMap).load()
+
+ private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME
+
+ private def scanOf(plan: SparkPlan): BatchScanExec =
+ collect(plan) { case b: BatchScanExec => b }.headOption
+ .getOrElse(fail(s"no BatchScanExec in plan:\n$plan"))
+
+ // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows
+ // the reader marks as checked.
+ private def bloomsAbove(plan: SparkPlan): (Int, Int) = {
+ val conditions = collect(plan) { case f: FilterExec => f.condition }
+ val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum
+ val skipped = conditions.map(_.collect {
+ case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain)
+ if a.name == checkedName => 1
+ }.sum).sum
+ (all - skipped, skipped)
+ }
+
+ // The scan offered nothing, and the Filter evaluates the bloom on every row.
+ private def assertNothingOffered(plan: SparkPlan): Unit = {
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.isEmpty, s"the scan must get no storage filter:\n$plan")
+ assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan")
+ }
+
+ // The rows the Filter on the scan returned, counted before the join.
+ private def rowsOutOfFilter(plan: SparkPlan): Long = {
+ val filters = collect(plan) {
+ case f: FilterExec if f.exists(_.isInstanceOf[BatchScanExec]) => f
+ }
+ assert(filters.size == 1, s"expected one Filter on the scan:\n$plan")
+ filters.head.metrics("numOutputRows").value
+ }
+
+ // Runs `df`, checks that it returns `expected`, and returns its plan.
+ private def checkJoin(df: DataFrame, expected: Set[Seq[Long]]): SparkPlan = {
+ val result = df.collect().map(_.toSeq.map(_.asInstanceOf[Long])).toSet
+ assert(result == expected, s"got $result")
+ df.queryExecution.executedPlan
+ }
+
+ // The build side keeps `bv = 5`, so the join returns that one row.
+ private val joinResult = Set(Seq(5L, 50L))
+
+ // Runs the bloom join over the v2 source read with `options`, checks its rows, and returns its
+ // plan.
+ private def runBloomJoin(options: (String, String)*): SparkPlan = {
+ StorageFilterTestSource.reset()
+ checkJoin(bloomJoin(v2Probe(options: _*)), joinResult)
+ }
+
+ for (columnar <- Seq(false, true); aqe <- Seq(false, true)) {
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind scan gets the checked column, and the Filter skips the bloom on checked rows " +
+ s"(AQE $aqe)") {
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val plan = runBloomJoin("columnar" -> columnar.toString)
+ val scan = scanOf(plan)
+ assert(scan.supportsColumnar == columnar, s"the scan must read $kind:\n$plan")
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(scan.appendsCheckedColumn, s"the scan must get the checked column:\n$plan")
+ assert(bloomsAbove(plan) == (0, 1),
+ s"the Filter must skip the bloom on checked rows:\n$plan")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(),
+ s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " +
+ s"${StorageFilterTestSource.rowsReturned.get()}")
+ val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output }
+ assert(probeSide.exists(!_.exists(_.name == checkedName)),
+ s"the Project above the scan must drop the checked column; got $probeSide")
+ assert(StorageFilterTestSource.filteredReaders.get() == 1,
+ s"the one partition's reader must get the filter; " +
+ s"got ${StorageFilterTestSource.filteredReaders.get()}")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 1,
+ "the reader must be asked for the checked column")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ s"the reader must evaluate every row; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(false, true).foreach { columnar =>
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind reader that gives the filter up marks its rows unchecked, and the Filter " +
+ "filters them") {
+ withBloomJoin() {
+ val plan = runBloomJoin("columnar" -> columnar.toString, "giveUp" -> "true")
+ assert(bloomsAbove(plan) == (0, 1), s"the plan is that of a filtering reader:\n$plan")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == 0,
+ s"the reader must evaluate nothing; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() == StorageFilterTestSource.numRows,
+ s"the reader must return every row; got ${StorageFilterTestSource.rowsReturned.get()}")
+ assert(rowsOutOfFilter(plan) < StorageFilterTestSource.numRows,
+ s"the Filter must apply the bloom to unchecked rows; got ${rowsOutOfFilter(plan)}")
+ }
+ }
+
+ test(s"the Filter takes a $kind reader's word on the rows it marks checked") {
+ // This reader breaks the contract on purpose: it applies no filter but marks every row
+ // checked. Only skipping the bloom on those rows keeps all of them.
+ withBloomJoin() {
+ // The join must still filter, which `runBloomJoin` checks.
+ val plan = runBloomJoin("columnar" -> columnar.toString, "markAllChecked" -> "true")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.numRows,
+ s"the Filter must keep every row marked checked; got ${rowsOutOfFilter(plan)}")
+ }
+ }
+
+ Seq(false, true).foreach { aqe =>
+ test(s"a $kind reader that leaves out the checked column it was asked for fails the task, " +
+ s"from any row or batch (AQE $aqe)") {
+ // A reader that gives the filter up returns every row, so it returns several batches.
+ Seq(Seq("dropChecked" -> "true"), Seq("dropChecked" -> "later", "giveUp" -> "true"))
+ .foreach { options =>
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ StorageFilterTestSource.reset()
+ val df = bloomJoin(v2Probe(options :+ ("columnar" -> columnar.toString): _*))
+ // `k`, `v` and the checked column.
+ checkError(
+ exception = intercept[SparkException](df.collect()),
+ condition = "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING",
+ sqlState = Some("KD000"),
+ parameters = Map(
+ "class" -> classOf[StorageFilterTestReaderFactory].getName,
+ "expected" -> "3",
+ "actual" -> "2"))
+ }
+ }
+ }
+ }
+
+ test(s"a $kind reader may return more columns after the checked column") {
+ withBloomJoin() {
+ val plan = runBloomJoin("columnar" -> columnar.toString, "checkedLayout" -> "extra")
+ assert(bloomsAbove(plan) == (0, 1),
+ s"the Filter must skip the bloom on checked rows:\n$plan")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(),
+ s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " +
+ s"${StorageFilterTestSource.rowsReturned.get()}")
+ }
+ }
+
+ test(s"a $kind test reader gives a filter up when it throws, and returns its rows unchecked") {
+ // `1 / k` under ANSI throws at k = 0, the partition's first row, so the reader gives the
+ // filter up before it keeps or drops anything. No runtime bloom throws, so the reader is
+ // driven by hand.
+ val k = AttributeReference("k", LongType)()
+ val inverse = Divide(Literal(1.0), Cast(k, DoubleType), NumericEvalContext(EvalMode.ANSI))
+ val filter = StorageFilterImpl(GreaterThan(inverse, Literal(0.0)), Seq(k))
+ val partition = RangeInputPartition(0, StorageFilterTestSource.numRows)
+ val schema = new StructType().add("k", LongType).add("v", LongType)
+ val rows = mutable.ArrayBuffer.empty[(Long, Boolean)]
+ if (columnar) {
+ val reader = new StorageFilterTestColumnarReader(
+ partition, schema, Some(filter), appendsChecked = true)
+ try {
+ while (reader.next()) {
+ val batch = reader.get()
+ (0 until batch.numRows()).foreach { i =>
+ rows += ((batch.column(0).getLong(i), batch.column(2).getBoolean(i)))
+ }
+ }
+ } finally {
+ reader.close()
+ }
+ } else {
+ val reader =
+ new StorageFilterTestReader(partition, schema, Some(filter), appendsChecked = true)
+ while (reader.next()) {
+ rows += ((reader.get().getLong(0), reader.get().getBoolean(2)))
+ }
+ }
+ assert(rows == (0L until StorageFilterTestSource.numRows.toLong).map(_ -> false),
+ s"expected every row, unchecked; got ${rows.take(5)} and ${rows.size} rows in all")
+ }
+ }
+
+ test("a columnar reader that returns the checked column ahead of the others fails the task") {
+ // The batch is as wide as asked for, but its last column is not a boolean. A row carries no
+ // types, so a row-based reader's mistake goes unnoticed.
+ withBloomJoin() {
+ StorageFilterTestSource.reset()
+ val df = bloomJoin(v2Probe("columnar" -> "true", "checkedLayout" -> "front"))
+ checkError(
+ exception = intercept[SparkException](df.collect()),
+ condition = "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING",
+ sqlState = Some("KD000"),
+ parameters = Map(
+ "class" -> classOf[StorageFilterTestReaderFactory].getName,
+ "expected" -> "3",
+ "actual" -> "3"))
+ }
+ }
+
+ // Under a limit of two fields, the checked column would make three.
+ Seq(
+ ("too wide for whole-stage codegen", SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> "2"),
+ ("without whole-stage codegen", SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false")
+ ).foreach { case (when, conf) =>
+ test(s"a scan $when gets the bloom without the checked column") {
+ withBloomJoin(conf) {
+ val plan = runBloomJoin("columnar" -> "true")
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(!scan.appendsCheckedColumn, s"the scan must not get the checked column:\n$plan")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 0,
+ "the reader must not be asked for the checked column")
+ assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(
+ ("with the conf off", Seq(SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "false"), Nil),
+ ("to a scan that declares no column", Nil, Seq("declared" -> "")),
+ ("to a scan that declares only another column", Nil, Seq("declared" -> "v")),
+ ("to a scan that declares only a nested field", Nil, Seq("declared" -> "k.x"))
+ ).foreach { case (when, conf, options) =>
+ test(s"nothing is offered $when") {
+ withBloomJoin(conf: _*) {
+ val plan = runBloomJoin(options :+ ("columnar" -> "true"): _*)
+ assertNothingOffered(plan)
+ assert(StorageFilterTestSource.filteredReaders.get() == 0,
+ "no reader may get a storage filter")
+ }
+ }
+ }
+
+ test("a scan that reads only the key and a metadata column is offered nothing") {
+ // Only the data columns count as columns left to skip, as on the file-source side.
+ withBloomJoin() {
+ val source = v2Probe("columnar" -> "true", "metadataColumn" -> "true")
+ val pos = StorageFilterTestSource.posName
+ StorageFilterTestSource.reset()
+ assertNothingOffered(
+ checkJoin(bloomJoin(source.select(source("k"), source.metadataColumn(pos))),
+ Set(Seq(5L, 5L))))
+ StorageFilterTestSource.reset()
+ val withValue = checkJoin(
+ bloomJoin(source.select(source("k"), source("v"), source.metadataColumn(pos))),
+ Set(Seq(5L, 50L, 5L)))
+ assert(scanOf(withValue).storageFilters.size == 1,
+ s"with `v` left to skip, the bloom must be offered:\n$withValue")
+ }
+ }
+
+ test("a variant scan that reads only its key column and a metadata column is offered nothing") {
+ // The key is an extraction of the declared variant column, so that column is the scan's only
+ // data column, and the scan's output has fresh exprIds.
+ withBloomJoin() {
+ val pos = StorageFilterTestSource.posName
+ val probe = v2Probe("columnar" -> "true", "variant" -> "true", "metadataColumn" -> "true",
+ "declared" -> "j").selectExpr("variant_get(j, '$.a', 'long') AS a", pos)
+ val build = spark.table("bf_build").where("bv = 5")
+ val df = probe.join(build, probe("a") === build("bk") * 10).select(probe("a"), probe(pos))
+ StorageFilterTestSource.reset()
+ assertNothingOffered(checkJoin(df, Set(Seq(50L, 5L))))
+ }
+ }
+
+ test("a scan that took a pushed variant extraction gets the bloom") {
+ // The extraction keeps the scan's rows.
+ withBloomJoin() {
+ StorageFilterTestSource.reset()
+ val probe = v2Probe("columnar" -> "true", "variant" -> "true")
+ .selectExpr("k", "variant_get(j, '$.a', 'long') AS a")
+ val plan = checkJoin(bloomJoin(probe), joinResult)
+ val scan = scanOf(plan)
+ assert(scan.output.exists(_.dataType.isInstanceOf[StructType]),
+ s"test setup: the scan must read the extraction:\n$plan")
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+
+ for (aqe <- Seq(false, true)) {
+ test("a scan with a column named like the checked column gets the filter, and runtime " +
+ s"filtering still resolves that column (AQE $aqe)") {
+ // Runtime filtering resolves the declared filter attributes by name, which the checked column
+ // the scan appends would make ambiguous. Dynamic partition pruning on `k` would take the
+ // bloom's place.
+ withBloomJoin(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "false",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val plan = runBloomJoin(
+ "columnar" -> "true", "valueName" -> checkedName, "runtimeFiltering" -> "v2")
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.size == 1 && scan.appendsCheckedColumn,
+ s"the scan must get the bloom and the checked column:\n$plan")
+ assert(StorageFilterTestSource.runtimeFilterCalls.get() > 0,
+ "test setup: the bloom must reach runtime filtering")
+ val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output }
+ assert(probeSide.contains(scan.output.init),
+ s"the Project above the scan must keep the table's column; got $probeSide")
+ }
+ }
+
+ test(s"a bloom the scan fully evaluates through runtime filtering is not offered (AQE $aqe)") {
+ // Dynamic partition pruning on `k` would take the bloom's place.
+ withBloomJoin(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "false",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val plan = runBloomJoin("columnar" -> "true", "runtimeFiltering" -> "catalyst")
+ val scan = scanOf(plan)
+ assert(scan.runtimeFilters.exists(_.exists(_.isInstanceOf[BloomFilterMightContain])) &&
+ StorageFilterTestSource.runtimeFilterCalls.get() > 0,
+ s"test setup: the scan must take the bloom as a runtime filter:\n$plan")
+ assert(scan.storageFilters.isEmpty, s"the scan must get no storage filter:\n$plan")
+ assert(bloomsAbove(plan) == (0, 0), s"no Filter may keep the bloom:\n$plan")
+ }
+ }
+
+ test(s"a two-key join offers both blooms in one filter, each bloom with its own bytes " +
+ s"(AQE $aqe)") {
+ // `MergeSubplans` merges the two blooms' aggregates into one subquery returning both blooms.
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ withTable("bf_build2") {
+ spark.range(0, 30).selectExpr("id AS bk", "id * 10 AS bv", "id AS flag")
+ .write.format("parquet").saveAsTable("bf_build2")
+ StorageFilterTestSource.reset()
+ val probe = v2Probe("columnar" -> "true", "wide" -> "true", "declared" -> "k,v")
+ .where("w % 3 != 1")
+ val build = spark.table("bf_build2").where("flag = 5")
+ val df = probe.join(build, probe("k") === build("bk") && probe("v") === build("bv"))
+ .select(probe("k"), probe("v"), probe("w"))
+ val plan = checkJoin(df, Set(Seq(5L, 50L, 500L)))
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.size == 2, s"both blooms must be on the scan:\n$plan")
+ assert(scan.storageFilters.forall(_.exists {
+ case GetStructField(_: ExecSubqueryExpression, _, _) => true
+ case _ => false
+ }), s"test setup: the blooms must come from one merged subquery:\n$plan")
+ assert(bloomsAbove(plan) == (0, 2), s"the Filter must skip both blooms:\n$plan")
+ val filter = collect(plan) {
+ case f: FilterExec if f.child.exists(_.isInstanceOf[BatchScanExec]) => f
+ }.head
+ assert(splitConjunctivePredicates(filter.condition).exists { c =>
+ c.references.map(_.name).toSet == Set("w") && !c.isInstanceOf[If]
+ }, s"the Filter must evaluate the conjunct over `w` on every row:\n$plan")
+ assert(StorageFilterTestSource.filterReferences == Set(Seq("k", "v")),
+ s"one filter over both keys; got ${StorageFilterTestSource.filterReferences}")
+ val bloomBytes = scan.storageFilters.map(_.collectFirst {
+ case m: BloomFilterMightContain =>
+ m.bloomFilterExpression.eval().asInstanceOf[Array[Byte]].length
+ }.get).sum
+ assert(StorageFilterTestSource.filterBytes.get() < bloomBytes * 3 / 2,
+ s"each bloom must ship once; the filter takes ${StorageFilterTestSource.filterBytes} " +
+ s"bytes for $bloomBytes bytes of blooms")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ "the reader must evaluate every row; got " +
+ StorageFilterTestSource.rowsEvaluated.get())
+ }
+ }
+ }
+
+ test(s"the readers of one task share the filter the task deserialized (AQE $aqe)") {
+ // `coalesce(1)` reads both partitions in one task.
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ StorageFilterTestSource.reset()
+ val probe = v2Probe("columnar" -> "true", "numPartitions" -> "2").coalesce(1)
+ val plan = checkJoin(bloomJoin(probe), joinResult)
+ assert(collect(plan) { case c: CoalesceExec => c }.nonEmpty, s"test setup:\n$plan")
+ assert(scanOf(plan).storageFilters.size == 1, s"the bloom must be on the scan:\n$plan")
+ assert(StorageFilterTestSource.filteredReaders.get() == 2,
+ s"both readers must get the filter; got ${StorageFilterTestSource.filteredReaders.get()}")
+ assert(StorageFilterTestSource.filterIdentities.size == 1,
+ "both readers must get the same instance")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ "the readers must evaluate every row; got " + StorageFilterTestSource.rowsEvaluated.get())
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the readers must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ test("the storage filter shows in the scan description and separates canonicalized plans") {
+ withBloomJoin() {
+ val plan = bloomJoin(v2Probe("columnar" -> "true")).queryExecution.executedPlan
+ val scan = scanOf(plan)
+ assert(scan.simpleString(100).contains("StorageFilters:"),
+ s"the scan description must name the filters: ${scan.simpleString(100)}")
+ assert(scan.verboseStringWithOperatorId().contains("StorageFilters:"),
+ s"EXPLAIN FORMATTED must name the filters: ${scan.verboseStringWithOperatorId()}")
+
+ intercept[IllegalArgumentException](scan.copy(storageFilters = Nil))
+ // As the scan would be without the checked column, too wide or without codegen.
+ val unchecked = scan.copy(output = scan.output.init, appendsCheckedColumn = false)
+ assert(scan != unchecked, "equals must see the checked column")
+ val plain = unchecked.copy(storageFilters = Nil)
+ assert(unchecked != plain && unchecked.hashCode != plain.hashCode,
+ "equals and hashCode must see storageFilters")
+ assert(!unchecked.sameResult(plain),
+ s"a filtering scan must not be reusable as a plain one:\n${unchecked.canonicalized}\n" +
+ s"${plain.canonicalized}")
+ assert(!plain.simpleString(100).contains("StorageFilters:"),
+ s"a plain scan's description must not name storage filters: ${plain.simpleString(100)}")
+ assert(!plain.verboseStringWithOperatorId().contains("StorageFilters:"),
+ s"nor its EXPLAIN FORMATTED: ${plain.verboseStringWithOperatorId()}")
+
+ // As another relation instance of the same table gets, with fresh exprIds.
+ val fresh = AttributeMap(scan.output.map(a => a -> a.newInstance()))
+ val twin = scan.copy(
+ output = scan.output.map(fresh(_)),
+ storageFilters = scan.storageFilters.map(_.transform {
+ case a: Attribute => fresh.getOrElse(a, a)
+ }))
+ assert(scan.sameResult(twin),
+ s"the same filtering scan must be reusable:\n${scan.canonicalized}\n${twin.canonicalized}")
+ }
+ }
+
+ // The reader would filter ahead of the operation, which then picks other rows.
+ Seq[(String, DataFrame => DataFrame)](
+ "limit" -> (_.limit(100)),
+ "offset" -> (_.offset(10)),
+ "topN" -> (_.orderBy("v").limit(100)),
+ "sample" -> (_.sample(0.5))
+ ).foreach { case (op, apply) =>
+ test(s"nothing is offered to a scan that took a pushed $op") {
+ withBloomJoin() {
+ val probe = apply(v2Probe("columnar" -> "true", "push" -> op))
+ val plan = bloomJoin(probe).queryExecution.executedPlan
+ assert(collect(plan) {
+ case p @ (_: BaseLimitExec | _: TakeOrderedAndProjectExec | _: SampleExec) => p
+ }.isEmpty, s"the $op must be pushed into the scan whole:\n$plan")
+ assertNothingOffered(plan)
+ }
+ }
+ }
+
+ test("the scan's pushdowns decide on storage filters, not whether it can be merged") {
+ // A pushed variant extraction, for one, keeps the scan's rows but cannot be merged.
+ withBloomJoin() {
+ val optimized = bloomJoin(v2Probe("columnar" -> "true")).queryExecution.optimizedPlan
+ def planned(mergeable: Boolean, rowWise: Boolean): BatchScanExec = scanOf(
+ spark.sessionState.planner.plan(ReturnAnswer(optimized.transform {
Review Comment:
**10. [Nit] `QueryExecution.createSparkPlan` already does this**
`spark.sessionState.planner.plan(ReturnAnswer(...)).next()` is the body of `QueryExecution.createSparkPlan` (QueryExecution.scala:923-929), which `DynamicPartitionPruningSuite` calls the same way. `QueryExecution.createSparkPlan(spark.sessionState.planner, optimized.transform { ... })` would also let the `ReturnAnswer` import (L26) go.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,103 @@
+/*
+ * 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.internal.connector
+
+import org.apache.spark.SparkException
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression => CatalystExpression}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.read.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the predicate is bound to
+ * it. The predicate is built on first use, see [[StorageFilterPredicate]].
+ */
+class StorageFilterImpl private (
+ private val predicate: StorageFilterPredicate,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes, read on every row.
+ private val numReferences = attrs.length
+
+ // Whether this instance has logged a failure, so a task logs only its first.
+ @transient private var loggedFailure = false
+
+ /**
+ * Wraps an error in [[StorageFilterNotAppliedException]], unless
+ * [[StorageFilterErrors.rethrowIfMustPropagate]] rethrows it, as the file-source reader does: a
+ * task kill, or a fatal error in its cause chain. Each task logs the first one. A row of another
+ * width than [[references]] is a reader bug, which fails the task, since the predicate would read
+ * the wrong values.
+ */
+ @throws[StorageFilterNotAppliedException]
+ override def eval(row: InternalRow): Boolean = {
+ // Outside the `try`, so a failure to build the predicate, or a row of another width, fails the
+ // task.
+ if (row.numFields != numReferences) {
+ throw SparkException.internalError(
+ s"$this takes rows of $numReferences values, but got ${row.numFields}")
+ }
+ val prepared = predicate.prepared
Review Comment:
**8. [Nit] `eval` resolves the predicate's lazy val on every row**
`predicate.prepared` is a `@transient lazy val` (StorageFilterPredicate.scala:30), so every call reads its volatile initialization flag, also after the predicate is built. The V1 reader moves the same lazy val out of its loop on purpose, "since resolving it is a volatile read" (LateMaterializationParquetRecordReader.java:731-732), and since item 12 of my second review both resolve the same `StorageFilterPredicate`. A V2 reader cannot do that, since it only calls `eval`, and the default `select` calls `eval` for every row. The cost is small next to a bloom probe, but a plain field keeps V2 on V1's rule, as one thread uses the filter (StorageFilter.java:48-49):
```scala
@transient private var built: BasePredicate = _
...
if (built == null) built = predicate.prepared
try {
built.eval(row)
```
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -2397,6 +2397,18 @@
],
"sqlState" : "0A000"
},
+ "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING" : {
+ "message" : [
+ "A reader of <class> was asked for the column that marks the rows a storage filter checked, at position <expected>. That column is missing or not a boolean. The reader returned <actual> columns."
Review Comment:
**3. [Low] The new message calls the width a position, and its parameters keep their count names**
Since `5141077` the message says the column was asked for `at position <expected>`, but `<expected>` is still the width `BatchScanExec` passes (`output.length`, BatchScanExec.scala:142). So it is a 1-based position, and the message does not say so. A connector author counts from 0, as `ColumnarBatch.column(int)` and `InternalRow` ordinals do, and this PR's tests read the checked column of `k, v, checked` as `batch.column(2)` (StorageFilterPushDownSuite.scala:236). For the `front` layout the task fails with "... at position 3. That column is missing or not a boolean. The reader returned 3 columns." Read 0-based, that says a fourth column is missing, while the problem is the BIGINT at `column(2)`. The parameters also kept their names from when they were counts, so `getMessageParameters()`, the JSON error formats and the new tests report `expected=3, actual=3` for a failure.
The condition is new, so renaming it is free now and a contract once it is released. For example:
```
A reader of <class> was asked for the column that marks the rows a storage filter checked,
right after the <numScanColumns> columns of the scan's read schema. That column is missing or
not a boolean. The reader returned <numColumns> columns.
```
with `width - 1` passed as `numScanColumns` (DataSourceRDD.scala:238-240), and the same names in `storageFilterCheckedColumnMissingError` and in both `checkError` calls.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4381903677
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -127,6 +127,29 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
tableSpec.withNewLocation(newLoc)
}
+ /**
+ * The conjuncts of `filters` offered to the scan as storage filters, see
+ * `SupportsStorageFilters` and `StorageFilterPlanning.offeredConjuncts`. Only runtime blooms for
+ * now, as on the file-source side. The columns left to skip are the data columns, as there, so a
+ * scan that reads only the key and metadata columns is offered none.
+ */
+ private def storageFiltersFor(
+ relation: DataSourceV2ScanRelation,
+ filters: Seq[Expression]): Seq[Expression] = {
+ if (!session.sessionState.conf.v2StorageFilterPushdownEnabled) return Nil
+ if (relation.storageFilterAttrs.isEmpty) return Nil
Review Comment:
Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538), in that order.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -212,7 +236,8 @@ class DataSourceRDD(
private class PartitionIterator[T](
private var reader: PartitionReader[T],
rowCount: T => Int,
- taskState: TaskState) extends Iterator[T] {
+ taskState: TaskState,
+ check: Option[T => Unit] = None) extends Iterator[T] {
Review Comment:
Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538).
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,491 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GetStructField, GreaterThan, If, Literal, NumericEvalContext, PredicateHelper, Rand, XxHash64}
+import org.apache.spark.sql.catalyst.plans.logical.ReturnAnswer
+import org.apache.spark.sql.execution.{BaseLimitExec, CoalesceExec, ExecSubqueryExpression, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation}
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{BinaryType, DoubleType, LongType, StructType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper with PredicateHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
+ withSQLConf(bloomJoinConf ++ conf: _*) {
+ withTable("bf_build") {
+ spark.range(0, 30).selectExpr("id AS bk", "id AS bv")
+ .write.format("parquet").saveAsTable("bf_build")
+ body
+ }
+ }
+ }
+
+ // A join whose build side is small and carries a selective predicate, which is the shape
+ // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`.
+ private def bloomJoin(probe: DataFrame): DataFrame = {
+ val build = spark.table("bf_build").where("bv = 5")
+ probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*)
+ }
+
+ private def v2Probe(options: (String, String)*): DataFrame =
+ spark.read.format(sourceName).options(options.toMap).load()
+
+ private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME
+
+ private def scanOf(plan: SparkPlan): BatchScanExec =
+ collect(plan) { case b: BatchScanExec => b }.headOption
+ .getOrElse(fail(s"no BatchScanExec in plan:\n$plan"))
+
+ // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows
+ // the reader marks as checked.
+ private def bloomsAbove(plan: SparkPlan): (Int, Int) = {
+ val conditions = collect(plan) { case f: FilterExec => f.condition }
+ val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum
+ val skipped = conditions.map(_.collect {
+ case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain)
+ if a.name == checkedName => 1
+ }.sum).sum
+ (all - skipped, skipped)
+ }
+
+ // The scan offered nothing, and the Filter evaluates the bloom on every row.
+ private def assertNothingOffered(plan: SparkPlan): Unit = {
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.isEmpty && !scan.appendsCheckedColumn,
Review Comment:
Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538), in both places.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterTestSource.scala:
##########
@@ -0,0 +1,444 @@
+/*
+ * 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.{ByteArrayOutputStream, ObjectOutputStream}
+import java.util
+import java.util.OptionalLong
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow}
+import org.apache.spark.sql.connector.catalog.{MetadataColumn, SupportsMetadataColumns, SupportsRead, Table, TableCapability, TableProvider}
+import org.apache.spark.sql.connector.expressions.{Expressions, FieldReference, NamedReference, SortOrder, Transform}
+import org.apache.spark.sql.connector.expressions.filter.{Predicate, StorageFilter, StorageFilterNotAppliedException}
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.sql.types.{BooleanType, DataType, LongType, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+/**
+ * A read-only v2 source over a fixed table of longs that implements [[SupportsStorageFilters]]:
+ * the key `k`, the value `v` holding `k * 10`, and with `wide` also `w` holding `k * 100`. Its
+ * reader gets the filter in the task and drops the rows the filter rejects. It reads the columns
+ * Spark requires, and builds the filter's row from the filter's `references()`. Its options:
+ * - `columnar`: reads in batches, through `StorageFilter.select`, rather than row by row;
+ * - `giveUp`: reads without the filter, as a reader does for a unit it cannot use the filter on.
+ * Here the unit is the partition;
+ * - `markAllChecked`: reads without the filter but marks every row checked, which breaks the
+ * contract on purpose;
+ * - `dropChecked`: `true` never appends the checked column, and `later` appends it to the first
+ * row or batch only. Either breaks the contract on purpose;
+ * - `declared`: the comma-separated storage-filter columns, `k` by default, a dot separating a
+ * nested field;
+ * - `valueName`: the name of the value column, `v` by default;
+ * - `wide`: adds the column `w`;
+ * - `numPartitions`: splits the keys into that many input partitions, 1 by default;
+ * - `metadataColumn`: adds the metadata column `_pos`, which holds the key;
+ * - `runtimeFiltering`: `v2` makes the scan a `SupportsRuntimeV2Filtering` one over every column,
+ * with iterative pushdown on a table partitioned by `k`. `catalyst` makes it a
+ * `SupportsRuntimeCatalystFiltering` one that evaluates the filters on `k` fully. Neither
+ * applies the filters it takes, so they serve planning tests only;
+ * - `push`: the comma-separated pushed operations it takes whole without applying them, of
+ * `limit`, `offset`, `topN` and `sample`, for planning tests only.
+ *
+ * A filter that throws `StorageFilterNotAppliedException` is given up for the rest of the
+ * partition too. The counters below record what the readers got, evaluated and returned.
+ */
+class StorageFilterTestSource extends TableProvider {
+
+ override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+ val schema = new StructType().add("k", LongType).add(options.getOrDefault("valueName", "v"),
+ LongType)
+ if (options.getBoolean("wide", false)) schema.add("w", LongType) else schema
+ }
+
+ override def getTable(
+ schema: StructType,
+ partitioning: Array[Transform],
+ properties: util.Map[String, String]): Table = {
+ val options = new CaseInsensitiveStringMap(properties)
+ if (options.getBoolean("metadataColumn", false)) {
+ new StorageFilterTestTable(schema, options) with SupportsMetadataColumns {
+ override def metadataColumns(): Array[MetadataColumn] = Array(new MetadataColumn {
+ override def name(): String = StorageFilterTestSource.posName
+ override def dataType(): DataType = LongType
+ })
+ }
+ } else {
+ new StorageFilterTestTable(schema, options)
+ }
+ }
+}
+
+object StorageFilterTestSource {
+
+ /** Rows the source serves, `k` from 0 until numRows. */
+ val numRows: Int = 200
+
+ /** Rows a batch holds at most in a columnar read. */
+ val batchSize: Int = 64
+
+ /** The name of the metadata column the `metadataColumn` option adds. */
+ val posName: String = "_pos"
+
+ /** How many readers were created with a storage-filter request. */
+ val filteredReaders = new AtomicInteger(0)
+
+ /** How many of those requests asked for the checked column. */
+ val checkedColumnRequests = new AtomicInteger(0)
+
+ /** Rows the reader evaluated the storage filter against. */
+ val rowsEvaluated = new AtomicLong(0)
+
+ /** Rows the reader returned. */
+ val rowsReturned = new AtomicLong(0)
+
+ /** The serialized size of the last filter a reader got. */
+ val filterBytes = new AtomicLong(0)
+
+ /** The references of the filters the readers got, and the identities of those filters. */
+ val filterReferences = mutable.Set.empty[Seq[String]]
+ val filterIdentities = mutable.Set.empty[Int]
+
+ /** How many times a runtime-filtering scan was handed filters. */
+ val runtimeFilterCalls = new AtomicInteger(0)
+
+ def reset(): Unit = {
+ filteredReaders.set(0)
+ checkedColumnRequests.set(0)
+ rowsEvaluated.set(0)
+ rowsReturned.set(0)
+ filterBytes.set(0)
+ filterReferences.synchronized(filterReferences.clear())
+ filterIdentities.synchronized(filterIdentities.clear())
+ runtimeFilterCalls.set(0)
+ }
+
+ /** The value of column `name` in the row of `key`. Any column but `k`, `w` and `_pos` is `v`. */
+ def valueOf(name: String, key: Long): Long = name match {
+ case "k" | `posName` => key
+ case "w" => key * 100
+ case _ => key * 10
+ }
+
+ /** The columns of the row the filter evaluates, in that order. */
+ def columnsOf(filter: StorageFilter): Seq[String] = filter.references.map(_.fieldNames.head).toSeq
+
+ // Records the filter a reader got, before the reader evaluates it.
+ private[connector] def record(filter: StorageFilter): Unit = {
+ val bytes = new ByteArrayOutputStream()
Review Comment:
Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). The test source uses `Utils.serialize`, `RangeInputPartition` and `Utils.stringToSeq`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -187,11 +196,19 @@ class DataSourceRDD(
castPartition(split).inputPartition.iterator.flatMap { inputPartition =>
val iter = if (columnarReads) {
- val batchReader = partitionReaderFactory.createColumnarReader(inputPartition)
- new PartitionIterator[ColumnarBatch](batchReader, _.numRows, taskState)
+ val batchReader = storageFilterRequest match {
Review Comment:
Thanks. I kept it as is. Each branch calls another pair of factory methods, so a helper would take two functions to save a few lines.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4377262897
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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.internal.connector
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.{SparkEnv, SparkException, TaskContext}
+import org.apache.spark.executor.Executor
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.internal.config.KILL_ON_FATAL_ERROR_DEPTH
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to
+ * it. The binding is built on first use, so an instance created on the driver can be evaluated on
+ * an executor.
+ */
+class StorageFilterImpl private (
+ private val catalystExpr: CatalystExpression,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes.
+ private val numReferences = attrs.length
+
+ // Whether this instance has logged a failure, so a task logs only its first.
+ @transient private var loggedFailure = false
+
+ @transient private lazy val boundPredicate: InternalRow => Boolean = {
+ val bound = BindReferences.bindReference(catalystExpr, attrs)
+ CatalystPredicate.create(bound).eval
+ }
+
+ /**
+ * Wraps a non-fatal error in [[StorageFilterNotAppliedException]], unless the task is being
+ * killed or a fatal error is in its cause chain. Each task logs the first one. A row of another
+ * width than [[references]] is a reader bug, which fails the task, since the predicate would read
+ * the wrong values.
+ */
+ @throws[StorageFilterNotAppliedException]
+ override def eval(row: InternalRow): Boolean = {
+ // Outside the `try`, so a failure to build the predicate, or a row of another width, fails the
+ // task.
+ if (row.numFields != numReferences) {
+ throw SparkException.internalError(
+ s"A storage filter over ${references.mkString("[", ", ", "]")} takes rows of " +
+ s"$numReferences values, but got ${row.numFields}")
+ }
+ val predicate = boundPredicate
+ try {
+ predicate(row)
+ } catch {
+ case NonFatal(e) if !StorageFilterImpl.taskKilled && !StorageFilterImpl.hasFatalCause(e) =>
Review Comment:
Thanks, fixed in [`4b01898`](https://github.com/apache/spark/pull/59715/commits/4b01898da695d7f6b5cce1ddca29d5c040dcf788). I had looked only at `DataSourceRDD`, which has no `ignoreCorruptFiles` catch, and missed `FilePartitionReader`'s. The V1 rule now lives in the new `StorageFilterErrors`, and `eval` follows it with `fromRead = false`. `StorageFilterImplSuite` now pins the V1 outcomes for a fatal cause, a bare `InternalError`, a kill and a `LinkageError`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60465][SQL] Late-materialization storage-filter pushdown for data source V2 scans [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4381901611
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/StorageFilter.java:
##########
@@ -0,0 +1,105 @@
+/*
+ * 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.expressions.filter;
+
+import java.io.Serializable;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.read.Scan;
+import org.apache.spark.sql.connector.read.StorageFilterRequest;
+import org.apache.spark.sql.connector.read.SupportsStorageFilters;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+/**
+ * A row-level predicate a scan may evaluate inside its readers, to skip reading the other columns
+ * of the rows it rejects. The reader reads the columns of {@link #references()} first, evaluates
+ * the filter, and reads the remaining columns only for the rows it keeps.
+ * <p>
+ * The filter is a predicate on the rows the scan returns. Spark offers none to a scan that took a
+ * pushed operation that decides which rows it returns, such as a limit, an offset, a top-N, a
+ * sample, an aggregate or a join.
+ * <p>
+ * Spark implements the predicate, and the reader only supplies the values. For example, a scan
+ * does not have to reproduce the Spark sketch and hash of a runtime bloom filter. The values must
+ * be the ones the scan returns for those columns, in the types of {@link Scan#readSchema()}, nulls
+ * included.
+ * <p>
+ * Applying it is optional, see {@link SupportsStorageFilters}.
+ * <p>
+ * A row's result depends only on its own values. So with one reference, a reader may evaluate each
+ * distinct value once, such as each entry of a dictionary, and reuse the result for every row
+ * holding that value. A dictionary entry is evaluated as the value the scan returns for it, decoded
+ * and converted.
+ * <p>
+ * Spark hands the filter to each task in a {@link StorageFilterRequest}. The readers of a task
+ * share one instance, which is not thread-safe. A reader must not copy or serialize it, since an
+ * evaluated instance may have released what a copy would be built from.
+ *
+ * @since 5.0.0
+ */
+@Evolving
+public abstract class StorageFilter implements Serializable {
Review Comment:
Thanks, moved in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538), with `StorageFilterNotAppliedException`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,116 @@
+/*
+ * 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.internal.connector
+
+import org.apache.spark.SparkException
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to
+ * it. The binding is built on first use, so an instance created on the driver can be evaluated on
+ * an executor.
+ */
+class StorageFilterImpl private (
+ private var catalystExpr: CatalystExpression,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes, read on every row.
+ private val numReferences = attrs.length
+
+ // Whether this instance has logged a failure, so a task logs only its first.
+ @transient private var loggedFailure = false
+
+ @transient private lazy val boundPredicate: InternalRow => Boolean = {
+ // The expression is released below, so a copy made after a first evaluation would have nothing
+ // to build from. That is named here rather than left to an NPE.
+ if (catalystExpr == null) {
+ throw SparkException.internalError(
+ s"$this was copied after an evaluation released the expression of its predicate")
+ }
+ val predicate = CatalystPredicate.create(BindReferences.bindReference(catalystExpr, attrs))
+ // A bloom's bytes sit in the expression as a binary literal, while a generated predicate holds
+ // the deserialized filter instead, so dropping the expression lets the task reclaim them. An
+ // interpreted predicate, the fallback, keeps the expression. Every task deserializes its own
+ // copy, and this is the field's only reader.
+ catalystExpr = null
+ predicate.eval
+ }
+
+ /**
+ * Wraps an error in [[StorageFilterNotAppliedException]], unless
+ * [[StorageFilterErrors.rethrowIfMustPropagate]] rethrows it, as the file-source reader does: a
+ * task kill, or a fatal error in its cause chain. Each task logs the first one. A row of another
+ * width than [[references]] is a reader bug, which fails the task, since the predicate would read
+ * the wrong values.
+ */
+ @throws[StorageFilterNotAppliedException]
+ override def eval(row: InternalRow): Boolean = {
+ // Outside the `try`, so a failure to build the predicate, or a row of another width, fails the
+ // task.
+ if (row.numFields != numReferences) {
+ throw SparkException.internalError(
+ s"$this takes rows of $numReferences values, but got ${row.numFields}")
+ }
+ val predicate = boundPredicate
+ try {
+ predicate(row)
+ } catch {
+ case e: Throwable =>
+ StorageFilterErrors.rethrowIfMustPropagate(e, fromRead = false)
+ if (!loggedFailure) {
+ loggedFailure = true
+ logWarning(log"${MDC(LogKeys.EXPR, this)} could not be applied, so the reader reads " +
+ log"without it.", e)
+ }
+ throw new StorageFilterNotAppliedException(e)
+ }
+ }
+
+ @transient override lazy val references: Array[NamedReference] =
Review Comment:
Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). `references()` returns a copy. A test reorders one and checks that the next call is unchanged.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala:
##########
@@ -0,0 +1,491 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.{SparkException, SparkThrowable}
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GetStructField, GreaterThan, If, Literal, NumericEvalContext, PredicateHelper, Rand, XxHash64}
+import org.apache.spark.sql.catalyst.plans.logical.ReturnAnswer
+import org.apache.spark.sql.execution.{BaseLimitExec, CoalesceExec, ExecSubqueryExpression, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation}
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.StorageFilterImpl
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{BinaryType, DoubleType, LongType, StructType}
+
+/**
+ * Covers the storage-filter contract for a data source v2 scan. With
+ * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through
+ * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The
+ * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked.
+ */
+class StorageFilterPushDownSuite
+ extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper with PredicateHelper {
+
+ private def sourceName: String = classOf[StorageFilterTestSource].getName
+
+ // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table.
+ private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = {
+ val bloomJoinConf = Seq(
+ SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true",
+ SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
+ withSQLConf(bloomJoinConf ++ conf: _*) {
+ withTable("bf_build") {
+ spark.range(0, 30).selectExpr("id AS bk", "id AS bv")
+ .write.format("parquet").saveAsTable("bf_build")
+ body
+ }
+ }
+ }
+
+ // A join whose build side is small and carries a selective predicate, which is the shape
+ // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`.
+ private def bloomJoin(probe: DataFrame): DataFrame = {
+ val build = spark.table("bf_build").where("bv = 5")
+ probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*)
+ }
+
+ private def v2Probe(options: (String, String)*): DataFrame =
+ spark.read.format(sourceName).options(options.toMap).load()
+
+ private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME
+
+ private def scanOf(plan: SparkPlan): BatchScanExec =
+ collect(plan) { case b: BatchScanExec => b }.headOption
+ .getOrElse(fail(s"no BatchScanExec in plan:\n$plan"))
+
+ // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows
+ // the reader marks as checked.
+ private def bloomsAbove(plan: SparkPlan): (Int, Int) = {
+ val conditions = collect(plan) { case f: FilterExec => f.condition }
+ val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum
+ val skipped = conditions.map(_.collect {
+ case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain)
+ if a.name == checkedName => 1
+ }.sum).sum
+ (all - skipped, skipped)
+ }
+
+ // The scan offered nothing, and the Filter evaluates the bloom on every row.
+ private def assertNothingOffered(plan: SparkPlan): Unit = {
+ val scan = scanOf(plan)
+ assert(scan.storageFilters.isEmpty && !scan.appendsCheckedColumn,
+ s"the scan must get no storage filter:\n$plan")
+ assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan")
+ }
+
+ // The rows the Filter on the scan returned, counted before the join.
+ private def rowsOutOfFilter(plan: SparkPlan): Long = {
+ val filters = collect(plan) {
+ case f: FilterExec if f.exists(_.isInstanceOf[BatchScanExec]) => f
+ }
+ assert(filters.size == 1, s"expected one Filter on the scan:\n$plan")
+ filters.head.metrics("numOutputRows").value
+ }
+
+ // Runs `df`, checks that it returns `expected`, and returns its plan.
+ private def checkJoin(df: DataFrame, expected: Set[Seq[Long]]): SparkPlan = {
+ val result = df.collect().map(_.toSeq.map(_.asInstanceOf[Long])).toSet
+ assert(result == expected, s"got $result")
+ df.queryExecution.executedPlan
+ }
+
+ // The build side keeps `bv = 5`, so the join returns that one row.
+ private val joinResult = Set(Seq(5L, 50L))
+
+ // Runs the bloom join over the v2 source read with `options`, checks its rows, and returns its
+ // plan.
+ private def runBloomJoin(options: (String, String)*): SparkPlan = {
+ StorageFilterTestSource.reset()
+ checkJoin(bloomJoin(v2Probe(options: _*)), joinResult)
+ }
+
+ for (columnar <- Seq(false, true); aqe <- Seq(false, true)) {
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind scan gets the checked column, and the Filter skips the bloom on checked rows " +
+ s"(AQE $aqe)") {
+ withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val plan = runBloomJoin("columnar" -> columnar.toString)
+ val scan = scanOf(plan)
+ assert(scan.supportsColumnar == columnar, s"the scan must read $kind:\n$plan")
+ assert(scan.storageFilters.size == 1,
+ s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan")
+ assert(scan.appendsCheckedColumn, s"the scan must get the checked column:\n$plan")
+ assert(bloomsAbove(plan) == (0, 1),
+ s"the Filter must skip the bloom on checked rows:\n$plan")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(),
+ s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " +
+ s"${StorageFilterTestSource.rowsReturned.get()}")
+ val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output }
+ assert(probeSide.exists(!_.exists(_.name == checkedName)),
+ s"the Project above the scan must drop the checked column; got $probeSide")
+ assert(StorageFilterTestSource.filteredReaders.get() == 1,
+ s"the one partition's reader must get the filter; " +
+ s"got ${StorageFilterTestSource.filteredReaders.get()}")
+ assert(StorageFilterTestSource.checkedColumnRequests.get() == 1,
+ "the reader must be asked for the checked column")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows,
+ s"the reader must evaluate every row; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows,
+ "the reader must have dropped rows, otherwise the filter did nothing")
+ }
+ }
+ }
+
+ Seq(false, true).foreach { columnar =>
+ val kind = if (columnar) "columnar" else "row-based"
+ test(s"a $kind reader that gives the filter up marks its rows unchecked, and the Filter " +
+ "filters them") {
+ withBloomJoin() {
+ val plan = runBloomJoin("columnar" -> columnar.toString, "giveUp" -> "true")
+ assert(bloomsAbove(plan) == (0, 1), s"the plan is that of a filtering reader:\n$plan")
+ assert(StorageFilterTestSource.rowsEvaluated.get() == 0,
+ s"the reader must evaluate nothing; got ${StorageFilterTestSource.rowsEvaluated.get()}")
+ assert(StorageFilterTestSource.rowsReturned.get() == StorageFilterTestSource.numRows,
+ s"the reader must return every row; got ${StorageFilterTestSource.rowsReturned.get()}")
+ assert(rowsOutOfFilter(plan) < StorageFilterTestSource.numRows,
+ s"the Filter must apply the bloom to unchecked rows; got ${rowsOutOfFilter(plan)}")
+ }
+ }
+
+ test(s"the Filter takes a $kind reader's word on the rows it marks checked") {
+ // This reader breaks the contract on purpose: it applies no filter but marks every row
+ // checked. Only skipping the bloom on those rows keeps all of them.
+ withBloomJoin() {
+ // The join must still filter, which `runBloomJoin` checks.
+ val plan = runBloomJoin("columnar" -> columnar.toString, "markAllChecked" -> "true")
+ assert(rowsOutOfFilter(plan) == StorageFilterTestSource.numRows,
+ s"the Filter must keep every row marked checked; got ${rowsOutOfFilter(plan)}")
+ }
+ }
+
+ test(s"a $kind reader that leaves out the checked column it was asked for fails the task, " +
+ "from any row or batch") {
+ // A reader that gives the filter up returns every row, so it returns several batches.
+ Seq(Seq("dropChecked" -> "true"), Seq("dropChecked" -> "later", "giveUp" -> "true"))
+ .foreach { options =>
+ withBloomJoin() {
+ StorageFilterTestSource.reset()
+ val df = bloomJoin(v2Probe(options :+ ("columnar" -> columnar.toString): _*))
+ val e = intercept[SparkException](df.collect())
+ val missing =
+ Iterator.iterate[Throwable](e)(_.getCause).takeWhile(_ != null).collectFirst {
+ case t: SparkThrowable
+ if t.getCondition == "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING" => t
+ }
+ assert(missing.exists(_.getMessageParameters.get("class") ==
Review Comment:
Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). The test uses `checkError` with the condition, `KD000` and all three parameters. So does the new test for item 1 whose reader returns the column ahead of the others.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -957,7 +972,8 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
// changes with variant extraction. When validConstraints is wired up, this needs revisiting.
// Pushed variant extraction cannot be reproduced by rebuilding the scan, so this scan stays
// not-mergeable (the default `mergeableScan = false`).
- val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput)
+ val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput,
+ onlyRowWisePushdowns = !hasRowSelectingPushdown(holder))
Review Comment:
Thanks, added in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). The test source now implements `SupportsPushDownVariantExtractions`, and a new test runs a query of your example's shape over it. Dropping the `onlyRowWisePushdowns` argument at the variant build site fails it. The name match matters only for a scan that reads just its key columns and metadata columns. A variant scan gets there when the declared column is the variant column and the join key is an extraction of it. A second test reads only such an extraction and a metadata column, and checks that nothing is offered. Going back to the exprId match fails it.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala:
##########
@@ -0,0 +1,116 @@
+/*
+ * 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.internal.connector
+
+import org.apache.spark.SparkException
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException}
+
+/**
+ * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not
+ * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key
+ * rather than a data row.
+ *
+ * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to
+ * it. The binding is built on first use, so an instance created on the driver can be evaluated on
+ * an executor.
+ */
+class StorageFilterImpl private (
+ private var catalystExpr: CatalystExpression,
+ private val attrs: Seq[Attribute])
+ extends StorageFilter with Logging {
+
+ // The width of the row `eval` takes, read on every row.
+ private val numReferences = attrs.length
+
+ // Whether this instance has logged a failure, so a task logs only its first.
+ @transient private var loggedFailure = false
+
+ @transient private lazy val boundPredicate: InternalRow => Boolean = {
Review Comment:
Thanks, changed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). The new `StorageFilterPredicate`, next to `StorageFilterErrors`, holds the bound expression. It builds the predicate on first use and then releases the expression. Its `apply` binds the expression, so `StorageFilterImpl` binds on the driver now, and `ParquetStorageFilter` passes its expression already bound. The bind and the build happen at different times now, so the two-argument `create` no longer fits.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org