peter-toth commented on code in PR #59715:
URL: https://github.com/apache/spark/pull/59715#discussion_r4381903677
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -127,6 +127,29 @@ class DataSourceV2Strategy(session: SparkSession) extends
Strategy with Predicat
tableSpec.withNewLocation(newLoc)
}
+ /**
+ * The conjuncts of `filters` offered to the scan as storage filters, see
+ * `SupportsStorageFilters` and `StorageFilterPlanning.offeredConjuncts`.
Only runtime blooms for
+ * now, as on the file-source side. The columns left to skip are the data
columns, as there, so a
+ * scan that reads only the key and metadata columns is offered none.
+ */
+ private def storageFiltersFor(
+ relation: DataSourceV2ScanRelation,
+ filters: Seq[Expression]): Seq[Expression] = {
+ if (!session.sessionState.conf.v2StorageFilterPushdownEnabled) return Nil
+ if (relation.storageFilterAttrs.isEmpty) return Nil
Review Comment:
Thanks, fixed in
[`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538),
in that order.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -212,7 +236,8 @@ class DataSourceRDD(
private class PartitionIterator[T](
private var reader: PartitionReader[T],
rowCount: T => Int,
- taskState: TaskState) extends Iterator[T] {
+ taskState: TaskState,
+ check: Option[T => Unit] = None) extends Iterator[T] {
Review Comment:
Thanks, fixed in
[`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538).
##########
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,
Review Comment:
Thanks, fixed in
[`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538),
in both places.
##########
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()
Review Comment:
Thanks, fixed in
[`5141077`](https://github.com/apache/spark/pull/59715/commits/5141077cff2f505c2badf57dd6ba326a26610538).
The test source uses `Utils.serialize`, `RangeInputPartition` and
`Utils.stringToSeq`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceRDD.scala:
##########
@@ -187,11 +196,19 @@ class DataSourceRDD(
castPartition(split).inputPartition.iterator.flatMap { inputPartition =>
val iter = if (columnarReads) {
- val batchReader =
partitionReaderFactory.createColumnarReader(inputPartition)
- new PartitionIterator[ColumnarBatch](batchReader, _.numRows, taskState)
+ val batchReader = storageFilterRequest match {
Review Comment:
Thanks. I kept it as is. Each branch calls another pair of factory methods,
so a helper would take two functions to save a few lines.
--
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]
