peter-toth commented on code in PR #59715: URL: https://github.com/apache/spark/pull/59715#discussion_r4386746505
########## 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: Thanks, changed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb) to your version. The field is read straight from the subquery's result, so only that field is copied. ########## 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: Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). I went back to the exact width check, and kept the boolean check for a batch. So your example fails the task again, and so does a reader that returns more columns after the checked one. The doc and the description now say what Spark checks. It fails the task for a row or a batch of another width, or a batch whose last column is not a boolean. ########## 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: Thanks, changed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb) to your message and parameter names. It also says the column may not be the last one, since the check is exact again, see item 2. ########## 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: Thanks, added in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). The test now also checks that the reader excluded rows. It fails when the reader gets no filter. ########## 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: Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb). Both tests use `checkError` with the SQLSTATE. The one for `DATA_SOURCE_STORAGE_FILTER_READER_NOT_SUPPORTED` checks the factory's class too. ########## 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: Thanks, fixed in [`cb22f93`](https://github.com/apache/spark/pull/59715/commits/cb22f93ce6282c5a457942838d84819e55ad4fdb), with your wording. -- 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]
