dongjoon-hyun commented on code in PR #59715: URL: https://github.com/apache/spark/pull/59715#discussion_r4383481085
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/StorageFilterPlanning.scala: ########## @@ -0,0 +1,102 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.datasources + +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, Expression, ExpressionSet, GetStructField, If, Literal} +import org.apache.spark.sql.execution.{ScalarSubquery => ExecScalarSubquery, WholeStageCodegenExec} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.BooleanType + +/** + * The rules file-source and DSv2 scans share for storage filters, see + * `FileFormat.buildReaderWithStorageFilters` and `SupportsStorageFilters`. Each strategy decides + * on its own whether its scan can take any. + */ +object StorageFilterPlanning { + + /** + * The conjuncts of `filters` offered to the scan, each once, as an `ExpressionSet` counts them. + * Each one: + * - is deterministic, since the post-scan `Filter` takes the reader's word on the rows it + * checked; + * - references at least one column, and only `keyCandidates`; + * - is taken by `accepts`. + * + * Nothing is offered when the offered conjuncts reference every one of `readColumns`, since the + * reader then reads the same columns for the same rows either way. + */ + def offeredConjuncts( + filters: Seq[Expression], + keyCandidates: AttributeSet, + readColumns: AttributeSet, + accepts: Expression => Boolean): Seq[Expression] = { + val offered = ExpressionSet(filters).toSeq.filter { expr => + val refs = expr.references + expr.deterministic && refs.nonEmpty && refs.subsetOf(keyCandidates) && accepts(expr) + } + if (offered.isEmpty) return Nil + val storageKeyAttrs = AttributeSet.fromAttributeSets(offered.map(_.references)) + if (readColumns.subsetOf(storageKeyAttrs)) Nil else offered + } + + /** + * The column the reader marks the rows it checked in, when a filter is `offered`. The caller + * places it in the scan's output. Only where whole-stage codegen is on and `output` with the + * column still fits it, so the column keeps the scan in whole-stage codegen. Elsewhere the + * Project that drops the column could copy every surviving row. A post-scan conjunct that cannot + * be generated still puts that Project in a stage of its own, the trade-off SPARK-60034 accepted. + */ + def checkedColumn( + conf: SQLConf, + offered: Seq[Expression], + output: Seq[Attribute]): Option[AttributeReference] = { + if (offered.isEmpty) return None + val checked = AttributeReference(FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME, BooleanType)() + Option.when(conf.wholeStageEnabled && + !WholeStageCodegenExec.isTooManyFields(conf, (output :+ checked).toStructType))(checked) + } + + /** + * `filters` with each of `offered` skipped on the rows `checked` marks, or as they are without + * the column. Matched by `ExpressionSet`, as `offered` was built. An `If` keeps two things an + * `Or` would lose: + * - a null mark leaves the conjunct as it was, an `And` around it included; + * - subexpression elimination cannot hoist what the conjunct shares with another one ahead of + * the skip. + */ + def skipCheckedRows( + filters: Seq[Expression], + offered: Seq[Expression], + checked: Option[Attribute]): Seq[Expression] = checked match { + case Some(c) => + val offeredSet = ExpressionSet(offered) + filters.map(f => if (offeredSet.contains(f)) If(c, Literal.TrueLiteral, f) else f) + case None => filters + } + + /** + * `filter` with each scalar subquery replaced by its result, for the readers. A field read out of + * a result becomes a literal of its own. A subquery `MergeSubplans` merged returns the results of + * several subqueries in one struct, such as the blooms of several filters, so a field read out of + * it would otherwise carry the whole struct. + */ + def withSubqueryResults(filter: Expression): Expression = filter.transformUp { + case s: ExecScalarSubquery => s.toLiteral Review Comment: **1. [Low] Each conjunct still copies every bloom of a merged subquery on the driver** This is the shape I suggested in item 4 of my second review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4379088708)), and it has a cost I missed. `transformUp` turns the `ScalarSubquery` leaf into `s.toLiteral` before it reaches the `GetStructField` above it, and `toLiteral` is `Literal.create(result, dataType)` over the whole merged struct (subquery.scala:112-115). `Literal`'s constructor always validates its value (literals.scala:499), and for a struct it reads every field with `row.get(i, fieldDataType)` (literals.scala:286-292). `result` is a view into the `UnsafeRow` the subquery collected (subquery.scala:95), so each of those reads is `UnsafeRow.getBinary`, which copies the whole bloom (UnsafeRow.java:454-455), only for the copy to be dropped after a type check. So a conjunct over a merge of N blooms makes N throwaway copies before the one it keeps, and a scan that takes all N blooms allocates N x (N + 1) bloom-sized arrays on the driver each time it runs. In the two-key tests that is 4 MiB of garbage for two 1 MiB blooms. `BatchScanExec.storageFilterRequest` (BatchScanExec.scala:130-131) adds this, and `FileSourceScanExec.preparedStorageFilters` already paid the N x N part on master through its own `s.toLiteral`. Reading the field straight from the subquery never builds the struct literal: ```scala def withSubqueryResults(filter: Expression): Expression = filter.transformDown { case g @ GetStructField(_: ExecScalarSubquery | _: Literal, _, _) => Literal(g.eval(), g.dataType) case s: ExecScalarSubquery => s.toLiteral } ``` `g.eval()` reaches `ScalarSubquery.eval`, which returns `result` without building a `Literal`, so only `GetStructField` copies its one field. Unlike `transformUp`, this folds only the inner read of a nested `GetStructField(GetStructField(...))`, but `MergeSubplans` builds single-level reads only (MergeSubplans.scala:331-335). ########## sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterImplSuite.scala: ########## @@ -0,0 +1,172 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import org.apache.spark.{SparkArithmeticException, SparkConf, SparkException, SparkFunSuite, SparkUnsupportedOperationException, TaskContext} +import org.apache.spark.serializer.JavaSerializer +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, Cast, Divide, EqualTo, EvalMode, Expression, GreaterThan, Literal, NumericEvalContext, Remainder, UnaryExpression} +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory, StorageFilterNotAppliedException} +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector +import org.apache.spark.sql.internal.connector.{StorageFilterImpl, StorageFilterRequestImpl} +import org.apache.spark.sql.types.{BooleanType, DataType, DoubleType, LongType} +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** Tests the Catalyst-backed [[StorageFilterImpl]] a connector's reader evaluates. */ +class StorageFilterImplSuite extends SparkFunSuite { + + private val k = AttributeReference("k", LongType)() + private val v = AttributeReference("v", LongType)() + + // `k % 2 = 0` over a scan whose output has `v` before `k`. + private def evenKeys: StorageFilterImpl = + StorageFilterImpl(EqualTo(Remainder(k, Literal(2L)), Literal(0L)), Seq(v, k)) + + // `1 / k > 0`, which raises an error at `k = 0` under ANSI. + private def positiveInverse: StorageFilterImpl = StorageFilterImpl( + GreaterThan( + Divide(Literal(1.0), Cast(k, DoubleType), NumericEvalContext(EvalMode.ANSI)), + Literal(0.0)), + Seq(k)) + + private def withBatchOf(keys: Long*)(body: ColumnarBatch => Unit): Unit = { + val vector = new OnHeapColumnVector(keys.length, LongType) + keys.zipWithIndex.foreach { case (key, i) => vector.putLong(i, key) } + val batch = new ColumnarBatch(Array(vector), keys.length) + try body(batch) finally batch.close() + } + + test("the references are the columns the filter reads") { + assert(evenKeys.references.map(_.toString).toSeq == Seq("k")) + } + + test("a filter over two columns takes them in the scan output's order") { + val filter = + StorageFilterImpl(And(EqualTo(k, Literal(1L)), EqualTo(v, Literal(10L))), Seq(v, k)) + assert(filter.references.map(_.toString).toSeq == Seq("v", "k")) + assert(filter.eval(InternalRow(10L, 1L)), "the row holds v, then k") + assert(!filter.eval(InternalRow(1L, 10L))) + } + + test("a reader that reorders the references changes the layout of no other reader") { + val filter = + StorageFilterImpl(And(EqualTo(k, Literal(1L)), EqualTo(v, Literal(10L))), Seq(v, k)) + val references = filter.references() + references.sortInPlaceBy(_.toString) + assert(filter.references().map(_.toString).toSeq == Seq("v", "k")) + } + + test("a row of another width than the references fails, rather than being misread") { + // `evenKeys` reads `k` alone, so a row holding `v, k` would put `v` where `k` is read. + val e = intercept[SparkException](evenKeys.eval(InternalRow(3L, 4L))) + assert(e.getCondition == "INTERNAL_ERROR", s"got ${e.getCondition}") + } + + test("a filter over a column outside the scan output is an internal error") { + val e = intercept[SparkException](StorageFilterImpl(EqualTo(k, Literal(1L)), Seq(v))) + assert(e.getCondition == "INTERNAL_ERROR", s"got ${e.getCondition}") + } + + test("select writes the kept rows' positions in ascending order") { + withBatchOf(3L, 4L, 5L, 6L, 8L) { batch => + val selected = new Array[Int](batch.numRows) + val count = evenKeys.select(batch, selected) + assert(selected.take(count).toSeq == Seq(1, 3, 4)) + } + } + + test("an error evaluating the filter answers that it is not applied") { + val filter = positiveInverse + assert(filter.eval(InternalRow(2L))) + val e = intercept[StorageFilterNotAppliedException](filter.eval(InternalRow(0L))) + assert(e.getCondition == "STORAGE_FILTER_NOT_APPLIED", s"got ${e.getCondition}") + assert(e.getCause.isInstanceOf[SparkArithmeticException], s"got ${e.getCause}") + withBatchOf(1L, 0L) { batch => + intercept[StorageFilterNotAppliedException](filter.select(batch, new Array[Int](2))) + } + } + + test("a reader factory without the storage-filter methods cannot create a filtering reader") { + val factory = new PartitionReaderFactory { + override def createReader(partition: InputPartition): PartitionReader[InternalRow] = null + } + val request = StorageFilterRequestImpl(evenKeys, appendsCheckedColumn = false) + Seq( + () => factory.createReader(RangeInputPartition(0, 1), request), + () => factory.createColumnarReader(RangeInputPartition(0, 1), request) + ).foreach { create => + val e = intercept[SparkUnsupportedOperationException](create()) + assert(e.getCondition == "DATA_SOURCE_STORAGE_FILTER_READER_NOT_SUPPORTED", Review Comment: **5. [Nit] The SQLSTATEs of the two new conditions, and the `class` parameter, are checked nowhere** This line and L98 compare only `e.getCondition`, and no other test raises either condition, since the test readers catch `StorageFilterNotAppliedException` (StorageFilterTestSource.scala:372 and 457). So the defaults at PartitionReaderFactory.java:70 and 84 could pass another class as `class`, e.g. `partition.getClass().getName()`, or this condition's `0A000` could become the `KD000` of the entry above it, and every test would still pass. The same holds for the `KD000` the description gives `STORAGE_FILTER_NOT_APPLIED`. As in item 10 of my second review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4380305050)), could both use `checkError`? ```scala checkError( exception = intercept[SparkUnsupportedOperationException](create()), condition = "DATA_SOURCE_STORAGE_FILTER_READER_NOT_SUPPORTED", sqlState = Some("0A000"), parameters = Map("class" -> factory.getClass.getName)) ``` and `checkError(exception = e, condition = "STORAGE_FILTER_NOT_APPLIED", sqlState = Some("KD000"))` at L98. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala: ########## @@ -124,9 +124,10 @@ case class GroupPartitionsExec( * where this check passes, and `UnknownPartitioning`, where re-deriving reports the same thing. * * The comparison is by reference for a child that reports the same object, which the columnar and - * codegen wrappers do. `PlanAdaptiveDynamicPruningFilters` rebuilds `BatchScanExec`, so there it - * compares the partition keys by value, and reads equal only while the connector returns the same - * input partitions for the same scan. + * codegen wrappers do. `PlanAdaptiveDynamicPruningFilters` and `ReuseAdaptiveSubquery` rebuild a + * `BatchScanExec` whose runtime or storage filters hold a subquery, so there it compares the + * partition keys by value, and reads equal only while the connector returns the same input + * partitions for the same scan. Review Comment: **6. [Nit] Without AQE, `ReuseExchangeAndSubquery` rebuilds the scan too** This sentence, from the answer to item 7 of my first review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009038)), names only AQE rules. Without AQE, `ReuseExchangeAndSubquery` runs after `EnsureRequirements` (QueryExecution.scala:848 and 886). It walks the plan bottom-up and keeps the first of two equal subqueries, so a scan that holds one keeps it, but a second one in the same scan becomes a `ReusedSubqueryExec` (ReuseExchangeAndSubquery.scala:53-67), and the scan is copied under the `GroupPartitionsExec` that is already planned. A storage-filtered scan gets there in two ways: two blooms that `MergeSubplans` merged both sit in `storageFilters`, as in the two-key test, or, with dynamic partition pruning off, one bloom on a key the scan declares for both `SupportsRuntimeV2Filtering` and `SupportsStorageFilters` sits in `runtimeFilters` and `storageFilters`. Then `checkChildStillMatches` compares the partition keys by value too. Maybe: ```suggestion * codegen wrappers do. A rule that rewrites a subquery in the runtime or storage filters of a * `BatchScanExec` after `EnsureRequirements` rebuilds it, so there it compares the partition * keys by value, and reads equal only while the connector returns the same input partitions for * the same scan. Such rules are `PlanAdaptiveDynamicPruningFilters` and `ReuseAdaptiveSubquery` * with AQE, and `ReuseExchangeAndSubquery` without it. ``` ########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala: ########## @@ -1060,6 +1061,46 @@ class ParquetStorageFilterSuite extends QueryTest with SharedSparkSession } } + Seq(false, true).foreach { aqe => + test(s"a two-key join ships each bloom to the reader once (AQE $aqe)") { + // `MergeSubplans` merges the two blooms' aggregates into one subquery returning both blooms. + withTable("bf3", "bf4") { + spark.range(600).selectExpr("id AS k", "id * 10 AS v", "id * 100 AS w") + .write.format("parquet").saveAsTable("bf3") + spark.range(30).selectExpr("id AS k", "id * 10 AS v", "id AS flag") + .write.format("parquet").saveAsTable("bf4") + withSQLConf( + SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true", + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + val df = spark.sql( + "SELECT bf3.k, bf3.v, bf3.w FROM bf3 JOIN bf4 ON bf3.k = bf4.k AND bf3.v = bf4.v " + + "WHERE bf4.flag = 5") + val rows = df.collect().map(r => (r.getLong(0), r.getLong(1), r.getLong(2))).toSeq + assert(rows == Seq((5L, 50L, 500L)), s"got $rows") + val plan = df.queryExecution.executedPlan + val scan = collectFirst(plan) { + case s: FileSourceScanExec if s.storageFilters.nonEmpty => s + }.getOrElse(fail(s"no scan with storage filters:\n$plan")) + assert(scan.storageFilters.size == 2 && scan.storageFilters.forall(_.exists { + case GetStructField(_: ExecSubqueryExpression, _, _) => true + case _ => false + }), s"test setup: both blooms must come from one merged subquery:\n$plan") + val bloomBytes = scan.storageFilters.map(_.collectFirst { + case m: BloomFilterMightContain => + m.bloomFilterExpression.eval().asInstanceOf[Array[Byte]].length + }.get).sum + // The reader's function ships with the scan's RDD, and the blooms with it. + val shipped = Utils.serialize(scan.inputRDD).length + assert(shipped < bloomBytes * 3 / 2, Review Comment: **4. [Nit] The new V1 two-key test passes when no bloom reaches the reader** Thank you for adding it for item 3 of my second review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4379088695)). It bounds the shipped bytes only from above, so it also passes when the reader gets no filter at all: if `storageFilterOpt` were `None` on this path (ParquetFileFormat.scala:294-304), `shipped` would drop to a few KB. The answer (L1081) and the planning check (L1086-1089) hold either way, and so does this bound when the reader gets the blooms but gives them up. As the test above puts it, "The metrics are what say it did" (L1055-1058). Could this one also assert `rowsExcludedByStorageFilters(plan) > 0` (L988-994)? ########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsStorageFilters.java: ########## @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector.read; + +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.connector.expressions.NamedReference; + +/** + * A mix-in interface for {@link Scan}. A data source implements it if its readers can evaluate a + * {@link StorageFilter}, to skip reading the other columns of the rows the filter rejects. The + * readers get the filter in the task, through + * {@link PartitionReaderFactory#createReader(InputPartition, StorageFilterRequest)} or + * {@link PartitionReaderFactory#createColumnarReader(InputPartition, StorageFilterRequest)}, which + * the scan's reader factory must implement. Review Comment: **12. [Nit] `FilePartitionReaderFactory` has no route for a request** Spark's own base for file-based scans, `FilePartitionReaderFactory`, implements only the plain `createReader` and `createColumnarReader` (FilePartitionReaderFactory.scala:29-39). So a `FileScan` that adds `SupportsStorageFilters` fails every filtered task with `DATA_SOURCE_STORAGE_FILTER_READER_NOT_SUPPORTED` unless its factory wraps `FilePartitionReader` itself. That is what this sentence asks for, and no built-in `FileScan` implements the mix-in, so nothing fails today. But an adopter such as the V2 Parquet scan would also have to place the checked column after the partition values, which `PartitionReaderWithPartitionValues` appends last, and the metrics that L47-48 point to would not reach Spark, since `FilePartitionReader`, `PartitionedFileReader` and `PartitionReaderWithPartitionValues` forward only `next`, `get` and `close`. That last part is the same on master. Could the description's list of possible next steps name the V2 file sources? ########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/StorageFilterRequest.java: ########## @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector.read; + +import java.io.Serializable; + +import org.apache.spark.annotation.Evolving; + +/** + * What Spark hands a reader of a {@link SupportsStorageFilters} scan, through + * {@link PartitionReaderFactory#createReader(InputPartition, StorageFilterRequest)} or + * {@link PartitionReaderFactory#createColumnarReader(InputPartition, StorageFilterRequest)}. + * + * @since 5.0.0 + */ +@Evolving +public interface StorageFilterRequest extends Serializable { + + /** The filter the reader may apply. */ + StorageFilter filter(); + + /** + * Whether the reader must append a boolean column after the columns of + * {@link Scan#readSchema()}. It is true only on a row the filter was evaluated on and kept. It is + * false on every other row the reader returns, such as a row read without the filter, or one + * kept although the filter rejected it. Spark skips the filter's predicate above the scan on the + * rows marked true. A reader asked for the column that returns a row or a batch without it fails + * the task. The reader may return more columns after the checked column, which Spark does not + * read. Review Comment: **2. [Low] The `appendsCheckedColumn` doc still promises a failure the new check no longer gives** L42-43 say "A reader asked for the column that returns a row or a batch without it fails the task", and the new sentence after it lets the reader return more columns. Since the fix for item 1 of my second review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4379088674)), `DataSourceRDD` checks only that a row is wide enough (DataSourceRDD.scala:235), and that a batch is wide enough and holds a boolean at `width - 1` (DataSourceRDD.scala:229). A width check cannot tell the checked column from the reader's own first extra column, so a reader that leaves the column out but returns more columns passes it. Take a scan that reads `k, v` (width 3), and a reader whose new overload forwards to its plain reader, the easy mistake of item 2 of my first review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009016)), returning `k, v, _pos: long, _deleted: boolean`. At `cb01a6f`, `4 != 3` failed the task with `DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING`, as documented. Now the row passes, and `If(checked, true, bloom)` reads ordinal 2. From an `UnsafeRow` that is a byte of `_pos`, so the `Filter` skips the bloom on arbitrary rows with no error. From a `GenericInternalRow` it is a `ClassCastException` that names neither the connector nor the request. A columnar `k, v, deleted: boolean, pos` passes too, with `deleted` as the mark. The join still compares the keys, so a bloom's answer stays right. Could the doc say what Spark checks, e.g. ```suggestion * rows marked true. The reader may return more columns after the checked column, which Spark * does not read. Spark checks only that a row or a batch is wide enough for the column, and that * a batch holds a boolean there, and fails the task otherwise. A reader that leaves the column * out but returns more columns can pass that check, and Spark then reads the column at that * position as the checked column. ``` The description's "fails the task with `DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING`" could say the same. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterTestSource.scala: ########## @@ -0,0 +1,492 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import java.util +import java.util.OptionalLong +import java.util.concurrent.atomic.{AtomicInteger, AtomicLong} + +import scala.collection.mutable +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow} +import org.apache.spark.sql.connector.catalog.{MetadataColumn, SupportsMetadataColumns, SupportsRead, Table, TableCapability, TableProvider} +import org.apache.spark.sql.connector.expressions.{Expressions, FieldReference, NamedReference, SortOrder, Transform} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read._ +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector +import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering +import org.apache.spark.sql.types.{BooleanType, DataType, LongType, StructField, StructType, VariantType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +/** + * A read-only v2 source over a fixed table of longs that implements [[SupportsStorageFilters]]: + * the key `k`, the value `v` holding `k * 10`, and with `wide` also `w` holding `k * 100`. Its + * reader gets the filter in the task and drops the rows the filter rejects. It reads the columns + * Spark requires, and builds the filter's row from the filter's `references()`. Its options: + * - `columnar`: reads in batches, through `StorageFilter.select`, rather than row by row; + * - `giveUp`: reads without the filter, as a reader does for a unit it cannot use the filter on. + * Here the unit is the partition; + * - `markAllChecked`: reads without the filter but marks every row checked, which breaks the + * contract on purpose; + * - `dropChecked`: `true` never appends the checked column, and `later` appends it to the first + * row or batch only. Either breaks the contract on purpose; + * - `checkedLayout`: `front` returns the checked column ahead of the scan's columns, which breaks + * the contract on purpose. `extra` returns one more column after it; + * - `declared`: the comma-separated storage-filter columns, `k` by default, a dot separating a + * nested field; + * - `valueName`: the name of the value column, `v` by default; + * - `wide`: adds the column `w`; + * - `variant`: adds the variant column `j`, which the scan reads only through pushed extractions + * of a long. Each one holds `k * 10`; + * - `numPartitions`: splits the keys into that many input partitions, 1 by default; + * - `metadataColumn`: adds the metadata column `_pos`, which holds the key; + * - `runtimeFiltering`: `v2` makes the scan a `SupportsRuntimeV2Filtering` one over every column, + * with iterative pushdown on a table partitioned by `k`. `catalyst` makes it a + * `SupportsRuntimeCatalystFiltering` one that evaluates the filters on `k` fully. Neither + * applies the filters it takes, so they serve planning tests only; + * - `push`: the comma-separated pushed operations it takes whole without applying them, of + * `limit`, `offset`, `topN` and `sample`, for planning tests only. + * + * A filter that throws `StorageFilterNotAppliedException` is given up for the rest of the + * partition too. The counters below record what the readers got, evaluated and returned. + */ +class StorageFilterTestSource extends TableProvider { + + override def inferSchema(options: CaseInsensitiveStringMap): StructType = { + val schema = new StructType().add("k", LongType).add(options.getOrDefault("valueName", "v"), + LongType) + val wide = if (options.getBoolean("wide", false)) schema.add("w", LongType) else schema + if (options.getBoolean("variant", false)) wide.add("j", VariantType) else wide + } + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: util.Map[String, String]): Table = { + val options = new CaseInsensitiveStringMap(properties) + if (options.getBoolean("metadataColumn", false)) { + new StorageFilterTestTable(schema, options) with SupportsMetadataColumns { + override def metadataColumns(): Array[MetadataColumn] = Array(new MetadataColumn { + override def name(): String = StorageFilterTestSource.posName + override def dataType(): DataType = LongType + }) + } + } else { + new StorageFilterTestTable(schema, options) + } + } +} + +object StorageFilterTestSource { + + /** Rows the source serves, `k` from 0 until numRows. */ + val numRows: Int = 200 + + /** Rows a batch holds at most in a columnar read. */ + val batchSize: Int = 64 + + /** The name of the metadata column the `metadataColumn` option adds. */ + val posName: String = "_pos" + + /** How many readers were created with a storage-filter request. */ + val filteredReaders = new AtomicInteger(0) + + /** How many of those requests asked for the checked column. */ + val checkedColumnRequests = new AtomicInteger(0) + + /** Rows the reader evaluated the storage filter against. */ + val rowsEvaluated = new AtomicLong(0) + + /** Rows the reader returned. */ + val rowsReturned = new AtomicLong(0) + + /** The serialized size of the first filter a reader got, before it evaluated the filter. */ + val filterBytes = new AtomicLong(0) + + /** The references of the filters the readers got, and the identities of those filters. */ + val filterReferences = mutable.Set.empty[Seq[String]] + val filterIdentities = mutable.Set.empty[Int] + + /** How many times a runtime-filtering scan was handed filters. */ + val runtimeFilterCalls = new AtomicInteger(0) + + def reset(): Unit = { + filteredReaders.set(0) + checkedColumnRequests.set(0) + rowsEvaluated.set(0) + rowsReturned.set(0) + filterBytes.set(0) + filterReferences.synchronized(filterReferences.clear()) + filterIdentities.synchronized(filterIdentities.clear()) + runtimeFilterCalls.set(0) + } + + /** The value of column `name` in the row of `key`. Any column but `k`, `w` and `_pos` is `v`. */ + def valueOf(name: String, key: Long): Long = name match { + case "k" | `posName` => key + case "w" => key * 100 + case _ => key * 10 + } + + /** The value of `field` in the row of `key`, a struct for the extractions of a variant. */ + def valueOf(field: StructField, key: Long): Any = field.dataType match { + case struct: StructType => InternalRow.fromSeq(struct.map(_ => valueOf(field.name, key))) + case _ => valueOf(field.name, key) + } + + /** The columns of the row the filter evaluates, in that order. */ + def columnsOf(filter: StorageFilter): Seq[String] = filter.references.map(_.fieldNames.head).toSeq + + // Records the filter a reader got, before the reader evaluates it. Only the first reader + // serializes it, since a later reader of the same task may get it evaluated, see `StorageFilter`. + private[connector] def record(filter: StorageFilter): Unit = { + if (filterBytes.get() == 0) filterBytes.compareAndSet(0, Utils.serialize(filter).length) + filterReferences.synchronized(filterReferences += columnsOf(filter)) + filterIdentities.synchronized(filterIdentities += System.identityHashCode(filter)) + } +} + +class StorageFilterTestTable(tableSchema: StructType, tableOptions: CaseInsensitiveStringMap) + extends Table with SupportsRead { + + override def name(): String = "storage_filter_test" + + override def schema(): StructType = tableSchema + + override def partitioning(): Array[Transform] = + if (tableOptions.getOrDefault("runtimeFiltering", "") == "v2") { + Array(Expressions.identity("k")) + } else { + Array.empty + } + + override def capabilities(): util.Set[TableCapability] = + Set(TableCapability.BATCH_READ).asJava + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = + new StorageFilterTestScanBuilder(tableSchema, options) +} + +/** + * Prunes the columns Spark does not read, and takes the pushed operations named in `push` whole. + * Its scan does not apply them, so they serve planning tests only. It takes every variant + * extraction, and reads each into a struct field of the variant's column named by its position. + */ +class StorageFilterTestScanBuilder(tableSchema: StructType, options: CaseInsensitiveStringMap) + extends SupportsPushDownRequiredColumns + with SupportsPushDownLimit + with SupportsPushDownOffset + with SupportsPushDownTopN + with SupportsPushDownTableSample + with SupportsPushDownVariantExtractions { + + private val pushes = Utils.stringToSeq(options.getOrDefault("push", "")).toSet + private var required = tableSchema + private var extractions = Array.empty[VariantExtraction] + + override def pruneColumns(requiredSchema: StructType): Unit = required = requiredSchema + + override def pushVariantExtractions(pushed: Array[VariantExtraction]): Array[Boolean] = { + extractions = pushed + pushed.map(_ => true) + } + + override def build(): Scan = { + val byColumn = extractions.groupBy(_.columnName().head) + val schema = StructType(required.map { field => + byColumn.get(field.name).map { columnExtractions => + field.copy(dataType = StructType(columnExtractions.zipWithIndex.map { case (e, i) => + StructField(i.toString, e.expectedDataType(), nullable = true, e.metadata()) + })) + }.getOrElse(field) + }) + val factory = StorageFilterTestReaderFactory( + schema, + columnar = options.getBoolean("columnar", false), + giveUp = options.getBoolean("giveUp", false), + marksAllChecked = options.getBoolean("markAllChecked", false), + dropChecked = options.getOrDefault("dropChecked", ""), + checkedLayout = options.getOrDefault("checkedLayout", "")) + val declared = Utils.stringToSeq(options.getOrDefault("declared", "k")) + val numPartitions = options.getInt("numPartitions", 1) + options.getOrDefault("runtimeFiltering", "") match { + case "v2" => new StorageFilterTestScan(schema, factory, declared, numPartitions) + with SupportsRuntimeV2Filtering { + override def filterAttributes(): Array[NamedReference] = + readSchema().fieldNames.map(FieldReference.column).toArray[NamedReference] + override def filter(predicates: Array[Predicate]): Unit = + StorageFilterTestSource.runtimeFilterCalls.incrementAndGet() + override def supportsIterativePushdown(): Boolean = true + } + case "catalyst" => new StorageFilterTestScan(schema, factory, declared, numPartitions) + with SupportsRuntimeCatalystFiltering { + override def filterAttributes(): Array[NamedReference] = Array(FieldReference.column("k")) + override def fullyPushedFilterAttributes(): Array[NamedReference] = filterAttributes() + override def filter(expressions: Array[Expression]): Unit = + StorageFilterTestSource.runtimeFilterCalls.incrementAndGet() + } + case _ => new StorageFilterTestScan(schema, factory, declared, numPartitions) + } + } + + override def pushLimit(limit: Int): Boolean = pushes("limit") + override def pushOffset(offset: Int): Boolean = pushes("offset") + override def pushTopN(orders: Array[SortOrder], limit: Int): Boolean = pushes("topN") + override def pushTableSample( + lowerBound: Double, + upperBound: Double, + withReplacement: Boolean, + seed: Long, + sampleMethod: SampleMethod): Boolean = pushes("sample") + override def isPartiallyPushed: Boolean = false +} + +class StorageFilterTestScan( + schema: StructType, + factory: StorageFilterTestReaderFactory, + declared: Seq[String], + numPartitions: Int) + extends Scan with Batch with SupportsReportStatistics with SupportsStorageFilters { + + override def readSchema(): StructType = schema + + override def toBatch: Batch = this + + // Over the suite's application-side scan size threshold, so InjectRuntimeFilter considers this + // side worth filtering. + override def estimateStatistics(): Statistics = new Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(10L * 1024 * 1024) + override def numRows(): OptionalLong = OptionalLong.of(StorageFilterTestSource.numRows.toLong) + } + + override def storageFilterAttributes(): Array[NamedReference] = + declared.map(name => FieldReference(name.split('.').toSeq)).toArray[NamedReference] + + override def planInputPartitions(): Array[InputPartition] = { + val n = StorageFilterTestSource.numRows + (0 until numPartitions).map { i => + RangeInputPartition(i * n / numPartitions, (i + 1) * n / numPartitions) + }.toArray + } + + override def createReaderFactory(): PartitionReaderFactory = factory +} + +case class StorageFilterTestReaderFactory( + schema: StructType, + columnar: Boolean, + giveUp: Boolean, + marksAllChecked: Boolean, + dropChecked: String, + checkedLayout: String) + extends PartitionReaderFactory { + + override def supportColumnarReads(partition: InputPartition): Boolean = columnar + + // The filter the reader applies, none when it gives the filter up or only marks the rows. + private def filterOf(request: StorageFilterRequest): Option[StorageFilter] = { + StorageFilterTestSource.filteredReaders.incrementAndGet() + StorageFilterTestSource.record(request.filter()) + if (request.appendsCheckedColumn()) { + StorageFilterTestSource.checkedColumnRequests.incrementAndGet() + } + if (giveUp || marksAllChecked) None else Some(request.filter()) + } + + override def createReader(partition: InputPartition): PartitionReader[InternalRow] = + new StorageFilterTestReader(keys(partition), schema, None, appendsChecked = false) + + override def createReader( + partition: InputPartition, + request: StorageFilterRequest): PartitionReader[InternalRow] = + new StorageFilterTestReader(keys(partition), schema, filterOf(request), + request.appendsCheckedColumn() && dropChecked != "true", marksAllChecked, + dropsCheckedLater = dropChecked == "later", checkedLayout) + + override def createColumnarReader(partition: InputPartition): PartitionReader[ColumnarBatch] = + new StorageFilterTestColumnarReader(keys(partition), schema, None, appendsChecked = false) + + override def createColumnarReader( + partition: InputPartition, + request: StorageFilterRequest): PartitionReader[ColumnarBatch] = + new StorageFilterTestColumnarReader(keys(partition), schema, filterOf(request), + request.appendsCheckedColumn() && dropChecked != "true", marksAllChecked, + dropsCheckedLater = dropChecked == "later", checkedLayout) + + // The keys from `start` until `end`. + private def keys(partition: InputPartition): RangeInputPartition = + partition.asInstanceOf[RangeInputPartition] +} + +/** + * Serves the rows of `partition` row by row, holding the columns of `schema`, and drops any row + * `filter` rejects, with the checked column last when `appendsChecked`. The filter's row holds + * the values of its `references()`. `filter` is cleared when it cannot be applied. + * `marksAllChecked` marks every row checked. `dropsCheckedLater` appends the checked column to + * the first row only. `checkedLayout` places the checked column as the source's option says. + */ +class StorageFilterTestReader( + partition: RangeInputPartition, + schema: StructType, + private var filter: Option[StorageFilter], + appendsChecked: Boolean, + marksAllChecked: Boolean = false, + dropsCheckedLater: Boolean = false, + checkedLayout: String = "") + extends PartitionReader[InternalRow] { + + private val referenceNames = filter.map(StorageFilterTestSource.columnsOf).getOrElse(Nil) + private var current: Long = partition.start - 1L + private var row: InternalRow = _ + private var returnedAny = false + + override def next(): Boolean = { + var found = false + while (!found && current + 1 < partition.end) { + current += 1 + val key = current + val passes = filter.forall { f => + StorageFilterTestSource.rowsEvaluated.incrementAndGet() + val keyRow = new GenericInternalRow( + referenceNames.map(StorageFilterTestSource.valueOf(_, key)).toArray[Any]) + try f.eval(keyRow) catch { + case _: StorageFilterNotAppliedException => + filter = None + true + } + } + if (passes) { + val values = schema.fields.map(StorageFilterTestSource.valueOf(_, key)) + val checked = filter.isDefined || marksAllChecked + val withChecked = appendsChecked && !(dropsCheckedLater && returnedAny) + row = new GenericInternalRow( + if (!withChecked) { + values + } else { + checkedLayout match { + case "front" => checked +: values Review Comment: **11. [Nit] The row reader's `front` layout runs in no test** Only the columnar test sets `checkedLayout -> front` (StorageFilterPushDownSuite.scala:259), while the option's doc (L52-53) and this reader's doc (L345) offer it to both readers. A row-based `front` read would not end in the contract error either: `checkRow` checks only the width, and the `Filter` above would fail with a `ClassCastException`, as the comment at StorageFilterPushDownSuite.scala:255-256 suggests. Could the docs say `front` is for columnar reads, or this reader reject it? ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterPredicate.scala: ########## @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.internal.connector + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.expressions.{Attribute, BasePredicate, BindReferences, Expression, Predicate} + +/** + * The predicate of a storage filter, built on first use from `boundExpression`, since a generated + * one can be awkward to serialize. The file-source Parquet reader and [[StorageFilterImpl]] share + * it. + */ +class StorageFilterPredicate(private var boundExpression: Expression) extends Serializable { + + @transient lazy val prepared: BasePredicate = { + // The expression is released below, so a copy made after a first evaluation would have nothing + // to build from. That is named here rather than left to an NPE. + if (boundExpression == null) { + throw SparkException.internalError("A storage filter was copied after an evaluation " + Review Comment: **7. [Nit] The copied-filter error no longer names the filter** Before the move for item 12 of my second review ([thread](https://github.com/apache/spark/pull/59715#discussion_r4380305098)), this was `s"$this was copied after an evaluation released the expression of its predicate"`, which names the columns, as item 20 of my first review suggested ([thread](https://github.com/apache/spark/pull/59715#discussion_r4372009086)). `StorageFilterPredicate` does not know its filter, so the message is now a constant, while the width error just before it in `StorageFilterImpl.eval` (StorageFilterImpl.scala:57-60) still names the filter. With two storage-filtered scans in a stage, a connector author who hits it cannot tell which filter a reader copied. The V1 message also lost its hint, "a storage filter must only be evaluated on an executor". Could `StorageFilterPredicate` take a short description for this message, e.g. `StorageFilterImpl.toString` and the V1 filter's key columns? ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala: ########## @@ -0,0 +1,103 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.internal.connector + +import org.apache.spark.SparkException +import org.apache.spark.internal.{Logging, LogKeys} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression => CatalystExpression} +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference} +import org.apache.spark.sql.connector.read.{StorageFilter, StorageFilterNotAppliedException} + +/** + * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not + * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key + * rather than a data row. + * + * `attrs` is the layout of the row [[eval]] expects, in that order, and the predicate is bound to + * it. The predicate is built on first use, see [[StorageFilterPredicate]]. + */ +class StorageFilterImpl private ( + private val predicate: StorageFilterPredicate, + private val attrs: Seq[Attribute]) Review Comment: **9. [Nit] `attrs` is now kept only for its length and the column names** Since `apply` binds the expression on the driver (L101), `attrs` is read only by `numReferences` (L41) and to rebuild the column names in the transient lazy `referenceArray` (L76-77). So each task deserializes full attributes for their names, and the class holds the layout twice. Could the constructor take the references instead, e.g. `private val referenceArray: Array[NamedReference]` with `numReferences = referenceArray.length`, and `apply` pass `referenced.map(a => FieldReference.column(a.name)).toArray`? `references()` keeps its clone. The one-caller `StorageFilterPredicate.apply` could then go too, as `ParquetStorageFilter` builds `new StorageFilterPredicate(remapped)` itself (ParquetStorageFilter.scala:185). ########## sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala: ########## @@ -0,0 +1,560 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import scala.collection.mutable + +import org.apache.spark.SparkException +import org.apache.spark.sql.{DataFrame, QueryTest} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GetStructField, GreaterThan, If, LessThanOrEqual, Literal, NumericEvalContext, PredicateHelper, Rand, XxHash64} +import org.apache.spark.sql.catalyst.plans.logical.ReturnAnswer +import org.apache.spark.sql.execution.{BaseLimitExec, CoalesceExec, ExecSubqueryExpression, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning} +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation} +import org.apache.spark.sql.execution.joins.SortMergeJoinExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.connector.StorageFilterImpl +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{BinaryType, DoubleType, LongType, StructType} + +/** + * Covers the storage-filter contract for a data source v2 scan. With + * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through + * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The + * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked. + */ +class StorageFilterPushDownSuite + extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper with PredicateHelper { + + private def sourceName: String = classOf[StorageFilterTestSource].getName + + // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table. + private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = { + val bloomJoinConf = Seq( + SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") + withSQLConf(bloomJoinConf ++ conf: _*) { + withTable("bf_build") { + spark.range(0, 30).selectExpr("id AS bk", "id AS bv") + .write.format("parquet").saveAsTable("bf_build") + body + } + } + } + + // A join whose build side is small and carries a selective predicate, which is the shape + // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`. + private def bloomJoin(probe: DataFrame): DataFrame = { + val build = spark.table("bf_build").where("bv = 5") + probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*) + } + + private def v2Probe(options: (String, String)*): DataFrame = + spark.read.format(sourceName).options(options.toMap).load() + + private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME + + private def scanOf(plan: SparkPlan): BatchScanExec = + collect(plan) { case b: BatchScanExec => b }.headOption + .getOrElse(fail(s"no BatchScanExec in plan:\n$plan")) + + // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows + // the reader marks as checked. + private def bloomsAbove(plan: SparkPlan): (Int, Int) = { + val conditions = collect(plan) { case f: FilterExec => f.condition } + val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum + val skipped = conditions.map(_.collect { + case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain) + if a.name == checkedName => 1 + }.sum).sum + (all - skipped, skipped) + } + + // The scan offered nothing, and the Filter evaluates the bloom on every row. + private def assertNothingOffered(plan: SparkPlan): Unit = { + val scan = scanOf(plan) + assert(scan.storageFilters.isEmpty, s"the scan must get no storage filter:\n$plan") + assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan") + } + + // The rows the Filter on the scan returned, counted before the join. + private def rowsOutOfFilter(plan: SparkPlan): Long = { + val filters = collect(plan) { + case f: FilterExec if f.exists(_.isInstanceOf[BatchScanExec]) => f + } + assert(filters.size == 1, s"expected one Filter on the scan:\n$plan") + filters.head.metrics("numOutputRows").value + } + + // Runs `df`, checks that it returns `expected`, and returns its plan. + private def checkJoin(df: DataFrame, expected: Set[Seq[Long]]): SparkPlan = { + val result = df.collect().map(_.toSeq.map(_.asInstanceOf[Long])).toSet + assert(result == expected, s"got $result") + df.queryExecution.executedPlan + } + + // The build side keeps `bv = 5`, so the join returns that one row. + private val joinResult = Set(Seq(5L, 50L)) + + // Runs the bloom join over the v2 source read with `options`, checks its rows, and returns its + // plan. + private def runBloomJoin(options: (String, String)*): SparkPlan = { + StorageFilterTestSource.reset() + checkJoin(bloomJoin(v2Probe(options: _*)), joinResult) + } + + for (columnar <- Seq(false, true); aqe <- Seq(false, true)) { + val kind = if (columnar) "columnar" else "row-based" + test(s"a $kind scan gets the checked column, and the Filter skips the bloom on checked rows " + + s"(AQE $aqe)") { + withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + val plan = runBloomJoin("columnar" -> columnar.toString) + val scan = scanOf(plan) + assert(scan.supportsColumnar == columnar, s"the scan must read $kind:\n$plan") + assert(scan.storageFilters.size == 1, + s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan") + assert(scan.appendsCheckedColumn, s"the scan must get the checked column:\n$plan") + assert(bloomsAbove(plan) == (0, 1), + s"the Filter must skip the bloom on checked rows:\n$plan") + assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(), + s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " + + s"${StorageFilterTestSource.rowsReturned.get()}") + val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output } + assert(probeSide.exists(!_.exists(_.name == checkedName)), + s"the Project above the scan must drop the checked column; got $probeSide") + assert(StorageFilterTestSource.filteredReaders.get() == 1, + s"the one partition's reader must get the filter; " + + s"got ${StorageFilterTestSource.filteredReaders.get()}") + assert(StorageFilterTestSource.checkedColumnRequests.get() == 1, + "the reader must be asked for the checked column") + assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows, + s"the reader must evaluate every row; got ${StorageFilterTestSource.rowsEvaluated.get()}") + assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows, + "the reader must have dropped rows, otherwise the filter did nothing") + } + } + } + + Seq(false, true).foreach { columnar => + val kind = if (columnar) "columnar" else "row-based" + test(s"a $kind reader that gives the filter up marks its rows unchecked, and the Filter " + + "filters them") { + withBloomJoin() { + val plan = runBloomJoin("columnar" -> columnar.toString, "giveUp" -> "true") + assert(bloomsAbove(plan) == (0, 1), s"the plan is that of a filtering reader:\n$plan") + assert(StorageFilterTestSource.rowsEvaluated.get() == 0, + s"the reader must evaluate nothing; got ${StorageFilterTestSource.rowsEvaluated.get()}") + assert(StorageFilterTestSource.rowsReturned.get() == StorageFilterTestSource.numRows, + s"the reader must return every row; got ${StorageFilterTestSource.rowsReturned.get()}") + assert(rowsOutOfFilter(plan) < StorageFilterTestSource.numRows, + s"the Filter must apply the bloom to unchecked rows; got ${rowsOutOfFilter(plan)}") + } + } + + test(s"the Filter takes a $kind reader's word on the rows it marks checked") { + // This reader breaks the contract on purpose: it applies no filter but marks every row + // checked. Only skipping the bloom on those rows keeps all of them. + withBloomJoin() { + // The join must still filter, which `runBloomJoin` checks. + val plan = runBloomJoin("columnar" -> columnar.toString, "markAllChecked" -> "true") + assert(rowsOutOfFilter(plan) == StorageFilterTestSource.numRows, + s"the Filter must keep every row marked checked; got ${rowsOutOfFilter(plan)}") + } + } + + Seq(false, true).foreach { aqe => + test(s"a $kind reader that leaves out the checked column it was asked for fails the task, " + + s"from any row or batch (AQE $aqe)") { + // A reader that gives the filter up returns every row, so it returns several batches. + Seq(Seq("dropChecked" -> "true"), Seq("dropChecked" -> "later", "giveUp" -> "true")) + .foreach { options => + withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + StorageFilterTestSource.reset() + val df = bloomJoin(v2Probe(options :+ ("columnar" -> columnar.toString): _*)) + // `k`, `v` and the checked column. + checkError( + exception = intercept[SparkException](df.collect()), + condition = "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING", + sqlState = Some("KD000"), + parameters = Map( + "class" -> classOf[StorageFilterTestReaderFactory].getName, + "expected" -> "3", + "actual" -> "2")) + } + } + } + } + + test(s"a $kind reader may return more columns after the checked column") { + withBloomJoin() { + val plan = runBloomJoin("columnar" -> columnar.toString, "checkedLayout" -> "extra") + assert(bloomsAbove(plan) == (0, 1), + s"the Filter must skip the bloom on checked rows:\n$plan") + assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(), + s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " + + s"${StorageFilterTestSource.rowsReturned.get()}") + } + } + + test(s"a $kind test reader gives a filter up when it throws, and returns its rows unchecked") { + // `1 / k` under ANSI throws at k = 0, the partition's first row, so the reader gives the + // filter up before it keeps or drops anything. No runtime bloom throws, so the reader is + // driven by hand. + val k = AttributeReference("k", LongType)() + val inverse = Divide(Literal(1.0), Cast(k, DoubleType), NumericEvalContext(EvalMode.ANSI)) + val filter = StorageFilterImpl(GreaterThan(inverse, Literal(0.0)), Seq(k)) + val partition = RangeInputPartition(0, StorageFilterTestSource.numRows) + val schema = new StructType().add("k", LongType).add("v", LongType) + val rows = mutable.ArrayBuffer.empty[(Long, Boolean)] + if (columnar) { + val reader = new StorageFilterTestColumnarReader( + partition, schema, Some(filter), appendsChecked = true) + try { + while (reader.next()) { + val batch = reader.get() + (0 until batch.numRows()).foreach { i => + rows += ((batch.column(0).getLong(i), batch.column(2).getBoolean(i))) + } + } + } finally { + reader.close() + } + } else { + val reader = + new StorageFilterTestReader(partition, schema, Some(filter), appendsChecked = true) + while (reader.next()) { + rows += ((reader.get().getLong(0), reader.get().getBoolean(2))) + } + } + assert(rows == (0L until StorageFilterTestSource.numRows.toLong).map(_ -> false), + s"expected every row, unchecked; got ${rows.take(5)} and ${rows.size} rows in all") + } + } + + test("a columnar reader that returns the checked column ahead of the others fails the task") { + // The batch is as wide as asked for, but its last column is not a boolean. A row carries no + // types, so a row-based reader's mistake goes unnoticed. + withBloomJoin() { + StorageFilterTestSource.reset() + val df = bloomJoin(v2Probe("columnar" -> "true", "checkedLayout" -> "front")) + checkError( + exception = intercept[SparkException](df.collect()), + condition = "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING", + sqlState = Some("KD000"), + parameters = Map( + "class" -> classOf[StorageFilterTestReaderFactory].getName, + "expected" -> "3", + "actual" -> "3")) + } + } + + // Under a limit of two fields, the checked column would make three. + Seq( + ("too wide for whole-stage codegen", SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> "2"), + ("without whole-stage codegen", SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") + ).foreach { case (when, conf) => + test(s"a scan $when gets the bloom without the checked column") { + withBloomJoin(conf) { + val plan = runBloomJoin("columnar" -> "true") + val scan = scanOf(plan) + assert(scan.storageFilters.size == 1, + s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan") + assert(!scan.appendsCheckedColumn, s"the scan must not get the checked column:\n$plan") + assert(StorageFilterTestSource.checkedColumnRequests.get() == 0, + "the reader must not be asked for the checked column") + assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan") + assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows, + "the reader must have dropped rows, otherwise the filter did nothing") + } + } + } + + Seq( + ("with the conf off", Seq(SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "false"), Nil), + ("to a scan that declares no column", Nil, Seq("declared" -> "")), + ("to a scan that declares only another column", Nil, Seq("declared" -> "v")), + ("to a scan that declares only a nested field", Nil, Seq("declared" -> "k.x")) + ).foreach { case (when, conf, options) => + test(s"nothing is offered $when") { + withBloomJoin(conf: _*) { + val plan = runBloomJoin(options :+ ("columnar" -> "true"): _*) + assertNothingOffered(plan) + assert(StorageFilterTestSource.filteredReaders.get() == 0, + "no reader may get a storage filter") + } + } + } + + test("a scan that reads only the key and a metadata column is offered nothing") { + // Only the data columns count as columns left to skip, as on the file-source side. + withBloomJoin() { + val source = v2Probe("columnar" -> "true", "metadataColumn" -> "true") + val pos = StorageFilterTestSource.posName + StorageFilterTestSource.reset() + assertNothingOffered( + checkJoin(bloomJoin(source.select(source("k"), source.metadataColumn(pos))), + Set(Seq(5L, 5L)))) + StorageFilterTestSource.reset() + val withValue = checkJoin( + bloomJoin(source.select(source("k"), source("v"), source.metadataColumn(pos))), + Set(Seq(5L, 50L, 5L))) + assert(scanOf(withValue).storageFilters.size == 1, + s"with `v` left to skip, the bloom must be offered:\n$withValue") + } + } + + test("a variant scan that reads only its key column and a metadata column is offered nothing") { + // The key is an extraction of the declared variant column, so that column is the scan's only + // data column, and the scan's output has fresh exprIds. + withBloomJoin() { + val pos = StorageFilterTestSource.posName + val probe = v2Probe("columnar" -> "true", "variant" -> "true", "metadataColumn" -> "true", + "declared" -> "j").selectExpr("variant_get(j, '$.a', 'long') AS a", pos) + val build = spark.table("bf_build").where("bv = 5") + val df = probe.join(build, probe("a") === build("bk") * 10).select(probe("a"), probe(pos)) + StorageFilterTestSource.reset() + assertNothingOffered(checkJoin(df, Set(Seq(50L, 5L)))) + } + } + + test("a scan that took a pushed variant extraction gets the bloom") { + // The extraction keeps the scan's rows. + withBloomJoin() { + StorageFilterTestSource.reset() + val probe = v2Probe("columnar" -> "true", "variant" -> "true") + .selectExpr("k", "variant_get(j, '$.a', 'long') AS a") + val plan = checkJoin(bloomJoin(probe), joinResult) + val scan = scanOf(plan) + assert(scan.output.exists(_.dataType.isInstanceOf[StructType]), + s"test setup: the scan must read the extraction:\n$plan") + assert(scan.storageFilters.size == 1, + s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan") + assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows, + "the reader must have dropped rows, otherwise the filter did nothing") + } + } + + for (aqe <- Seq(false, true)) { + test("a scan with a column named like the checked column gets the filter, and runtime " + + s"filtering still resolves that column (AQE $aqe)") { + // Runtime filtering resolves the declared filter attributes by name, which the checked column + // the scan appends would make ambiguous. Dynamic partition pruning on `k` would take the + // bloom's place. + withBloomJoin(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + val plan = runBloomJoin( + "columnar" -> "true", "valueName" -> checkedName, "runtimeFiltering" -> "v2") + val scan = scanOf(plan) + assert(scan.storageFilters.size == 1 && scan.appendsCheckedColumn, + s"the scan must get the bloom and the checked column:\n$plan") + assert(StorageFilterTestSource.runtimeFilterCalls.get() > 0, + "test setup: the bloom must reach runtime filtering") + val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output } + assert(probeSide.contains(scan.output.init), + s"the Project above the scan must keep the table's column; got $probeSide") + } + } + + test(s"a bloom the scan fully evaluates through runtime filtering is not offered (AQE $aqe)") { + // Dynamic partition pruning on `k` would take the bloom's place. + withBloomJoin(SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + val plan = runBloomJoin("columnar" -> "true", "runtimeFiltering" -> "catalyst") + val scan = scanOf(plan) + assert(scan.runtimeFilters.exists(_.exists(_.isInstanceOf[BloomFilterMightContain])) && + StorageFilterTestSource.runtimeFilterCalls.get() > 0, + s"test setup: the scan must take the bloom as a runtime filter:\n$plan") + assert(scan.storageFilters.isEmpty, s"the scan must get no storage filter:\n$plan") + assert(bloomsAbove(plan) == (0, 0), s"no Filter may keep the bloom:\n$plan") + } + } + + test(s"a two-key join offers both blooms in one filter, each bloom with its own bytes " + + s"(AQE $aqe)") { + // `MergeSubplans` merges the two blooms' aggregates into one subquery returning both blooms. + withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + withTable("bf_build2") { + spark.range(0, 30).selectExpr("id AS bk", "id * 10 AS bv", "id AS flag") + .write.format("parquet").saveAsTable("bf_build2") + StorageFilterTestSource.reset() + val probe = v2Probe("columnar" -> "true", "wide" -> "true", "declared" -> "k,v") + .where("w % 3 != 1") + val build = spark.table("bf_build2").where("flag = 5") + val df = probe.join(build, probe("k") === build("bk") && probe("v") === build("bv")) + .select(probe("k"), probe("v"), probe("w")) + val plan = checkJoin(df, Set(Seq(5L, 50L, 500L))) + val scan = scanOf(plan) + assert(scan.storageFilters.size == 2, s"both blooms must be on the scan:\n$plan") + assert(scan.storageFilters.forall(_.exists { + case GetStructField(_: ExecSubqueryExpression, _, _) => true + case _ => false + }), s"test setup: the blooms must come from one merged subquery:\n$plan") + assert(bloomsAbove(plan) == (0, 2), s"the Filter must skip both blooms:\n$plan") + val filter = collect(plan) { + case f: FilterExec if f.child.exists(_.isInstanceOf[BatchScanExec]) => f + }.head + assert(splitConjunctivePredicates(filter.condition).exists { c => + c.references.map(_.name).toSet == Set("w") && !c.isInstanceOf[If] + }, s"the Filter must evaluate the conjunct over `w` on every row:\n$plan") + assert(StorageFilterTestSource.filterReferences == Set(Seq("k", "v")), + s"one filter over both keys; got ${StorageFilterTestSource.filterReferences}") + val bloomBytes = scan.storageFilters.map(_.collectFirst { + case m: BloomFilterMightContain => + m.bloomFilterExpression.eval().asInstanceOf[Array[Byte]].length + }.get).sum + assert(StorageFilterTestSource.filterBytes.get() < bloomBytes * 3 / 2, + s"each bloom must ship once; the filter takes ${StorageFilterTestSource.filterBytes} " + + s"bytes for $bloomBytes bytes of blooms") + assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows, + "the reader must evaluate every row; got " + + StorageFilterTestSource.rowsEvaluated.get()) + } + } + } + + test(s"the readers of one task share the filter the task deserialized (AQE $aqe)") { + // `coalesce(1)` reads both partitions in one task. + withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + StorageFilterTestSource.reset() + val probe = v2Probe("columnar" -> "true", "numPartitions" -> "2").coalesce(1) + val plan = checkJoin(bloomJoin(probe), joinResult) + assert(collect(plan) { case c: CoalesceExec => c }.nonEmpty, s"test setup:\n$plan") + assert(scanOf(plan).storageFilters.size == 1, s"the bloom must be on the scan:\n$plan") + assert(StorageFilterTestSource.filteredReaders.get() == 2, + s"both readers must get the filter; got ${StorageFilterTestSource.filteredReaders.get()}") + assert(StorageFilterTestSource.filterIdentities.size == 1, + "both readers must get the same instance") + assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows, + "the readers must evaluate every row; got " + StorageFilterTestSource.rowsEvaluated.get()) + assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows, + "the readers must have dropped rows, otherwise the filter did nothing") + } + } + } + + test("the storage filter shows in the scan description and separates canonicalized plans") { + withBloomJoin() { + val plan = bloomJoin(v2Probe("columnar" -> "true")).queryExecution.executedPlan + val scan = scanOf(plan) + assert(scan.simpleString(100).contains("StorageFilters:"), + s"the scan description must name the filters: ${scan.simpleString(100)}") + assert(scan.verboseStringWithOperatorId().contains("StorageFilters:"), + s"EXPLAIN FORMATTED must name the filters: ${scan.verboseStringWithOperatorId()}") + + intercept[IllegalArgumentException](scan.copy(storageFilters = Nil)) + // As the scan would be without the checked column, too wide or without codegen. + val unchecked = scan.copy(output = scan.output.init, appendsCheckedColumn = false) + assert(scan != unchecked, "equals must see the checked column") + val plain = unchecked.copy(storageFilters = Nil) + assert(unchecked != plain && unchecked.hashCode != plain.hashCode, + "equals and hashCode must see storageFilters") + assert(!unchecked.sameResult(plain), + s"a filtering scan must not be reusable as a plain one:\n${unchecked.canonicalized}\n" + + s"${plain.canonicalized}") + assert(!plain.simpleString(100).contains("StorageFilters:"), + s"a plain scan's description must not name storage filters: ${plain.simpleString(100)}") + assert(!plain.verboseStringWithOperatorId().contains("StorageFilters:"), + s"nor its EXPLAIN FORMATTED: ${plain.verboseStringWithOperatorId()}") + + // As another relation instance of the same table gets, with fresh exprIds. + val fresh = AttributeMap(scan.output.map(a => a -> a.newInstance())) + val twin = scan.copy( + output = scan.output.map(fresh(_)), + storageFilters = scan.storageFilters.map(_.transform { + case a: Attribute => fresh.getOrElse(a, a) + })) + assert(scan.sameResult(twin), + s"the same filtering scan must be reusable:\n${scan.canonicalized}\n${twin.canonicalized}") + } + } + + // The reader would filter ahead of the operation, which then picks other rows. + Seq[(String, DataFrame => DataFrame)]( + "limit" -> (_.limit(100)), + "offset" -> (_.offset(10)), + "topN" -> (_.orderBy("v").limit(100)), + "sample" -> (_.sample(0.5)) + ).foreach { case (op, apply) => + test(s"nothing is offered to a scan that took a pushed $op") { + withBloomJoin() { + val probe = apply(v2Probe("columnar" -> "true", "push" -> op)) + val plan = bloomJoin(probe).queryExecution.executedPlan + assert(collect(plan) { + case p @ (_: BaseLimitExec | _: TakeOrderedAndProjectExec | _: SampleExec) => p + }.isEmpty, s"the $op must be pushed into the scan whole:\n$plan") + assertNothingOffered(plan) + } + } + } + + test("the scan's pushdowns decide on storage filters, not whether it can be merged") { + // A pushed variant extraction, for one, keeps the scan's rows but cannot be merged. + withBloomJoin() { + val optimized = bloomJoin(v2Probe("columnar" -> "true")).queryExecution.optimizedPlan + def planned(mergeable: Boolean, rowWise: Boolean): BatchScanExec = scanOf( + spark.sessionState.planner.plan(ReturnAnswer(optimized.transform { Review Comment: **10. [Nit] `QueryExecution.createSparkPlan` already does this** `spark.sessionState.planner.plan(ReturnAnswer(...)).next()` is the body of `QueryExecution.createSparkPlan` (QueryExecution.scala:923-929), which `DynamicPartitionPruningSuite` calls the same way. `QueryExecution.createSparkPlan(spark.sessionState.planner, optimized.transform { ... })` would also let the `ReturnAnswer` import (L26) go. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala: ########## @@ -0,0 +1,103 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.internal.connector + +import org.apache.spark.SparkException +import org.apache.spark.internal.{Logging, LogKeys} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression => CatalystExpression} +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference} +import org.apache.spark.sql.connector.read.{StorageFilter, StorageFilterNotAppliedException} + +/** + * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not + * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key + * rather than a data row. + * + * `attrs` is the layout of the row [[eval]] expects, in that order, and the predicate is bound to + * it. The predicate is built on first use, see [[StorageFilterPredicate]]. + */ +class StorageFilterImpl private ( + private val predicate: StorageFilterPredicate, + private val attrs: Seq[Attribute]) + extends StorageFilter with Logging { + + // The width of the row `eval` takes, read on every row. + private val numReferences = attrs.length + + // Whether this instance has logged a failure, so a task logs only its first. + @transient private var loggedFailure = false + + /** + * Wraps an error in [[StorageFilterNotAppliedException]], unless + * [[StorageFilterErrors.rethrowIfMustPropagate]] rethrows it, as the file-source reader does: a + * task kill, or a fatal error in its cause chain. Each task logs the first one. A row of another + * width than [[references]] is a reader bug, which fails the task, since the predicate would read + * the wrong values. + */ + @throws[StorageFilterNotAppliedException] + override def eval(row: InternalRow): Boolean = { + // Outside the `try`, so a failure to build the predicate, or a row of another width, fails the + // task. + if (row.numFields != numReferences) { + throw SparkException.internalError( + s"$this takes rows of $numReferences values, but got ${row.numFields}") + } + val prepared = predicate.prepared Review Comment: **8. [Nit] `eval` resolves the predicate's lazy val on every row** `predicate.prepared` is a `@transient lazy val` (StorageFilterPredicate.scala:30), so every call reads its volatile initialization flag, also after the predicate is built. The V1 reader moves the same lazy val out of its loop on purpose, "since resolving it is a volatile read" (LateMaterializationParquetRecordReader.java:731-732), and since item 12 of my second review both resolve the same `StorageFilterPredicate`. A V2 reader cannot do that, since it only calls `eval`, and the default `select` calls `eval` for every row. The cost is small next to a bloom probe, but a plain field keeps V2 on V1's rule, as one thread uses the filter (StorageFilter.java:48-49): ```scala @transient private var built: BasePredicate = _ ... if (built == null) built = predicate.prepared try { built.eval(row) ``` ########## common/utils/src/main/resources/error/error-conditions.json: ########## @@ -2397,6 +2397,18 @@ ], "sqlState" : "0A000" }, + "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING" : { + "message" : [ + "A reader of <class> was asked for the column that marks the rows a storage filter checked, at position <expected>. That column is missing or not a boolean. The reader returned <actual> columns." Review Comment: **3. [Low] The new message calls the width a position, and its parameters keep their count names** Since `5141077` the message says the column was asked for `at position <expected>`, but `<expected>` is still the width `BatchScanExec` passes (`output.length`, BatchScanExec.scala:142). So it is a 1-based position, and the message does not say so. A connector author counts from 0, as `ColumnarBatch.column(int)` and `InternalRow` ordinals do, and this PR's tests read the checked column of `k, v, checked` as `batch.column(2)` (StorageFilterPushDownSuite.scala:236). For the `front` layout the task fails with "... at position 3. That column is missing or not a boolean. The reader returned 3 columns." Read 0-based, that says a fourth column is missing, while the problem is the BIGINT at `column(2)`. The parameters also kept their names from when they were counts, so `getMessageParameters()`, the JSON error formats and the new tests report `expected=3, actual=3` for a failure. The condition is new, so renaming it is free now and a contract once it is released. For example: ``` A reader of <class> was asked for the column that marks the rows a storage filter checked, right after the <numScanColumns> columns of the scan's read schema. That column is missing or not a boolean. The reader returned <numColumns> columns. ``` with `width - 1` passed as `numScanColumns` (DataSourceRDD.scala:238-240), and the same names in `storageFilterCheckedColumnMissingError` and in both `checkError` calls. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [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]
