peter-toth commented on code in PR #59715: URL: https://github.com/apache/spark/pull/59715#discussion_r4381901611
########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/StorageFilter.java: ########## @@ -0,0 +1,105 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector.expressions.filter; + +import java.io.Serializable; + +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.connector.expressions.NamedReference; +import org.apache.spark.sql.connector.read.Scan; +import org.apache.spark.sql.connector.read.StorageFilterRequest; +import org.apache.spark.sql.connector.read.SupportsStorageFilters; +import org.apache.spark.sql.vectorized.ColumnarBatch; + +/** + * A row-level predicate a scan may evaluate inside its readers, to skip reading the other columns + * of the rows it rejects. The reader reads the columns of {@link #references()} first, evaluates + * the filter, and reads the remaining columns only for the rows it keeps. + * <p> + * The filter is a predicate on the rows the scan returns. Spark offers none to a scan that took a + * pushed operation that decides which rows it returns, such as a limit, an offset, a top-N, a + * sample, an aggregate or a join. + * <p> + * Spark implements the predicate, and the reader only supplies the values. For example, a scan + * does not have to reproduce the Spark sketch and hash of a runtime bloom filter. The values must + * be the ones the scan returns for those columns, in the types of {@link Scan#readSchema()}, nulls + * included. + * <p> + * Applying it is optional, see {@link SupportsStorageFilters}. + * <p> + * A row's result depends only on its own values. So with one reference, a reader may evaluate each + * distinct value once, such as each entry of a dictionary, and reuse the result for every row + * holding that value. A dictionary entry is evaluated as the value the scan returns for it, decoded + * and converted. + * <p> + * Spark hands the filter to each task in a {@link StorageFilterRequest}. The readers of a task + * share one instance, which is not thread-safe. A reader must not copy or serialize it, since an + * evaluated instance may have released what a copy would be built from. + * + * @since 5.0.0 + */ +@Evolving +public abstract class StorageFilter implements Serializable { Review Comment: Thanks, moved in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538), with `StorageFilterNotAppliedException`. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala: ########## @@ -0,0 +1,116 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.internal.connector + +import org.apache.spark.SparkException +import org.apache.spark.internal.{Logging, LogKeys} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate} +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference} +import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException} + +/** + * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not + * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key + * rather than a data row. + * + * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to + * it. The binding is built on first use, so an instance created on the driver can be evaluated on + * an executor. + */ +class StorageFilterImpl private ( + private var catalystExpr: CatalystExpression, + private val attrs: Seq[Attribute]) + extends StorageFilter with Logging { + + // The width of the row `eval` takes, read on every row. + private val numReferences = attrs.length + + // Whether this instance has logged a failure, so a task logs only its first. + @transient private var loggedFailure = false + + @transient private lazy val boundPredicate: InternalRow => Boolean = { + // The expression is released below, so a copy made after a first evaluation would have nothing + // to build from. That is named here rather than left to an NPE. + if (catalystExpr == null) { + throw SparkException.internalError( + s"$this was copied after an evaluation released the expression of its predicate") + } + val predicate = CatalystPredicate.create(BindReferences.bindReference(catalystExpr, attrs)) + // A bloom's bytes sit in the expression as a binary literal, while a generated predicate holds + // the deserialized filter instead, so dropping the expression lets the task reclaim them. An + // interpreted predicate, the fallback, keeps the expression. Every task deserializes its own + // copy, and this is the field's only reader. + catalystExpr = null + predicate.eval + } + + /** + * Wraps an error in [[StorageFilterNotAppliedException]], unless + * [[StorageFilterErrors.rethrowIfMustPropagate]] rethrows it, as the file-source reader does: a + * task kill, or a fatal error in its cause chain. Each task logs the first one. A row of another + * width than [[references]] is a reader bug, which fails the task, since the predicate would read + * the wrong values. + */ + @throws[StorageFilterNotAppliedException] + override def eval(row: InternalRow): Boolean = { + // Outside the `try`, so a failure to build the predicate, or a row of another width, fails the + // task. + if (row.numFields != numReferences) { + throw SparkException.internalError( + s"$this takes rows of $numReferences values, but got ${row.numFields}") + } + val predicate = boundPredicate + try { + predicate(row) + } catch { + case e: Throwable => + StorageFilterErrors.rethrowIfMustPropagate(e, fromRead = false) + if (!loggedFailure) { + loggedFailure = true + logWarning(log"${MDC(LogKeys.EXPR, this)} could not be applied, so the reader reads " + + log"without it.", e) + } + throw new StorageFilterNotAppliedException(e) + } + } + + @transient override lazy val references: Array[NamedReference] = Review Comment: Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). `references()` returns a copy. A test reorders one and checks that the next call is unchanged. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterPushDownSuite.scala: ########## @@ -0,0 +1,491 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import scala.collection.mutable + +import org.apache.spark.{SparkException, SparkThrowable} +import org.apache.spark.sql.{DataFrame, QueryTest} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, BloomFilterMightContain, Cast, Divide, EvalMode, GetStructField, GreaterThan, If, Literal, NumericEvalContext, PredicateHelper, Rand, XxHash64} +import org.apache.spark.sql.catalyst.plans.logical.ReturnAnswer +import org.apache.spark.sql.execution.{BaseLimitExec, CoalesceExec, ExecSubqueryExpression, FilterExec, SampleExec, SparkPlan, TakeOrderedAndProjectExec} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.datasources.{FileFormat, StorageFilterPlanning} +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation} +import org.apache.spark.sql.execution.joins.SortMergeJoinExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.connector.StorageFilterImpl +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{BinaryType, DoubleType, LongType, StructType} + +/** + * Covers the storage-filter contract for a data source v2 scan. With + * `V2_STORAGE_FILTER_PUSHDOWN_ENABLED` on, a scan that declares columns through + * `SupportsStorageFilters` gets the runtime bloom filter over those columns in every reader. The + * bloom stays in the Filter above the scan, skipped on the rows the reader marks as checked. + */ +class StorageFilterPushDownSuite + extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper with PredicateHelper { + + private def sourceName: String = classOf[StorageFilterTestSource].getName + + // Runs `body` with runtime blooms and V2 storage filters on, `conf` on top, and the build table. + private def withBloomJoin(conf: (String, String)*)(body: => Unit): Unit = { + val bloomJoinConf = Seq( + SQLConf.V2_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") + withSQLConf(bloomJoinConf ++ conf: _*) { + withTable("bf_build") { + spark.range(0, 30).selectExpr("id AS bk", "id AS bv") + .write.format("parquet").saveAsTable("bf_build") + body + } + } + } + + // A join whose build side is small and carries a selective predicate, which is the shape + // InjectRuntimeFilter builds a bloom filter for. The bloom lands on `probe`'s `k`. + private def bloomJoin(probe: DataFrame): DataFrame = { + val build = spark.table("bf_build").where("bv = 5") + probe.join(build, probe("k") === build("bk")).select(probe.columns.map(probe(_)).toSeq: _*) + } + + private def v2Probe(options: (String, String)*): DataFrame = + spark.read.format(sourceName).options(options.toMap).load() + + private val checkedName = FileFormat.STORAGE_FILTER_CHECKED_COLUMN_NAME + + private def scanOf(plan: SparkPlan): BatchScanExec = + collect(plan) { case b: BatchScanExec => b }.headOption + .getOrElse(fail(s"no BatchScanExec in plan:\n$plan")) + + // The blooms the Filters above the scan evaluate on every row, and those they skip on the rows + // the reader marks as checked. + private def bloomsAbove(plan: SparkPlan): (Int, Int) = { + val conditions = collect(plan) { case f: FilterExec => f.condition } + val all = conditions.map(_.collect { case _: BloomFilterMightContain => 1 }.sum).sum + val skipped = conditions.map(_.collect { + case If(a: Attribute, Literal.TrueLiteral, _: BloomFilterMightContain) + if a.name == checkedName => 1 + }.sum).sum + (all - skipped, skipped) + } + + // The scan offered nothing, and the Filter evaluates the bloom on every row. + private def assertNothingOffered(plan: SparkPlan): Unit = { + val scan = scanOf(plan) + assert(scan.storageFilters.isEmpty && !scan.appendsCheckedColumn, + s"the scan must get no storage filter:\n$plan") + assert(bloomsAbove(plan) == (1, 0), s"the Filter must keep the bloom as it is:\n$plan") + } + + // The rows the Filter on the scan returned, counted before the join. + private def rowsOutOfFilter(plan: SparkPlan): Long = { + val filters = collect(plan) { + case f: FilterExec if f.exists(_.isInstanceOf[BatchScanExec]) => f + } + assert(filters.size == 1, s"expected one Filter on the scan:\n$plan") + filters.head.metrics("numOutputRows").value + } + + // Runs `df`, checks that it returns `expected`, and returns its plan. + private def checkJoin(df: DataFrame, expected: Set[Seq[Long]]): SparkPlan = { + val result = df.collect().map(_.toSeq.map(_.asInstanceOf[Long])).toSet + assert(result == expected, s"got $result") + df.queryExecution.executedPlan + } + + // The build side keeps `bv = 5`, so the join returns that one row. + private val joinResult = Set(Seq(5L, 50L)) + + // Runs the bloom join over the v2 source read with `options`, checks its rows, and returns its + // plan. + private def runBloomJoin(options: (String, String)*): SparkPlan = { + StorageFilterTestSource.reset() + checkJoin(bloomJoin(v2Probe(options: _*)), joinResult) + } + + for (columnar <- Seq(false, true); aqe <- Seq(false, true)) { + val kind = if (columnar) "columnar" else "row-based" + test(s"a $kind scan gets the checked column, and the Filter skips the bloom on checked rows " + + s"(AQE $aqe)") { + withBloomJoin(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + val plan = runBloomJoin("columnar" -> columnar.toString) + val scan = scanOf(plan) + assert(scan.supportsColumnar == columnar, s"the scan must read $kind:\n$plan") + assert(scan.storageFilters.size == 1, + s"the bloom filter must be on the scan; got ${scan.storageFilters}\n$plan") + assert(scan.appendsCheckedColumn, s"the scan must get the checked column:\n$plan") + assert(bloomsAbove(plan) == (0, 1), + s"the Filter must skip the bloom on checked rows:\n$plan") + assert(rowsOutOfFilter(plan) == StorageFilterTestSource.rowsReturned.get(), + s"the Filter must keep every checked row; got ${rowsOutOfFilter(plan)} of " + + s"${StorageFilterTestSource.rowsReturned.get()}") + val probeSide = collectFirst(plan) { case j: SortMergeJoinExec => j.left.output } + assert(probeSide.exists(!_.exists(_.name == checkedName)), + s"the Project above the scan must drop the checked column; got $probeSide") + assert(StorageFilterTestSource.filteredReaders.get() == 1, + s"the one partition's reader must get the filter; " + + s"got ${StorageFilterTestSource.filteredReaders.get()}") + assert(StorageFilterTestSource.checkedColumnRequests.get() == 1, + "the reader must be asked for the checked column") + assert(StorageFilterTestSource.rowsEvaluated.get() == StorageFilterTestSource.numRows, + s"the reader must evaluate every row; got ${StorageFilterTestSource.rowsEvaluated.get()}") + assert(StorageFilterTestSource.rowsReturned.get() < StorageFilterTestSource.numRows, + "the reader must have dropped rows, otherwise the filter did nothing") + } + } + } + + Seq(false, true).foreach { columnar => + val kind = if (columnar) "columnar" else "row-based" + test(s"a $kind reader that gives the filter up marks its rows unchecked, and the Filter " + + "filters them") { + withBloomJoin() { + val plan = runBloomJoin("columnar" -> columnar.toString, "giveUp" -> "true") + assert(bloomsAbove(plan) == (0, 1), s"the plan is that of a filtering reader:\n$plan") + assert(StorageFilterTestSource.rowsEvaluated.get() == 0, + s"the reader must evaluate nothing; got ${StorageFilterTestSource.rowsEvaluated.get()}") + assert(StorageFilterTestSource.rowsReturned.get() == StorageFilterTestSource.numRows, + s"the reader must return every row; got ${StorageFilterTestSource.rowsReturned.get()}") + assert(rowsOutOfFilter(plan) < StorageFilterTestSource.numRows, + s"the Filter must apply the bloom to unchecked rows; got ${rowsOutOfFilter(plan)}") + } + } + + test(s"the Filter takes a $kind reader's word on the rows it marks checked") { + // This reader breaks the contract on purpose: it applies no filter but marks every row + // checked. Only skipping the bloom on those rows keeps all of them. + withBloomJoin() { + // The join must still filter, which `runBloomJoin` checks. + val plan = runBloomJoin("columnar" -> columnar.toString, "markAllChecked" -> "true") + assert(rowsOutOfFilter(plan) == StorageFilterTestSource.numRows, + s"the Filter must keep every row marked checked; got ${rowsOutOfFilter(plan)}") + } + } + + test(s"a $kind reader that leaves out the checked column it was asked for fails the task, " + + "from any row or batch") { + // A reader that gives the filter up returns every row, so it returns several batches. + Seq(Seq("dropChecked" -> "true"), Seq("dropChecked" -> "later", "giveUp" -> "true")) + .foreach { options => + withBloomJoin() { + StorageFilterTestSource.reset() + val df = bloomJoin(v2Probe(options :+ ("columnar" -> columnar.toString): _*)) + val e = intercept[SparkException](df.collect()) + val missing = + Iterator.iterate[Throwable](e)(_.getCause).takeWhile(_ != null).collectFirst { + case t: SparkThrowable + if t.getCondition == "DATA_SOURCE_STORAGE_FILTER_CHECKED_COLUMN_MISSING" => t + } + assert(missing.exists(_.getMessageParameters.get("class") == Review Comment: Thanks, fixed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). The test uses `checkError` with the condition, `KD000` and all three parameters. So does the new test for item 1 whose reader returns the column ahead of the others. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala: ########## @@ -957,7 +972,8 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { // changes with variant extraction. When validConstraints is wired up, this needs revisiting. // Pushed variant extraction cannot be reproduced by rebuilding the scan, so this scan stays // not-mergeable (the default `mergeableScan = false`). - val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput) + val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput, + onlyRowWisePushdowns = !hasRowSelectingPushdown(holder)) Review Comment: Thanks, added in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). The test source now implements `SupportsPushDownVariantExtractions`, and a new test runs a query of your example's shape over it. Dropping the `onlyRowWisePushdowns` argument at the variant build site fails it. The name match matters only for a scan that reads just its key columns and metadata columns. A variant scan gets there when the declared column is the variant column and the join key is an extraction of it. A second test reads only such an extraction and a metadata column, and checks that nothing is offered. Going back to the exprId match fails it. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/StorageFilterImpl.scala: ########## @@ -0,0 +1,116 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.internal.connector + +import org.apache.spark.SparkException +import org.apache.spark.internal.{Logging, LogKeys} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, BindReferences, Expression => CatalystExpression, Predicate => CatalystPredicate} +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference} +import org.apache.spark.sql.connector.expressions.filter.{StorageFilter, StorageFilterNotAppliedException} + +/** + * A [[StorageFilter]] wrapping a Catalyst expression, so a scan can apply a predicate it does not + * have to understand. It mirrors [[PartitionPredicateImpl]], which evaluates a partition key + * rather than a data row. + * + * `attrs` is the layout of the row [[eval]] expects, in that order, and the expression is bound to + * it. The binding is built on first use, so an instance created on the driver can be evaluated on + * an executor. + */ +class StorageFilterImpl private ( + private var catalystExpr: CatalystExpression, + private val attrs: Seq[Attribute]) + extends StorageFilter with Logging { + + // The width of the row `eval` takes, read on every row. + private val numReferences = attrs.length + + // Whether this instance has logged a failure, so a task logs only its first. + @transient private var loggedFailure = false + + @transient private lazy val boundPredicate: InternalRow => Boolean = { Review Comment: Thanks, changed in [`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538). The new `StorageFilterPredicate`, next to `StorageFilterErrors`, holds the bound expression. It builds the predicate on first use and then releases the expression. Its `apply` binds the expression, so `StorageFilterImpl` binds on the driver now, and `ParquetStorageFilter` passes its expression already bound. The bind and the build happen at different times now, so the two-argument `create` no longer fits. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [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]
