peter-toth commented on code in PR #59715: URL: https://github.com/apache/spark/pull/59715#discussion_r4386748262
########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterPredicate.scala: ########## @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.internal.connector + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.expressions.{Attribute, BasePredicate, BindReferences, Expression, Predicate} + +/** + * The predicate of a storage filter, built on first use from `boundExpression`, since a generated + * one can be awkward to serialize. The file-source Parquet reader and [[StorageFilterImpl]] share + * it. + */ +class StorageFilterPredicate(private var boundExpression: Expression) extends Serializable { + + @transient lazy val prepared: BasePredicate = { + // The expression is released below, so a copy made after a first evaluation would have nothing + // to build from. That is named here rather than left to an NPE. + if (boundExpression == null) { + throw SparkException.internalError("A storage filter was copied after an evaluation " + Review Comment: Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). `StorageFilterPredicate` takes the filter's name for this message. That is the `toString` of `StorageFilterImpl`, and the key columns for the V1 filter. The message also says that only a filter that was never evaluated can be copied. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala: ########## @@ -0,0 +1,103 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.internal.connector + +import org.apache.spark.SparkException +import org.apache.spark.internal.{Logging, LogKeys} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression => CatalystExpression} +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference} +import org.apache.spark.sql.connector.read.{StorageFilter, StorageFilterNotAppliedException} + +/** + * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not + * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key + * rather than a data row. + * + * `attrs` is the layout of the row [[eval]] expects, in that order, and the predicate is bound to + * it. The predicate is built on first use, see [[StorageFilterPredicate]]. + */ +class StorageFilterImpl private ( + private val predicate: StorageFilterPredicate, + private val attrs: Seq[Attribute]) + extends StorageFilter with Logging { + + // The width of the row `eval` takes, read on every row. + private val numReferences = attrs.length + + // Whether this instance has logged a failure, so a task logs only its first. + @transient private var loggedFailure = false + + /** + * Wraps an error in [[StorageFilterNotAppliedException]], unless + * [[StorageFilterErrors.rethrowIfMustPropagate]] rethrows it, as the file-source reader does: a + * task kill, or a fatal error in its cause chain. Each task logs the first one. A row of another + * width than [[references]] is a reader bug, which fails the task, since the predicate would read + * the wrong values. + */ + @throws[StorageFilterNotAppliedException] + override def eval(row: InternalRow): Boolean = { + // Outside the `try`, so a failure to build the predicate, or a row of another width, fails the + // task. + if (row.numFields != numReferences) { + throw SparkException.internalError( + s"$this takes rows of $numReferences values, but got ${row.numFields}") + } + val prepared = predicate.prepared Review Comment: Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). `eval` keeps the built predicate in a plain transient field. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala: ########## @@ -0,0 +1,103 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.internal.connector + +import org.apache.spark.SparkException +import org.apache.spark.internal.{Logging, LogKeys} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression => CatalystExpression} +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference} +import org.apache.spark.sql.connector.read.{StorageFilter, StorageFilterNotAppliedException} + +/** + * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not + * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key + * rather than a data row. + * + * `attrs` is the layout of the row [[eval]] expects, in that order, and the predicate is bound to + * it. The predicate is built on first use, see [[StorageFilterPredicate]]. + */ +class StorageFilterImpl private ( + private val predicate: StorageFilterPredicate, + private val attrs: Seq[Attribute]) Review Comment: Thanks, changed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). `StorageFilterImpl` keeps the references, and `StorageFilterPredicate.apply` is gone. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala: ########## @@ -0,0 +1,560 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import scala.collection.mutable + +import org.apache.spark.SparkException +import org.apache.spark.sql.{DataFrame, QueryTest} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GetStructField, GreaterThan, If, LessThanOrEqual, Literal, NumericEvalContext, PredicateHelper, Rand, XxHash64} +import org.apache.spark.sql.catalyst.plans.logical.ReturnAnswer +import org.apache.spark.sql.execution.{BaseLimitExec, CoalesceExec, ExecSubqueryExpression, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning} +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation} +import org.apache.spark.sql.execution.joins.SortMergeJoinExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.connector.StorageFilterImpl +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{BinaryType, DoubleType, LongType, StructType} + +/** + * Covers the storage-filter contract for a data source v2 scan. With + * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through + * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The + * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked. + */ +class StorageFilterPushDownSuite + extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper with PredicateHelper { + + private def sourceName: String = classOf[StorageFilterTestSource].getName + + // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table. + private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = { + val bloomJoinConf = Seq( + SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") + withSQLConf(bloomJoinConf ++ conf: _*) { + withTable("bf_build") { + spark.range(0, 30).selectExpr("id AS bk", "id AS bv") + .write.format("parquet").saveAsTable("bf_build") + body + } + } + } + + // A join whose build side is small and carries a selective predicate, which is the shape + // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`. + private def bloomJoin(probe: DataFrame): DataFrame = { + val build = spark.table("bf_build").where("bv = 5") + probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*) + } + + private def v2Probe(options: (String, String)*): DataFrame = + spark.read.format(sourceName).options(options.toMap).load() + + private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME + + private def scanOf(plan: SparkPlan): BatchScanExec = + collect(plan) { case b: BatchScanExec => b }.headOption + .getOrElse(fail(s"no BatchScanExec in plan:\n$plan")) + + // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows + // the reader marks as checked. + private def bloomsAbove(plan: SparkPlan): (Int, Int) = { + val conditions = collect(plan) { case f: FilterExec => f.condition } + val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum + val skipped = conditions.map(_.collect { + case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain) + if a.name == checkedName => 1 + }.sum).sum + (all - skipped, skipped) + } + + // The scan offered nothing, and the Filter evaluates the bloom on every row. + private def assertNothingOffered(plan: SparkPlan): Unit = { + val scan = scanOf(plan) + assert(scan.storageFilters.isEmpty, s"the scan must get no storage filter:\n$plan") + assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan") + } + + // The rows the Filter on the scan returned, counted before the join. + private def rowsOutOfFilter(plan: SparkPlan): Long = { + val filters = collect(plan) { + case f: FilterExec if f.exists(_.isInstanceOf[BatchScanExec]) => f + } + assert(filters.size == 1, s"expected one Filter on the scan:\n$plan") + filters.head.metrics("numOutputRows").value + } + + // Runs `df`, checks that it returns `expected`, and returns its plan. + private def checkJoin(df: DataFrame, expected: Set[Seq[Long]]): SparkPlan = { + val result = df.collect().map(_.toSeq.map(_.asInstanceOf[Long])).toSet + assert(result == expected, s"got $result") + df.queryExecution.executedPlan + } + + // The build side keeps `bv = 5`, so the join returns that one row. + private val joinResult = Set(Seq(5L, 50L)) + + // Runs the bloom join over the v2 source read with `options`, checks its rows, and returns its + // plan. + private def runBloomJoin(options: (String, String)*): SparkPlan = { + StorageFilterTestSource.reset() + checkJoin(bloomJoin(v2Probe(options: _*)), joinResult) + } + + for (columnar <- Seq(false, true); aqe <- Seq(false, true)) { + val kind = if (columnar) "columnar" else "row-based" + test(s"a $kind scan gets the checked column, and the Filter skips the bloom on checked rows " + + s"(AQE $aqe)") { + withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + val plan = runBloomJoin("columnar" -> columnar.toString) + val scan = scanOf(plan) + assert(scan.supportsColumnar == columnar, s"the scan must read $kind:\n$plan") + assert(scan.storageFilters.size == 1, + s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan") + assert(scan.appendsCheckedColumn, s"the scan must get the checked column:\n$plan") + assert(bloomsAbove(plan) == (0, 1), + s"the Filter must skip the bloom on checked rows:\n$plan") + assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(), + s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " + + s"${StorageFilterTestSource.rowsReturned.get()}") + val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output } + assert(probeSide.exists(!_.exists(_.name == checkedName)), + s"the Project above the scan must drop the checked column; got $probeSide") + assert(StorageFilterTestSource.filteredReaders.get() == 1, + s"the one partition's reader must get the filter; " + + s"got ${StorageFilterTestSource.filteredReaders.get()}") + assert(StorageFilterTestSource.checkedColumnRequests.get() == 1, + "the reader must be asked for the checked column") + assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows, + s"the reader must evaluate every row; got ${StorageFilterTestSource.rowsEvaluated.get()}") + assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows, + "the reader must have dropped rows, otherwise the filter did nothing") + } + } + } + + Seq(false, true).foreach { columnar => + val kind = if (columnar) "columnar" else "row-based" + test(s"a $kind reader that gives the filter up marks its rows unchecked, and the Filter " + + "filters them") { + withBloomJoin() { + val plan = runBloomJoin("columnar" -> columnar.toString, "giveUp" -> "true") + assert(bloomsAbove(plan) == (0, 1), s"the plan is that of a filtering reader:\n$plan") + assert(StorageFilterTestSource.rowsEvaluated.get() == 0, + s"the reader must evaluate nothing; got ${StorageFilterTestSource.rowsEvaluated.get()}") + assert(StorageFilterTestSource.rowsReturned.get() == StorageFilterTestSource.numRows, + s"the reader must return every row; got ${StorageFilterTestSource.rowsReturned.get()}") + assert(rowsOutOfFilter(plan) < StorageFilterTestSource.numRows, + s"the Filter must apply the bloom to unchecked rows; got ${rowsOutOfFilter(plan)}") + } + } + + test(s"the Filter takes a $kind reader's word on the rows it marks checked") { + // This reader breaks the contract on purpose: it applies no filter but marks every row + // checked. Only skipping the bloom on those rows keeps all of them. + withBloomJoin() { + // The join must still filter, which `runBloomJoin` checks. + val plan = runBloomJoin("columnar" -> columnar.toString, "markAllChecked" -> "true") + assert(rowsOutOfFilter(plan) == StorageFilterTestSource.numRows, + s"the Filter must keep every row marked checked; got ${rowsOutOfFilter(plan)}") + } + } + + Seq(false, true).foreach { aqe => + test(s"a $kind reader that leaves out the checked column it was asked for fails the task, " + + s"from any row or batch (AQE $aqe)") { + // A reader that gives the filter up returns every row, so it returns several batches. + Seq(Seq("dropChecked" -> "true"), Seq("dropChecked" -> "later", "giveUp" -> "true")) + .foreach { options => + withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + StorageFilterTestSource.reset() + val df = bloomJoin(v2Probe(options :+ ("columnar" -> columnar.toString): _*)) + // `k`, `v` and the checked column. + checkError( + exception = intercept[SparkException](df.collect()), + condition = "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING", + sqlState = Some("KD000"), + parameters = Map( + "class" -> classOf[StorageFilterTestReaderFactory].getName, + "expected" -> "3", + "actual" -> "2")) + } + } + } + } + + test(s"a $kind reader may return more columns after the checked column") { + withBloomJoin() { + val plan = runBloomJoin("columnar" -> columnar.toString, "checkedLayout" -> "extra") + assert(bloomsAbove(plan) == (0, 1), + s"the Filter must skip the bloom on checked rows:\n$plan") + assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(), + s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " + + s"${StorageFilterTestSource.rowsReturned.get()}") + } + } + + test(s"a $kind test reader gives a filter up when it throws, and returns its rows unchecked") { + // `1 / k` under ANSI throws at k = 0, the partition's first row, so the reader gives the + // filter up before it keeps or drops anything. No runtime bloom throws, so the reader is + // driven by hand. + val k = AttributeReference("k", LongType)() + val inverse = Divide(Literal(1.0), Cast(k, DoubleType), NumericEvalContext(EvalMode.ANSI)) + val filter = StorageFilterImpl(GreaterThan(inverse, Literal(0.0)), Seq(k)) + val partition = RangeInputPartition(0, StorageFilterTestSource.numRows) + val schema = new StructType().add("k", LongType).add("v", LongType) + val rows = mutable.ArrayBuffer.empty[(Long, Boolean)] + if (columnar) { + val reader = new StorageFilterTestColumnarReader( + partition, schema, Some(filter), appendsChecked = true) + try { + while (reader.next()) { + val batch = reader.get() + (0 until batch.numRows()).foreach { i => + rows += ((batch.column(0).getLong(i), batch.column(2).getBoolean(i))) + } + } + } finally { + reader.close() + } + } else { + val reader = + new StorageFilterTestReader(partition, schema, Some(filter), appendsChecked = true) + while (reader.next()) { + rows += ((reader.get().getLong(0), reader.get().getBoolean(2))) + } + } + assert(rows == (0L until StorageFilterTestSource.numRows.toLong).map(_ -> false), + s"expected every row, unchecked; got ${rows.take(5)} and ${rows.size} rows in all") + } + } + + test("a columnar reader that returns the checked column ahead of the others fails the task") { + // The batch is as wide as asked for, but its last column is not a boolean. A row carries no + // types, so a row-based reader's mistake goes unnoticed. + withBloomJoin() { + StorageFilterTestSource.reset() + val df = bloomJoin(v2Probe("columnar" -> "true", "checkedLayout" -> "front")) + checkError( + exception = intercept[SparkException](df.collect()), + condition = "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING", + sqlState = Some("KD000"), + parameters = Map( + "class" -> classOf[StorageFilterTestReaderFactory].getName, + "expected" -> "3", + "actual" -> "3")) + } + } + + // Under a limit of two fields, the checked column would make three. + Seq( + ("too wide for whole-stage codegen", SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> "2"), + ("without whole-stage codegen", SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") + ).foreach { case (when, conf) => + test(s"a scan $when gets the bloom without the checked column") { + withBloomJoin(conf) { + val plan = runBloomJoin("columnar" -> "true") + val scan = scanOf(plan) + assert(scan.storageFilters.size == 1, + s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan") + assert(!scan.appendsCheckedColumn, s"the scan must not get the checked column:\n$plan") + assert(StorageFilterTestSource.checkedColumnRequests.get() == 0, + "the reader must not be asked for the checked column") + assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan") + assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows, + "the reader must have dropped rows, otherwise the filter did nothing") + } + } + } + + Seq( + ("with the conf off", Seq(SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "false"), Nil), + ("to a scan that declares no column", Nil, Seq("declared" -> "")), + ("to a scan that declares only another column", Nil, Seq("declared" -> "v")), + ("to a scan that declares only a nested field", Nil, Seq("declared" -> "k.x")) + ).foreach { case (when, conf, options) => + test(s"nothing is offered $when") { + withBloomJoin(conf: _*) { + val plan = runBloomJoin(options :+ ("columnar" -> "true"): _*) + assertNothingOffered(plan) + assert(StorageFilterTestSource.filteredReaders.get() == 0, + "no reader may get a storage filter") + } + } + } + + test("a scan that reads only the key and a metadata column is offered nothing") { + // Only the data columns count as columns left to skip, as on the file-source side. + withBloomJoin() { + val source = v2Probe("columnar" -> "true", "metadataColumn" -> "true") + val pos = StorageFilterTestSource.posName + StorageFilterTestSource.reset() + assertNothingOffered( + checkJoin(bloomJoin(source.select(source("k"), source.metadataColumn(pos))), + Set(Seq(5L, 5L)))) + StorageFilterTestSource.reset() + val withValue = checkJoin( + bloomJoin(source.select(source("k"), source("v"), source.metadataColumn(pos))), + Set(Seq(5L, 50L, 5L))) + assert(scanOf(withValue).storageFilters.size == 1, + s"with `v` left to skip, the bloom must be offered:\n$withValue") + } + } + + test("a variant scan that reads only its key column and a metadata column is offered nothing") { + // The key is an extraction of the declared variant column, so that column is the scan's only + // data column, and the scan's output has fresh exprIds. + withBloomJoin() { + val pos = StorageFilterTestSource.posName + val probe = v2Probe("columnar" -> "true", "variant" -> "true", "metadataColumn" -> "true", + "declared" -> "j").selectExpr("variant_get(j, '$.a', 'long') AS a", pos) + val build = spark.table("bf_build").where("bv = 5") + val df = probe.join(build, probe("a") === build("bk") * 10).select(probe("a"), probe(pos)) + StorageFilterTestSource.reset() + assertNothingOffered(checkJoin(df, Set(Seq(50L, 5L)))) + } + } + + test("a scan that took a pushed variant extraction gets the bloom") { + // The extraction keeps the scan's rows. + withBloomJoin() { + StorageFilterTestSource.reset() + val probe = v2Probe("columnar" -> "true", "variant" -> "true") + .selectExpr("k", "variant_get(j, '$.a', 'long') AS a") + val plan = checkJoin(bloomJoin(probe), joinResult) + val scan = scanOf(plan) + assert(scan.output.exists(_.dataType.isInstanceOf[StructType]), + s"test setup: the scan must read the extraction:\n$plan") + assert(scan.storageFilters.size == 1, + s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan") + assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows, + "the reader must have dropped rows, otherwise the filter did nothing") + } + } + + for (aqe <- Seq(false, true)) { + test("a scan with a column named like the checked column gets the filter, and runtime " + + s"filtering still resolves that column (AQE $aqe)") { + // Runtime filtering resolves the declared filter attributes by name, which the checked column + // the scan appends would make ambiguous. Dynamic partition pruning on `k` would take the + // bloom's place. + withBloomJoin(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + val plan = runBloomJoin( + "columnar" -> "true", "valueName" -> checkedName, "runtimeFiltering" -> "v2") + val scan = scanOf(plan) + assert(scan.storageFilters.size == 1 && scan.appendsCheckedColumn, + s"the scan must get the bloom and the checked column:\n$plan") + assert(StorageFilterTestSource.runtimeFilterCalls.get() > 0, + "test setup: the bloom must reach runtime filtering") + val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output } + assert(probeSide.contains(scan.output.init), + s"the Project above the scan must keep the table's column; got $probeSide") + } + } + + test(s"a bloom the scan fully evaluates through runtime filtering is not offered (AQE $aqe)") { + // Dynamic partition pruning on `k` would take the bloom's place. + withBloomJoin(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + val plan = runBloomJoin("columnar" -> "true", "runtimeFiltering" -> "catalyst") + val scan = scanOf(plan) + assert(scan.runtimeFilters.exists(_.exists(_.isInstanceOf[BloomFilterMightContain])) && + StorageFilterTestSource.runtimeFilterCalls.get() > 0, + s"test setup: the scan must take the bloom as a runtime filter:\n$plan") + assert(scan.storageFilters.isEmpty, s"the scan must get no storage filter:\n$plan") + assert(bloomsAbove(plan) == (0, 0), s"no Filter may keep the bloom:\n$plan") + } + } + + test(s"a two-key join offers both blooms in one filter, each bloom with its own bytes " + + s"(AQE $aqe)") { + // `MergeSubplans` merges the two blooms' aggregates into one subquery returning both blooms. + withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + withTable("bf_build2") { + spark.range(0, 30).selectExpr("id AS bk", "id * 10 AS bv", "id AS flag") + .write.format("parquet").saveAsTable("bf_build2") + StorageFilterTestSource.reset() + val probe = v2Probe("columnar" -> "true", "wide" -> "true", "declared" -> "k,v") + .where("w % 3 != 1") + val build = spark.table("bf_build2").where("flag = 5") + val df = probe.join(build, probe("k") === build("bk") && probe("v") === build("bv")) + .select(probe("k"), probe("v"), probe("w")) + val plan = checkJoin(df, Set(Seq(5L, 50L, 500L))) + val scan = scanOf(plan) + assert(scan.storageFilters.size == 2, s"both blooms must be on the scan:\n$plan") + assert(scan.storageFilters.forall(_.exists { + case GetStructField(_: ExecSubqueryExpression, _, _) => true + case _ => false + }), s"test setup: the blooms must come from one merged subquery:\n$plan") + assert(bloomsAbove(plan) == (0, 2), s"the Filter must skip both blooms:\n$plan") + val filter = collect(plan) { + case f: FilterExec if f.child.exists(_.isInstanceOf[BatchScanExec]) => f + }.head + assert(splitConjunctivePredicates(filter.condition).exists { c => + c.references.map(_.name).toSet == Set("w") && !c.isInstanceOf[If] + }, s"the Filter must evaluate the conjunct over `w` on every row:\n$plan") + assert(StorageFilterTestSource.filterReferences == Set(Seq("k", "v")), + s"one filter over both keys; got ${StorageFilterTestSource.filterReferences}") + val bloomBytes = scan.storageFilters.map(_.collectFirst { + case m: BloomFilterMightContain => + m.bloomFilterExpression.eval().asInstanceOf[Array[Byte]].length + }.get).sum + assert(StorageFilterTestSource.filterBytes.get() < bloomBytes * 3 / 2, + s"each bloom must ship once; the filter takes ${StorageFilterTestSource.filterBytes} " + + s"bytes for $bloomBytes bytes of blooms") + assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows, + "the reader must evaluate every row; got " + + StorageFilterTestSource.rowsEvaluated.get()) + } + } + } + + test(s"the readers of one task share the filter the task deserialized (AQE $aqe)") { + // `coalesce(1)` reads both partitions in one task. + withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + StorageFilterTestSource.reset() + val probe = v2Probe("columnar" -> "true", "numPartitions" -> "2").coalesce(1) + val plan = checkJoin(bloomJoin(probe), joinResult) + assert(collect(plan) { case c: CoalesceExec => c }.nonEmpty, s"test setup:\n$plan") + assert(scanOf(plan).storageFilters.size == 1, s"the bloom must be on the scan:\n$plan") + assert(StorageFilterTestSource.filteredReaders.get() == 2, + s"both readers must get the filter; got ${StorageFilterTestSource.filteredReaders.get()}") + assert(StorageFilterTestSource.filterIdentities.size == 1, + "both readers must get the same instance") + assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows, + "the readers must evaluate every row; got " + StorageFilterTestSource.rowsEvaluated.get()) + assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows, + "the readers must have dropped rows, otherwise the filter did nothing") + } + } + } + + test("the storage filter shows in the scan description and separates canonicalized plans") { + withBloomJoin() { + val plan = bloomJoin(v2Probe("columnar" -> "true")).queryExecution.executedPlan + val scan = scanOf(plan) + assert(scan.simpleString(100).contains("StorageFilters:"), + s"the scan description must name the filters: ${scan.simpleString(100)}") + assert(scan.verboseStringWithOperatorId().contains("StorageFilters:"), + s"EXPLAIN FORMATTED must name the filters: ${scan.verboseStringWithOperatorId()}") + + intercept[IllegalArgumentException](scan.copy(storageFilters = Nil)) + // As the scan would be without the checked column, too wide or without codegen. + val unchecked = scan.copy(output = scan.output.init, appendsCheckedColumn = false) + assert(scan != unchecked, "equals must see the checked column") + val plain = unchecked.copy(storageFilters = Nil) + assert(unchecked != plain && unchecked.hashCode != plain.hashCode, + "equals and hashCode must see storageFilters") + assert(!unchecked.sameResult(plain), + s"a filtering scan must not be reusable as a plain one:\n${unchecked.canonicalized}\n" + + s"${plain.canonicalized}") + assert(!plain.simpleString(100).contains("StorageFilters:"), + s"a plain scan's description must not name storage filters: ${plain.simpleString(100)}") + assert(!plain.verboseStringWithOperatorId().contains("StorageFilters:"), + s"nor its EXPLAIN FORMATTED: ${plain.verboseStringWithOperatorId()}") + + // As another relation instance of the same table gets, with fresh exprIds. + val fresh = AttributeMap(scan.output.map(a => a -> a.newInstance())) + val twin = scan.copy( + output = scan.output.map(fresh(_)), + storageFilters = scan.storageFilters.map(_.transform { + case a: Attribute => fresh.getOrElse(a, a) + })) + assert(scan.sameResult(twin), + s"the same filtering scan must be reusable:\n${scan.canonicalized}\n${twin.canonicalized}") + } + } + + // The reader would filter ahead of the operation, which then picks other rows. + Seq[(String, DataFrame => DataFrame)]( + "limit" -> (_.limit(100)), + "offset" -> (_.offset(10)), + "topN" -> (_.orderBy("v").limit(100)), + "sample" -> (_.sample(0.5)) + ).foreach { case (op, apply) => + test(s"nothing is offered to a scan that took a pushed $op") { + withBloomJoin() { + val probe = apply(v2Probe("columnar" -> "true", "push" -> op)) + val plan = bloomJoin(probe).queryExecution.executedPlan + assert(collect(plan) { + case p @ (_: BaseLimitExec | _: TakeOrderedAndProjectExec | _: SampleExec) => p + }.isEmpty, s"the $op must be pushed into the scan whole:\n$plan") + assertNothingOffered(plan) + } + } + } + + test("the scan's pushdowns decide on storage filters, not whether it can be merged") { + // A pushed variant extraction, for one, keeps the scan's rows but cannot be merged. + withBloomJoin() { + val optimized = bloomJoin(v2Probe("columnar" -> "true")).queryExecution.optimizedPlan + def planned(mergeable: Boolean, rowWise: Boolean): BatchScanExec = scanOf( + spark.sessionState.planner.plan(ReturnAnswer(optimized.transform { Review Comment: Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). ########## sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterTestSource.scala: ########## @@ -0,0 +1,492 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import java.util +import java.util.OptionalLong +import java.util.concurrent.atomic.{AtomicInteger, AtomicLong} + +import scala.collection.mutable +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow} +import org.apache.spark.sql.connector.catalog.{MetadataColumn, SupportsMetadataColumns, SupportsRead, Table, TableCapability, TableProvider} +import org.apache.spark.sql.connector.expressions.{Expressions, FieldReference, NamedReference, SortOrder, Transform} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read._ +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector +import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering +import org.apache.spark.sql.types.{BooleanType, DataType, LongType, StructField, StructType, VariantType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +/** + * A read-only v2 source over a fixed table of longs that implements [[SupportsStorageFilters]]: + * the key `k`, the value `v` holding `k * 10`, and with `wide` also `w` holding `k * 100`. Its + * reader gets the filter in the task and drops the rows the filter rejects. It reads the columns + * Spark requires, and builds the filter's row from the filter's `references()`. Its options: + * - `columnar`: reads in batches, through `StorageFilter.select`, rather than row by row; + * - `giveUp`: reads without the filter, as a reader does for a unit it cannot use the filter on. + * Here the unit is the partition; + * - `markAllChecked`: reads without the filter but marks every row checked, which breaks the + * contract on purpose; + * - `dropChecked`: `true` never appends the checked column, and `later` appends it to the first + * row or batch only. Either breaks the contract on purpose; + * - `checkedLayout`: `front` returns the checked column ahead of the scan's columns, which breaks + * the contract on purpose. `extra` returns one more column after it; + * - `declared`: the comma-separated storage-filter columns, `k` by default, a dot separating a + * nested field; + * - `valueName`: the name of the value column, `v` by default; + * - `wide`: adds the column `w`; + * - `variant`: adds the variant column `j`, which the scan reads only through pushed extractions + * of a long. Each one holds `k * 10`; + * - `numPartitions`: splits the keys into that many input partitions, 1 by default; + * - `metadataColumn`: adds the metadata column `_pos`, which holds the key; + * - `runtimeFiltering`: `v2` makes the scan a `SupportsRuntimeV2Filtering` one over every column, + * with iterative pushdown on a table partitioned by `k`. `catalyst` makes it a + * `SupportsRuntimeCatalystFiltering` one that evaluates the filters on `k` fully. Neither + * applies the filters it takes, so they serve planning tests only; + * - `push`: the comma-separated pushed operations it takes whole without applying them, of + * `limit`, `offset`, `topN` and `sample`, for planning tests only. + * + * A filter that throws `StorageFilterNotAppliedException` is given up for the rest of the + * partition too. The counters below record what the readers got, evaluated and returned. + */ +class StorageFilterTestSource extends TableProvider { + + override def inferSchema(options: CaseInsensitiveStringMap): StructType = { + val schema = new StructType().add("k", LongType).add(options.getOrDefault("valueName", "v"), + LongType) + val wide = if (options.getBoolean("wide", false)) schema.add("w", LongType) else schema + if (options.getBoolean("variant", false)) wide.add("j", VariantType) else wide + } + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: util.Map[String, String]): Table = { + val options = new CaseInsensitiveStringMap(properties) + if (options.getBoolean("metadataColumn", false)) { + new StorageFilterTestTable(schema, options) with SupportsMetadataColumns { + override def metadataColumns(): Array[MetadataColumn] = Array(new MetadataColumn { + override def name(): String = StorageFilterTestSource.posName + override def dataType(): DataType = LongType + }) + } + } else { + new StorageFilterTestTable(schema, options) + } + } +} + +object StorageFilterTestSource { + + /** Rows the source serves, `k` from 0 until numRows. */ + val numRows: Int = 200 + + /** Rows a batch holds at most in a columnar read. */ + val batchSize: Int = 64 + + /** The name of the metadata column the `metadataColumn` option adds. */ + val posName: String = "_pos" + + /** How many readers were created with a storage-filter request. */ + val filteredReaders = new AtomicInteger(0) + + /** How many of those requests asked for the checked column. */ + val checkedColumnRequests = new AtomicInteger(0) + + /** Rows the reader evaluated the storage filter against. */ + val rowsEvaluated = new AtomicLong(0) + + /** Rows the reader returned. */ + val rowsReturned = new AtomicLong(0) + + /** The serialized size of the first filter a reader got, before it evaluated the filter. */ + val filterBytes = new AtomicLong(0) + + /** The references of the filters the readers got, and the identities of those filters. */ + val filterReferences = mutable.Set.empty[Seq[String]] + val filterIdentities = mutable.Set.empty[Int] + + /** How many times a runtime-filtering scan was handed filters. */ + val runtimeFilterCalls = new AtomicInteger(0) + + def reset(): Unit = { + filteredReaders.set(0) + checkedColumnRequests.set(0) + rowsEvaluated.set(0) + rowsReturned.set(0) + filterBytes.set(0) + filterReferences.synchronized(filterReferences.clear()) + filterIdentities.synchronized(filterIdentities.clear()) + runtimeFilterCalls.set(0) + } + + /** The value of column `name` in the row of `key`. Any column but `k`, `w` and `_pos` is `v`. */ + def valueOf(name: String, key: Long): Long = name match { + case "k" | `posName` => key + case "w" => key * 100 + case _ => key * 10 + } + + /** The value of `field` in the row of `key`, a struct for the extractions of a variant. */ + def valueOf(field: StructField, key: Long): Any = field.dataType match { + case struct: StructType => InternalRow.fromSeq(struct.map(_ => valueOf(field.name, key))) + case _ => valueOf(field.name, key) + } + + /** The columns of the row the filter evaluates, in that order. */ + def columnsOf(filter: StorageFilter): Seq[String] = filter.references.map(_.fieldNames.head).toSeq + + // Records the filter a reader got, before the reader evaluates it. Only the first reader + // serializes it, since a later reader of the same task may get it evaluated, see `StorageFilter`. + private[connector] def record(filter: StorageFilter): Unit = { + if (filterBytes.get() == 0) filterBytes.compareAndSet(0, Utils.serialize(filter).length) + filterReferences.synchronized(filterReferences += columnsOf(filter)) + filterIdentities.synchronized(filterIdentities += System.identityHashCode(filter)) + } +} + +class StorageFilterTestTable(tableSchema: StructType, tableOptions: CaseInsensitiveStringMap) + extends Table with SupportsRead { + + override def name(): String = "storage_filter_test" + + override def schema(): StructType = tableSchema + + override def partitioning(): Array[Transform] = + if (tableOptions.getOrDefault("runtimeFiltering", "") == "v2") { + Array(Expressions.identity("k")) + } else { + Array.empty + } + + override def capabilities(): util.Set[TableCapability] = + Set(TableCapability.BATCH_READ).asJava + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = + new StorageFilterTestScanBuilder(tableSchema, options) +} + +/** + * Prunes the columns Spark does not read, and takes the pushed operations named in `push` whole. + * Its scan does not apply them, so they serve planning tests only. It takes every variant + * extraction, and reads each into a struct field of the variant's column named by its position. + */ +class StorageFilterTestScanBuilder(tableSchema: StructType, options: CaseInsensitiveStringMap) + extends SupportsPushDownRequiredColumns + with SupportsPushDownLimit + with SupportsPushDownOffset + with SupportsPushDownTopN + with SupportsPushDownTableSample + with SupportsPushDownVariantExtractions { + + private val pushes = Utils.stringToSeq(options.getOrDefault("push", "")).toSet + private var required = tableSchema + private var extractions = Array.empty[VariantExtraction] + + override def pruneColumns(requiredSchema: StructType): Unit = required = requiredSchema + + override def pushVariantExtractions(pushed: Array[VariantExtraction]): Array[Boolean] = { + extractions = pushed + pushed.map(_ => true) + } + + override def build(): Scan = { + val byColumn = extractions.groupBy(_.columnName().head) + val schema = StructType(required.map { field => + byColumn.get(field.name).map { columnExtractions => + field.copy(dataType = StructType(columnExtractions.zipWithIndex.map { case (e, i) => + StructField(i.toString, e.expectedDataType(), nullable = true, e.metadata()) + })) + }.getOrElse(field) + }) + val factory = StorageFilterTestReaderFactory( + schema, + columnar = options.getBoolean("columnar", false), + giveUp = options.getBoolean("giveUp", false), + marksAllChecked = options.getBoolean("markAllChecked", false), + dropChecked = options.getOrDefault("dropChecked", ""), + checkedLayout = options.getOrDefault("checkedLayout", "")) + val declared = Utils.stringToSeq(options.getOrDefault("declared", "k")) + val numPartitions = options.getInt("numPartitions", 1) + options.getOrDefault("runtimeFiltering", "") match { + case "v2" => new StorageFilterTestScan(schema, factory, declared, numPartitions) + with SupportsRuntimeV2Filtering { + override def filterAttributes(): Array[NamedReference] = + readSchema().fieldNames.map(FieldReference.column).toArray[NamedReference] + override def filter(predicates: Array[Predicate]): Unit = + StorageFilterTestSource.runtimeFilterCalls.incrementAndGet() + override def supportsIterativePushdown(): Boolean = true + } + case "catalyst" => new StorageFilterTestScan(schema, factory, declared, numPartitions) + with SupportsRuntimeCatalystFiltering { + override def filterAttributes(): Array[NamedReference] = Array(FieldReference.column("k")) + override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes() + override def filter(expressions: Array[Expression]): Unit = + StorageFilterTestSource.runtimeFilterCalls.incrementAndGet() + } + case _ => new StorageFilterTestScan(schema, factory, declared, numPartitions) + } + } + + override def pushLimit(limit: Int): Boolean = pushes("limit") + override def pushOffset(offset: Int): Boolean = pushes("offset") + override def pushTopN(orders: Array[SortOrder], limit: Int): Boolean = pushes("topN") + override def pushTableSample( + lowerBound: Double, + upperBound: Double, + withReplacement: Boolean, + seed: Long, + sampleMethod: SampleMethod): Boolean = pushes("sample") + override def isPartiallyPushed: Boolean = false +} + +class StorageFilterTestScan( + schema: StructType, + factory: StorageFilterTestReaderFactory, + declared: Seq[String], + numPartitions: Int) + extends Scan with Batch with SupportsReportStatistics with SupportsStorageFilters { + + override def readSchema(): StructType = schema + + override def toBatch: Batch = this + + // Over the suite's application-side scan size threshold, so InjectRuntimeFilter considers this + // side worth filtering. + override def estimateStatistics(): Statistics = new Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(10L * 1024 * 1024) + override def numRows(): OptionalLong = OptionalLong.of(StorageFilterTestSource.numRows.toLong) + } + + override def storageFilterAttributes(): Array[NamedReference] = + declared.map(name => FieldReference(name.split('.').toSeq)).toArray[NamedReference] + + override def planInputPartitions(): Array[InputPartition] = { + val n = StorageFilterTestSource.numRows + (0 until numPartitions).map { i => + RangeInputPartition(i * n / numPartitions, (i + 1) * n / numPartitions) + }.toArray + } + + override def createReaderFactory(): PartitionReaderFactory = factory +} + +case class StorageFilterTestReaderFactory( + schema: StructType, + columnar: Boolean, + giveUp: Boolean, + marksAllChecked: Boolean, + dropChecked: String, + checkedLayout: String) + extends PartitionReaderFactory { + + override def supportColumnarReads(partition: InputPartition): Boolean = columnar + + // The filter the reader applies, none when it gives the filter up or only marks the rows. + private def filterOf(request: StorageFilterRequest): Option[StorageFilter] = { + StorageFilterTestSource.filteredReaders.incrementAndGet() + StorageFilterTestSource.record(request.filter()) + if (request.appendsCheckedColumn()) { + StorageFilterTestSource.checkedColumnRequests.incrementAndGet() + } + if (giveUp || marksAllChecked) None else Some(request.filter()) + } + + override def createReader(partition: InputPartition): PartitionReader[InternalRow] = + new StorageFilterTestReader(keys(partition), schema, None, appendsChecked = false) + + override def createReader( + partition: InputPartition, + request: StorageFilterRequest): PartitionReader[InternalRow] = + new StorageFilterTestReader(keys(partition), schema, filterOf(request), + request.appendsCheckedColumn() && dropChecked != "true", marksAllChecked, + dropsCheckedLater = dropChecked == "later", checkedLayout) + + override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = + new StorageFilterTestColumnarReader(keys(partition), schema, None, appendsChecked = false) + + override def createColumnarReader( + partition: InputPartition, + request: StorageFilterRequest): PartitionReader[ColumnarBatch] = + new StorageFilterTestColumnarReader(keys(partition), schema, filterOf(request), + request.appendsCheckedColumn() && dropChecked != "true", marksAllChecked, + dropsCheckedLater = dropChecked == "later", checkedLayout) + + // The keys from `start` until `end`. + private def keys(partition: InputPartition): RangeInputPartition = + partition.asInstanceOf[RangeInputPartition] +} + +/** + * Serves the rows of `partition` row by row, holding the columns of `schema`, and drops any row + * `filter` rejects, with the checked column last when `appendsChecked`. The filter's row holds + * the values of its `references()`. `filter` is cleared when it cannot be applied. + * `marksAllChecked` marks every row checked. `dropsCheckedLater` appends the checked column to + * the first row only. `checkedLayout` places the checked column as the source's option says. + */ +class StorageFilterTestReader( + partition: RangeInputPartition, + schema: StructType, + private var filter: Option[StorageFilter], + appendsChecked: Boolean, + marksAllChecked: Boolean = false, + dropsCheckedLater: Boolean = false, + checkedLayout: String = "") + extends PartitionReader[InternalRow] { + + private val referenceNames = filter.map(StorageFilterTestSource.columnsOf).getOrElse(Nil) + private var current: Long = partition.start - 1L + private var row: InternalRow = _ + private var returnedAny = false + + override def next(): Boolean = { + var found = false + while (!found && current + 1 < partition.end) { + current += 1 + val key = current + val passes = filter.forall { f => + StorageFilterTestSource.rowsEvaluated.incrementAndGet() + val keyRow = new GenericInternalRow( + referenceNames.map(StorageFilterTestSource.valueOf(_, key)).toArray[Any]) + try f.eval(keyRow) catch { + case _: StorageFilterNotAppliedException => + filter = None + true + } + } + if (passes) { + val values = schema.fields.map(StorageFilterTestSource.valueOf(_, key)) + val checked = filter.isDefined || marksAllChecked + val withChecked = appendsChecked && !(dropsCheckedLater && returnedAny) + row = new GenericInternalRow( + if (!withChecked) { + values + } else { + checkedLayout match { + case "front" => checked +: values Review Comment: Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). The docs say `front` is for columnar reads, and the row reader rejects it. ########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsStorageFilters.java: ########## @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector.read; + +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.connector.expressions.NamedReference; + +/** + * A mix-in interface for {@link Scan}. A data source implements it if its readers can evaluate a + * {@link StorageFilter}, to skip reading the other columns of the rows the filter rejects. The + * readers get the filter in the task, through + * {@link PartitionReaderFactory#createReader(InputPartition, StorageFilterRequest)} or + * {@link PartitionReaderFactory#createColumnarReader(InputPartition, StorageFilterRequest)}, which + * the scan's reader factory must implement. Review Comment: Thanks, the description's possible follow-ups now name the V2 file sources, with the three gaps you list. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [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]
