peter-toth commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4381899032
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -204,6 +221,13 @@ class DataSourceRDD(
}
}
+ private def checkWidth(expected: Int, actual: Int): Unit = {
+ if (actual != expected) {
Review Comment:
Thanks, fixed in
[`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538).
A batch now needs at least the expected number of columns, with a boolean at
the checked column's position. A row needs at least that many fields. So a
reader may return more columns after the checked one. The message says that the
column is missing or not a boolean. This stays a best-effort guard for a reader
that breaks the documented layout, which I'd expect to be rare. A row carries
no names or types, and a misplaced boolean column still passes. New tests cover
a columnar reader that returns the column ahead of the others, and readers that
return one more column after it.
##########
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")
Review Comment:
Thanks, fixed in
[`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538).
The two-key, two-reader and runtime-filtering tests, and the test of a reader
that leaves out the checked column, now run with AQE on and off.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala:
##########
@@ -817,8 +817,8 @@ case class FileSourceScanExec(
protected lazy val preparedStorageFilters: Seq[Expression] = {
val requestedDataAttrs = output.take(requiredSchema.length)
storageFilters.map { expr =>
- val subqueryReplaced = expr.transform { case s: execution.ScalarSubquery
=> s.toLiteral }
- BindReferences.bindReference(subqueryReplaced, requestedDataAttrs)
+ BindReferences.bindReference(
+ StorageFilterPlanning.withSubqueryResults(expr), requestedDataAttrs)
Review Comment:
Thanks, added in
[`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538).
`ParquetStorageFilterSuite` has a two-key join whose blooms come from one
merged subquery, with AQE on and off. It checks that the scan's RDD, which
carries the reader's function, takes less than 1.5 times the bytes of the two
blooms. With `s.toLiteral` alone it takes about twice the bytes of the two
blooms, and the test fails.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/StorageFilterPlanning.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.execution.datasources
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute,
AttributeReference, AttributeSet, BloomFilterMightContain, Expression,
ExpressionSet, If, Literal}
+import org.apache.spark.sql.execution.{ScalarSubquery => ExecScalarSubquery,
WholeStageCodegenExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BinaryType, 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. Each bloom filter
+ * becomes a literal of its own bytes. A subquery `MergeSubplans` merged
returns the blooms of
+ * several filters in one struct, so a bloom read out of it would otherwise
carry the whole
+ * struct.
+ */
+ def withSubqueryResults(filter: Expression): Expression = filter.transformUp
{
+ case s: ExecScalarSubquery => s.toLiteral
+ case m @ BloomFilterMightContain(bloom, _) if bloom.foldable &&
!bloom.isInstanceOf[Literal] =>
Review Comment:
Thanks, changed in
[`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538)
to your version. The unit test covers a long field too. One correction on the
link: SPARK-60448's key-column conditions hold no subquery. The ones that do
come with SPARK-60500, so the description now names both.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -254,6 +279,7 @@ private class PartitionIterator[T](
}
valuePrepared = false
val result = reader.get()
+ check.foreach(_(result))
Review Comment:
Thanks, fixed in
[`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538).
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/Scan.java:
##########
@@ -70,9 +70,10 @@ default String description() {
* {@link Table} that creates this scan returns {@link
TableCapability#BATCH_READ} support in its
* {@link Table#capabilities()}.
* <p>
- * If the scan supports runtime filtering and implements {@link
SupportsRuntimeFiltering},
- * this method may be called multiple times. Therefore, implementations can
cache some state
- * to avoid planning the job twice.
+ * If the scan supports runtime filtering and implements {@link
SupportsRuntimeFiltering}, or
+ * Spark offers it a storage filter through {@link SupportsStorageFilters},
this method may be
Review Comment:
Thanks, fixed in
[`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538),
with your wording.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/StorageFilterTestSource.scala:
##########
@@ -0,0 +1,444 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector
+
+import java.io.{ByteArrayOutputStream, ObjectOutputStream}
+import java.util
+import java.util.OptionalLong
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Expression,
GenericInternalRow}
+import org.apache.spark.sql.connector.catalog.{MetadataColumn,
SupportsMetadataColumns, SupportsRead, Table, TableCapability, TableProvider}
+import org.apache.spark.sql.connector.expressions.{Expressions,
FieldReference, NamedReference, SortOrder, Transform}
+import org.apache.spark.sql.connector.expressions.filter.{Predicate,
StorageFilter, StorageFilterNotAppliedException}
+import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.sql.types.{BooleanType, DataType, LongType, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+/**
+ * A read-only v2 source over a fixed table of longs that implements
[[SupportsStorageFilters]]:
+ * the key `k`, the value `v` holding `k * 10`, and with `wide` also `w`
holding `k * 100`. Its
+ * reader gets the filter in the task and drops the rows the filter rejects.
It reads the columns
+ * Spark requires, and builds the filter's row from the filter's
`references()`. Its options:
+ * - `columnar`: reads in batches, through `StorageFilter.select`, rather
than row by row;
+ * - `giveUp`: reads without the filter, as a reader does for a unit it
cannot use the filter on.
+ * Here the unit is the partition;
+ * - `markAllChecked`: reads without the filter but marks every row checked,
which breaks the
+ * contract on purpose;
+ * - `dropChecked`: `true` never appends the checked column, and `later`
appends it to the first
+ * row or batch only. Either breaks the contract on purpose;
+ * - `declared`: the comma-separated storage-filter columns, `k` by default,
a dot separating a
+ * nested field;
+ * - `valueName`: the name of the value column, `v` by default;
+ * - `wide`: adds the column `w`;
+ * - `numPartitions`: splits the keys into that many input partitions, 1 by
default;
+ * - `metadataColumn`: adds the metadata column `_pos`, which holds the key;
+ * - `runtimeFiltering`: `v2` makes the scan a `SupportsRuntimeV2Filtering`
one over every column,
+ * with iterative pushdown on a table partitioned by `k`. `catalyst` makes
it a
+ * `SupportsRuntimeCatalystFiltering` one that evaluates the filters on `k`
fully. Neither
+ * applies the filters it takes, so they serve planning tests only;
+ * - `push`: the comma-separated pushed operations it takes whole without
applying them, of
+ * `limit`, `offset`, `topN` and `sample`, for planning tests only.
+ *
+ * A filter that throws `StorageFilterNotAppliedException` is given up for the
rest of the
+ * partition too. The counters below record what the readers got, evaluated
and returned.
+ */
+class StorageFilterTestSource extends TableProvider {
+
+ override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+ val schema = new StructType().add("k",
LongType).add(options.getOrDefault("valueName", "v"),
+ LongType)
+ if (options.getBoolean("wide", false)) schema.add("w", LongType) else
schema
+ }
+
+ override def getTable(
+ schema: StructType,
+ partitioning: Array[Transform],
+ properties: util.Map[String, String]): Table = {
+ val options = new CaseInsensitiveStringMap(properties)
+ if (options.getBoolean("metadataColumn", false)) {
+ new StorageFilterTestTable(schema, options) with SupportsMetadataColumns
{
+ override def metadataColumns(): Array[MetadataColumn] = Array(new
MetadataColumn {
+ override def name(): String = StorageFilterTestSource.posName
+ override def dataType(): DataType = LongType
+ })
+ }
+ } else {
+ new StorageFilterTestTable(schema, options)
+ }
+ }
+}
+
+object StorageFilterTestSource {
+
+ /** Rows the source serves, `k` from 0 until numRows. */
+ val numRows: Int = 200
+
+ /** Rows a batch holds at most in a columnar read. */
+ val batchSize: Int = 64
+
+ /** The name of the metadata column the `metadataColumn` option adds. */
+ val posName: String = "_pos"
+
+ /** How many readers were created with a storage-filter request. */
+ val filteredReaders = new AtomicInteger(0)
+
+ /** How many of those requests asked for the checked column. */
+ val checkedColumnRequests = new AtomicInteger(0)
+
+ /** Rows the reader evaluated the storage filter against. */
+ val rowsEvaluated = new AtomicLong(0)
+
+ /** Rows the reader returned. */
+ val rowsReturned = new AtomicLong(0)
+
+ /** The serialized size of the last filter a reader got. */
+ val filterBytes = new AtomicLong(0)
+
+ /** The references of the filters the readers got, and the identities of
those filters. */
+ val filterReferences = mutable.Set.empty[Seq[String]]
+ val filterIdentities = mutable.Set.empty[Int]
+
+ /** How many times a runtime-filtering scan was handed filters. */
+ val runtimeFilterCalls = new AtomicInteger(0)
+
+ def reset(): Unit = {
+ filteredReaders.set(0)
+ checkedColumnRequests.set(0)
+ rowsEvaluated.set(0)
+ rowsReturned.set(0)
+ filterBytes.set(0)
+ filterReferences.synchronized(filterReferences.clear())
+ filterIdentities.synchronized(filterIdentities.clear())
+ runtimeFilterCalls.set(0)
+ }
+
+ /** The value of column `name` in the row of `key`. Any column but `k`, `w`
and `_pos` is `v`. */
+ def valueOf(name: String, key: Long): Long = name match {
+ case "k" | `posName` => key
+ case "w" => key * 100
+ case _ => key * 10
+ }
+
+ /** The columns of the row the filter evaluates, in that order. */
+ def columnsOf(filter: StorageFilter): Seq[String] =
filter.references.map(_.fieldNames.head).toSeq
+
+ // Records the filter a reader got, before the reader evaluates it.
+ private[connector] def record(filter: StorageFilter): Unit = {
+ val bytes = new ByteArrayOutputStream()
+ val out = new ObjectOutputStream(bytes)
+ try out.writeObject(filter) finally out.close()
Review Comment:
Thanks, fixed in
[`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538).
Only the first reader serializes the filter now, before it evaluates it.
--
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]
