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: [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]
