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