This is an automated email from the ASF dual-hosted git repository.
szehon-ho pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 91333d9d513d [SPARK-58523][SQL] Add Catalyst runtime filtering
interface for DSv2 scans
91333d9d513d is described below
commit 91333d9d513deebef34f81b3acecd716d585bce5
Author: Szehon Ho <[email protected]>
AuthorDate: Mon Aug 10 16:07:33 2026 -0700
[SPARK-58523][SQL] Add Catalyst runtime filtering interface for DSv2 scans
### What changes were proposed in this pull request?
This PR adds an internal `SupportsRuntimeCatalystFiltering` mix-in for DSv2
`Scan`s that receives runtime filters as Catalyst `Expression`s instead of
connector `Predicate`s.
The interface is internal because it is aimed at Spark-integrated sources
that already evaluate Catalyst expressions for partition pruning (the same
class of sources that bind and interpret expressions, e.g. via
`PartitionPredicateImpl`), rather than as a general connector API that speaks
V2 `Predicate`s.
- `SupportsRuntimeCatalystFiltering` (in
`org.apache.spark.sql.internal.connector`) declares `filterAttributes()`,
`filter(Array[Expression])`, and `fullyPushedFilterAttributes()`. It is an
alternative to `SupportsRuntimeFiltering` / `SupportsRuntimeV2Filtering`, not
an extension of them: Spark takes exactly one runtime filtering path per scan,
and only one runtime filtering interface may be implemented. The name follows
`SupportsRuntimeV2Filtering`, since in DSv2 the `SupportsPushDow [...]
- `PushDownUtils.pushRuntimeFilters` gains a branch for the new interface.
All runtime filters are pushed in a single `filter` call, with no translation
to connector predicates, so filters that have no V2 translation still reach the
source. Unavailable DPP filters (`TrueLiteral` / result-unavailable) are
dropped explicitly, and filters are screened with the same pushability guard as
the V2 `PartitionPredicate` path (`isPushablePartitionFilter`). The V2 path is
unchanged and is matched first.
- The unwrapping logic in `createRuntimePartitionPredicates` (unwrap DPP,
literalize scalar subqueries) is extracted into `unwrapRuntimeFilterExpression`
and shared by both paths.
- `DataSourceV2ScanRelation.runtimeFilterAttrs` /
`fullyPushedRuntimeFilterAttrs`, `PartitionPruning.getFilterableTableScan`, and
`RowLevelOperationRuntimeGroupFiltering` recognize the new interface, so DPP,
scalar subquery, and row-level group runtime filters are derived for these
scans. Row-level group filtering covers both group-based and delta-based
operations, matching the coverage the V2 interface already has.
- A scan can report `fullyPushedFilterAttributes()`. A runtime filter that
only references those attributes is considered fully pushed and
`DataSourceV2Strategy` drops it from the post-scan `FilterExec`, instead of
both pushing it and re-evaluating it. Dynamic pruning filters were already
excluded from the post-scan filter list. That drop applies the same pushability
test as pushdown, so a filter can never be removed from the post-scan filters
and then refused at the scan. The Javadoc [...]
- Because Spark resolves the two interfaces in favor of
`SupportsRuntimeV2Filtering` wherever it dispatches on them, a scan
implementing both is rejected with an internal error rather than silently
taking the V2 push path while its `fullyPushedFilterAttributes()` are trusted.
The check runs from both runtime filter attribute sets on
`DataSourceV2ScanRelation` and from `pushRuntimeFilters`.
- Javadoc on `SupportsRuntimeFiltering` / `SupportsRuntimeV2Filtering`
records the top-level-attribute requirement for `filterAttributes()`, and
clarifies that `pushedPredicates()` reports predicates that fully or partially
help pruning rather than predicates Spark can skip evaluating.
Dropping the scalar subquery filter from the post-scan `FilterExec` is
limited to `SupportsRuntimeCatalystFiltering` in this PR. Doing the same for
`SupportsRuntimeFiltering` / `SupportsRuntimeV2Filtering` is deferred to a
subsequent PR, because on those interfaces the filter is pushed as a translated
connector `Predicate` and the source decides what it can fully evaluate from
the semantics of that `Predicate`, not from the attribute alone. A source
therefore cannot promise to fully e [...]
### Why are the changes needed?
The existing runtime filtering interfaces receive connector predicates, so
a runtime filter is only pushed if it can be translated to a V2 `Predicate`. A
filter with no translation - `RLIKE` against a subquery result, a UDF over a
partition column - is silently dropped and never reaches the data source, even
when the source could use it to prune input partitions. Sources that already
work with Catalyst expressions have no way to receive them.
The interface also gives sources a way to state that they fully evaluate a
runtime filter, so Spark can skip the redundant post-scan evaluation of that
filter.
### Does this PR introduce _any_ user-facing change?
No. `SupportsRuntimeCatalystFiltering` lives in
`org.apache.spark.sql.internal.connector` and is not public API. The changes to
the public `SupportsRuntimeFiltering` and `SupportsRuntimeV2Filtering`
interfaces are documentation only.
### How was this patch tested?
Added `DataSourceV2CatalystRuntimeFilterSuite`, backed by
`InMemoryCatalystRuntimeFilterTable` /
`InMemoryTableCatalystRuntimeFilterCatalog`, covering:
- a scalar subquery on a partition column pushed as a Catalyst expression,
and still evaluated after the scan because the attribute is not declared fully
pushed;
- the same predicate no longer evaluated after the scan once the attribute
is declared in `fullyPushedFilterAttributes`, with matching and nonmatching
partitions so the fixture must prune itself;
- a non-deterministic predicate over a fully pushed attribute, which stays
in the post-scan `FilterExec`;
- a predicate that references both a fully pushed attribute and one that is
not, so it stays in the post-scan `FilterExec`;
- a filter with no V2 translation (`part RLIKE (subquery)`) pushed instead
of dropped, asserting that `translateScalarSubqueryFilterV2` finds no
translation for it;
- arithmetic around a scalar subquery (`part > (subquery) + 1`), checking
that the subquery is literalized and the surrounding expression reaches the
scan intact - this one does translate to a V2 predicate, so it covers
expression shape rather than untranslatability;
- a DPP filter pushed as the raw `InSubqueryExec` expression;
- a scan implementing both runtime filtering interfaces being rejected;
- a filter on a column outside `filterAttributes` not being pushed;
- no runtime filter meaning `filter()` is never called.
The fixture scan binds and interprets any pushed expression whose
references are all partition columns and prunes its partitions accordingly, the
same way `PartitionPredicateImpl` does, so a fully pushed predicate is enforced
by the source alone once Spark drops the post-scan `FilterExec`. That logic
lives in a `CatalystRuntimeFilteringScan` trait in `InMemoryBaseTable` and is
shared by the runtime filter and row-level operation fixtures.
Added `RowLevelOperationCatalystRuntimeFilterSuiteBase`, with
`GroupBasedRowLevelOperationCatalystRuntimeFilterSuite` and
`DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite`, covering
`RowLevelOperationRuntimeGroupFiltering` for Catalyst runtime-filtering scans:
- UPDATE and MERGE for both group-based and delta-based operations. Each
asserts that the scan declares the group key in `filterAttributes`, that every
scan node carries a dynamic pruning filter keyed on it, that the filter
subquery projects only the columns the row-level condition needs and resolves
to the expected groups, that the same Catalyst expression reaches the connector
once per scan node, and that the scan then reads only those groups.
- group-based DELETE, which additionally checks the replaced partition and
the delete metrics;
- delta-based DELETE, which asserts that no filter is injected: that plan
scans the row ID, the condition columns and the metadata columns, so the group
key is not available to filter on.
```
build/sbt 'sql/testOnly *DataSourceV2CatalystRuntimeFilterSuite
*RowLevelOperationCatalystRuntimeFilterSuite
*DataSourceV2EnhancedRuntimePartitionFilterSuite'
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Cursor
Closes #57727 from szehon-ho/catalyst-runtime-filter-pushdown.
Authored-by: Szehon Ho <[email protected]>
Signed-off-by: Szehon Ho <[email protected]>
---
.../connector/read/SupportsRuntimeFiltering.java | 4 +
.../connector/read/SupportsRuntimeV2Filtering.java | 18 +-
.../datasources/v2/DataSourceV2Relation.scala | 43 ++-
.../SupportsRuntimeCatalystFiltering.scala | 103 ++++++
.../sql/connector/catalog/InMemoryBaseTable.scala | 105 +++++-
.../InMemoryCatalystRuntimeFilterTable.scala | 111 ++++++
.../catalog/InMemoryRowLevelOperationTable.scala | 68 +++-
...InMemoryTableCatalystRuntimeFilterCatalog.scala | 50 +++
.../datasources/v2/DataSourceV2Strategy.scala | 13 +-
.../execution/datasources/v2/PushDownUtils.scala | 93 ++++-
.../dynamicpruning/PartitionPruning.scala | 9 +
.../RowLevelOperationRuntimeGroupFiltering.scala | 37 +-
.../DataSourceV2CatalystRuntimeFilterSuite.scala | 400 +++++++++++++++++++++
...wLevelOperationCatalystRuntimeFilterSuite.scala | 50 +++
...wLevelOperationCatalystRuntimeFilterSuite.scala | 60 ++++
...elOperationCatalystRuntimeFilterSuiteBase.scala | 229 ++++++++++++
16 files changed, 1331 insertions(+), 62 deletions(-)
diff --git
a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java
b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java
index 927d4a53e22f..067202a36270 100644
---
a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java
+++
b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java
@@ -38,6 +38,10 @@ public interface SupportsRuntimeFiltering extends
SupportsRuntimeV2Filtering {
* <p>
* Spark will call {@link #filter(Filter[])} if it can derive a runtime
* predicate for any of the filter attributes.
+ * <p>
+ * Each reference must be a top-level attribute present in {@link
Scan#readSchema()}.
+ * Nested references and attributes pruned out of the read schema fail to
resolve when
+ * Spark builds the scan relation.
*/
NamedReference[] filterAttributes();
diff --git
a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java
b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java
index 94dbc3865958..6b286f041b01 100644
---
a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java
+++
b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java
@@ -28,8 +28,9 @@ import org.apache.spark.sql.sources.Filter;
* filter initially planned {@link InputPartition}s using predicates Spark
infers at runtime.
* This interface is very similar to {@link SupportsRuntimeFiltering} except
it uses
* data source V2 {@link Predicate} instead of data source V1 {@link Filter}.
- * {@link SupportsRuntimeV2Filtering} is preferred over {@link
SupportsRuntimeFiltering}
- * and only one of them should be implemented by the data sources.
+ * {@link SupportsRuntimeV2Filtering} is preferred over {@link
SupportsRuntimeFiltering}.
+ * A scan must not implement SupportsRuntimeCatalystFiltering together with
this interface;
+ * Spark rejects such a scan.
* <p>
* <b>Iterative filtering:</b> When {@link #supportsIterativePushdown()}
returns true,
* {@link #filter(Predicate[])} may be called <i>multiple times</i> on the same
@@ -50,6 +51,10 @@ public interface SupportsRuntimeV2Filtering extends Scan {
* <p>
* Spark will call {@link #filter(Predicate[])} if it can derive a runtime
* predicate for any of the filter attributes.
+ * <p>
+ * Each reference must be a top-level attribute present in {@link
Scan#readSchema()}.
+ * Nested references and attributes pruned out of the read schema fail to
resolve when
+ * Spark builds the scan relation.
*/
NamedReference[] filterAttributes();
@@ -72,6 +77,11 @@ public interface SupportsRuntimeV2Filtering extends Scan {
* The implementation must accumulate state across all calls so that
* {@link #pushedPredicates()} can return predicates from all of them.
* <p>
+ * Independently of {@link #supportsIterativePushdown()}, this method may
also be called once
+ * per scan node when a plan holds several scan nodes sharing one {@link
Scan} instance (e.g.
+ * the two branches of a group-based UPDATE). Implementations must
accumulate state across
+ * those calls as well.
+ * <p>
* Note that Spark will call {@link Scan#toBatch()} again after filtering
the scan at runtime.
*
* @param predicates data source V2 predicates used to filter the scan at
runtime
@@ -82,6 +92,10 @@ public interface SupportsRuntimeV2Filtering extends Scan {
* Returns the predicates that are pushed to the data source via
* {@link #filter(Predicate[])}.
* <p>
+ * These are not fully pushed predicates: Spark may still evaluate them
after the scan.
+ * They are predicates that fully or partially help the data source prune
initially planned
+ * {@link InputPartition}s.
+ * <p>
* When iterative filtering is supported and {@link #filter(Predicate[])}
was called
* multiple times, this method must return predicates from <i>all</i> calls.
* <p>
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala
index 52c2674adc2c..1c7a7abe4a8a 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala
@@ -35,7 +35,7 @@ import
org.apache.spark.sql.connector.expressions.{FieldReference, NamedReferenc
import org.apache.spark.sql.connector.read.{Scan, Statistics => V2Statistics,
SupportsReportStatistics, SupportsRuntimeV2Filtering}
import org.apache.spark.sql.connector.read.colstats.{ColumnStatistics,
Histogram => V2Histogram, HistogramBin => V2HistogramBin}
import org.apache.spark.sql.connector.read.streaming.{Offset, SparkDataStream}
-import org.apache.spark.sql.internal.connector.V2StatisticsUtils
+import
org.apache.spark.sql.internal.connector.{SupportsRuntimeCatalystFiltering,
V2StatisticsUtils}
import org.apache.spark.sql.types.{DataType, StructType}
import org.apache.spark.sql.util.CaseInsensitiveStringMap
import org.apache.spark.util.ArrayImplicits._
@@ -196,14 +196,33 @@ case class DataSourceV2ScanRelation(
/**
* Resolved attributes that the scan declares for runtime filtering via
- * [[SupportsRuntimeV2Filtering.filterAttributes]]. Empty when the scan
- * does not implement [[SupportsRuntimeV2Filtering]] or exposes no
attributes.
+ * [[SupportsRuntimeV2Filtering.filterAttributes]] or
+ * [[SupportsRuntimeCatalystFiltering.filterAttributes]]. Empty when the scan
+ * implements neither interface or exposes no attributes.
*/
- lazy val runtimeFilterAttrs: AttributeSet = scan match {
- case s: SupportsRuntimeV2Filtering =>
- AttributeSet(V2ExpressionUtils.resolveRefs[Attribute](
- s.filterAttributes.toImmutableArraySeq, this))
- case _ => AttributeSet.empty
+ lazy val runtimeFilterAttrs: AttributeSet = {
+ checkRuntimeFilteringInterfaces()
+ val filterAttrs = scan match {
+ case s: SupportsRuntimeV2Filtering => s.filterAttributes
+ case s: SupportsRuntimeCatalystFiltering => s.filterAttributes()
+ case _ => Array.empty[NamedReference]
+ }
+ AttributeSet(V2ExpressionUtils.resolveRefs[Attribute](
+ filterAttrs.toImmutableArraySeq, this))
+ }
+
+ /**
+ * Resolved attributes for which a Catalyst runtime-filtering scan fully
evaluates predicates.
+ * Empty for a [[SupportsRuntimeV2Filtering]] scan, which keeps its
post-scan filters.
+ */
+ lazy val fullyPushedRuntimeFilterAttrs: AttributeSet = {
+ checkRuntimeFilteringInterfaces()
+ val filterAttrs = scan match {
+ case s: SupportsRuntimeCatalystFiltering =>
s.fullyPushedFilterAttributes()
+ case _ => Array.empty[NamedReference]
+ }
+ AttributeSet(V2ExpressionUtils.resolveRefs[Attribute](
+ filterAttrs.toImmutableArraySeq, this))
}
override def name: String = relation.name
@@ -252,6 +271,14 @@ case class DataSourceV2ScanRelation(
Statistics(sizeInBytes = conf.defaultSizeInBytes)
}
+ private def checkRuntimeFilteringInterfaces(): Unit = scan match {
+ case _: SupportsRuntimeV2Filtering with SupportsRuntimeCatalystFiltering =>
+ throw SparkException.internalError(
+ "A scan must not implement both SupportsRuntimeV2Filtering and " +
+ s"SupportsRuntimeCatalystFiltering, but ${scan.getClass.getName}
implements both.")
+ case _ =>
+ }
+
override def doCanonicalize(): DataSourceV2ScanRelation = {
this.copy(
relation = this.relation.copy(
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala
new file mode 100644
index 000000000000..5e5900a21196
--- /dev/null
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.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.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.connector.expressions.NamedReference
+import org.apache.spark.sql.connector.read.Scan
+
+/**
+ * A mix-in interface for [[Scan]]. Data sources can implement this interface
if they can
+ * filter initially planned
[[org.apache.spark.sql.connector.read.InputPartition]]s using
+ * Catalyst [[Expression]]s Spark infers at runtime.
+ * A scan must not implement this interface together with
+ * [[org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering]] or its
subinterface
+ * [[org.apache.spark.sql.connector.read.SupportsRuntimeFiltering]]; Spark
rejects such a scan.
+ *
+ * Spark considers a runtime predicate fully pushed when all attributes
referenced by the
+ * predicate are returned by [[fullyPushedFilterAttributes]]. Fully pushed
predicates are not
+ * evaluated again after the scan.
+ *
+ * Note that Spark will push runtime filters only if they are beneficial.
+ */
+trait SupportsRuntimeCatalystFiltering extends Scan {
+
+ /**
+ * Returns attributes this scan can be filtered by at runtime.
+ *
+ * Spark will call [[filter]] if it can derive a runtime filter for any of
these attributes.
+ * Each reference must be a top-level attribute present in
[[Scan.readSchema]]. Nested
+ * references and attributes pruned out of the read schema fail to resolve
when Spark builds
+ * the scan relation.
+ */
+ def filterAttributes(): Array[NamedReference]
+
+ /**
+ * Returns attributes for which this scan fully evaluates runtime predicates.
+ *
+ * Any runtime predicate that references only attributes in this set is
considered fully pushed
+ * and will not be evaluated again after the scan. These attributes must
also be returned by
+ * [[filterAttributes]]. Each attribute's value must therefore be fixed
within every
+ * [[org.apache.spark.sql.connector.read.InputPartition]] the scan returns,
since pruning
+ * partitions cannot fully evaluate a predicate on a column that varies
within a partition.
+ *
+ * Spark relies on the scan alone here, so the scan must return only
partitions it has proven
+ * satisfy such a predicate. Spark passes these expressions through as they
are, leaving both
+ * translation and capability checking to the scan, so evaluating one can
fail, for example on
+ * an ANSI cast or overflow error, or on a nested access that the scan
matches differently
+ * against its partition layout. Declare an attribute only when the scan can
evaluate every
+ * predicate over it.
+ *
+ * Each reference must be a top-level attribute present in
[[Scan.readSchema]]. Nested
+ * references and attributes pruned out of the read schema fail to resolve
when Spark builds
+ * the scan relation.
+ */
+ def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty
+
+ /**
+ * Filters this scan using runtime Catalyst expressions.
+ *
+ * The provided expressions must be interpreted as a set of predicates that
are ANDed together.
+ * Implementations may use the expressions to prune initially planned
+ * [[org.apache.spark.sql.connector.read.InputPartition]]s.
+ *
+ * An expression may access nested fields of an attribute returned by
[[filterAttributes]], as
+ * that attribute is required to be top-level. The scan is responsible for
matching such
+ * accesses against its own partition layout.
+ *
+ * Spark may call this method more than once for the same scan instance: a
plan can hold several
+ * scan nodes sharing one scan (e.g. the two branches of a group-based
UPDATE), and each pushes
+ * its own copy of the runtime filters. Implementations must treat
successive calls as additive,
+ * ANDing the new expressions with those already pushed rather than
replacing them.
+ *
+ * If the scan also implements
+ * [[org.apache.spark.sql.connector.read.SupportsReportPartitioning]], it
must preserve
+ * the originally reported partitioning during runtime filtering. While
applying runtime
+ * predicates, the scan may detect that some
+ * [[org.apache.spark.sql.connector.read.InputPartition]]s have no matching
data, in which
+ * case it can either replace the initially planned
+ * [[org.apache.spark.sql.connector.read.InputPartition]]s that have no
matching data with
+ * empty [[org.apache.spark.sql.connector.read.InputPartition]]s, or report
only a subset of
+ * the original partition values (omitting those with no data) via
+ * [[org.apache.spark.sql.connector.read.Batch#planInputPartitions]]. The
scan must not
+ * report new partition values that were not present in the original
partitioning.
+ *
+ * Note that Spark will call [[Scan.toBatch]] again after filtering the scan
at runtime.
+ */
+ def filter(expressions: Array[Expression]): Unit
+}
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala
index fba80eb3d4cb..1b5d7bb75fe6 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala
@@ -25,11 +25,11 @@ import java.util.OptionalLong
import java.util.concurrent.atomic.AtomicLong
import scala.collection.mutable
-import scala.collection.mutable.ListBuffer
+import scala.collection.mutable.{ArrayBuffer, ListBuffer}
import scala.jdk.CollectionConverters._
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Cast, EvalMode,
GenericInternalRow, JoinedRow, Literal, MetadataStructFieldWithLogicalName}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
BindReferences, Cast, EvalMode, Expression => CatalystExpression,
GenericInternalRow, JoinedRow, Literal, MetadataStructFieldWithLogicalName,
Predicate => CatalystPredicate}
import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, ArrayData,
CaseInsensitiveMap, CharVarcharUtils, DateTimeUtils, GenericArrayData, MapData,
ResolveDefaultColumns}
import org.apache.spark.sql.connector.catalog.constraints.Constraint
import org.apache.spark.sql.connector.distributions.{Distribution,
Distributions}
@@ -43,7 +43,7 @@ import
org.apache.spark.sql.connector.read.streaming.{MicroBatchStream, Offset}
import org.apache.spark.sql.connector.write._
import
org.apache.spark.sql.connector.write.streaming.{StreamingDataWriterFactory,
StreamingWrite}
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.internal.connector.{ColumnImpl,
SupportsStreamingUpdateAsAppend}
+import org.apache.spark.sql.internal.connector.{ColumnImpl,
SupportsRuntimeCatalystFiltering, SupportsStreamingUpdateAsAppend}
import org.apache.spark.sql.sources._
import org.apache.spark.sql.types._
import org.apache.spark.sql.util.CaseInsensitiveStringMap
@@ -521,23 +521,35 @@ abstract class InMemoryBaseTable(
private var _pushedFilters: Array[Filter] = Array.empty
override def build: Scan = {
- val scan = if (InMemoryBaseTable.this.ordering.nonEmpty) {
- new InMemoryBatchScanWithOrdering(
- data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq,
schema, tableSchema,
- options)
- } else {
- InMemoryBatchScan(
- data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq,
schema, tableSchema,
- options)
- }
- if (evaluableFilters.nonEmpty) {
- scan.filter(evaluableFilters)
+ val scan = createScan(
+ data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, schema,
tableSchema, options)
+ scan match {
+ case s: InMemoryBatchScan =>
+ if (evaluableFilters.nonEmpty) {
+ s.filter(evaluableFilters)
+ }
+ s.pushedFilters = _pushedFilters
+ case _ =>
}
- scan.pushedFilters = _pushedFilters
recordScanEvent(_pushedFilters)
scan
}
+ /**
+ * Creates the batch scan for [[build]].
+ */
+ protected def createScan(
+ partitions: Seq[InputPartition],
+ readSchema: StructType,
+ tableSchema: StructType,
+ options: CaseInsensitiveStringMap): BatchScanBaseClass = {
+ if (InMemoryBaseTable.this.ordering.nonEmpty) {
+ new InMemoryBatchScanWithOrdering(partitions, readSchema, tableSchema,
options)
+ } else {
+ InMemoryBatchScan(partitions, readSchema, tableSchema, options)
+ }
+ }
+
override def pruneColumns(requiredSchema: StructType): Unit = {
// The required schema could contain conflict-renamed metadata columns,
so we need to match
// them by their logical (original) names, not their current names.
@@ -691,6 +703,69 @@ abstract class InMemoryBaseTable(
new InMemoryMicroBatchStream(readSchema, tableSchema)
}
+ /**
+ * Reference implementation of [[SupportsRuntimeCatalystFiltering.filter]]
for the in-memory
+ * fixtures: records what was pushed, and for expressions referencing only
partition columns
+ * binds them against the partition key and drops partitions that do not
match. Binding and
+ * interpreting rather than pattern matching a fixed set of operators is
what lets the fixture
+ * honor an arbitrary pushed expression, the same way
`PartitionPredicateImpl` does. Mixing
+ * classes supply their own `filterAttributes()`.
+ */
+ trait CatalystRuntimeFilteringScan extends SupportsRuntimeCatalystFiltering {
+ self: BatchScanBaseClass =>
+
+ /** The full table schema, used to locate partition columns pruned out of
`readSchema`. */
+ protected def tableSchema: StructType
+
+ private val catalystPredicates = ArrayBuffer.empty[CatalystExpression]
+
+ override def filter(expressions: Array[CatalystExpression]): Unit = {
+ catalystPredicates ++= expressions
+ val partAttrs = partitionAttributes
+ if (partAttrs.isEmpty) return
+
+ val resolver = SQLConf.get.resolver
+ expressions.foreach { expr =>
+ val remapped = expr.transform {
+ case a: AttributeReference =>
+ partAttrs.find(p => resolver(p.name, a.name)).getOrElse(a)
+ }
+ // Only evaluate expressions whose refs are all partition columns, so
we can bind
+ // against the partition key InternalRow (same approach as
PartitionPredicateImpl).
+ if (remapped.references.forall(r => partAttrs.exists(_.exprId ==
r.exprId))) {
+ val bound = BindReferences.bindReference(remapped, partAttrs)
+ val pred = CatalystPredicate.createInterpreted(bound)
+ self.data = self.data.filter { p =>
+ try {
+ pred.eval(p.asInstanceOf[BufferedRows].partitionKey())
+ } catch {
+ // Keep the partition on eval failure, which is safe here
because every predicate
+ // the fixture pushes evaluates cleanly.
`PartitionPredicateImpl` fails open for a
+ // reason of its own: Spark keeps the post-scan `FilterExec` on
that path, so
+ // failing open costs just a pruning opportunity. A scan
declaring an attribute in
+ // `fullyPushedFilterAttributes()` stands alone as the
evaluator, so keeping an
+ // unevaluated partition would return nonmatching rows.
+ case _: Exception => true
+ }
+ }
+ }
+ }
+ }
+
+ /** Predicates recorded by [[filter]], for test assertions only. */
+ def pushedCatalystPredicates: Seq[CatalystExpression] =
catalystPredicates.toSeq
+
+ /** AttributeReferences matching the partition-key InternalRow field
order. */
+ private def partitionAttributes: Seq[AttributeReference] = {
+ partitioning.flatMap(_.references()).flatMap { ref =>
+ val name = ref.fieldNames.mkString(".")
+ readSchema.find(_.name == name).orElse(tableSchema.find(_.name ==
name)).map { f =>
+ AttributeReference(f.name, f.dataType, f.nullable)()
+ }
+ }.toSeq
+ }
+ }
+
case class InMemoryBatchScan(
var _data: Seq[InputPartition],
readSchema: StructType,
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala
new file mode 100644
index 000000000000..0f32028f20a9
--- /dev/null
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala
@@ -0,0 +1,111 @@
+/*
+ * 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.catalog
+
+import java.util
+
+import InMemoryCatalystRuntimeFilterTable._
+
+import org.apache.spark.sql.connector.expressions.{NamedReference, Transform}
+import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder}
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import org.apache.spark.util.ArrayImplicits._
+
+/**
+ * In-memory table whose batch scan mixes in [[CatalystRuntimeFilteringScan]],
so runtime filters
+ * arrive as Catalyst expressions rather than connector predicates.
+ *
+ * Table properties:
+ * - `filter-attributes` (default: all partition cols): comma-separated list
of
+ * column names to expose from `filterAttributes`.
+ * - `fully-pushed-filter-attributes` (default: none): comma-separated list of
+ * column names to expose from `fullyPushedFilterAttributes`.
+ */
+class InMemoryCatalystRuntimeFilterTable(
+ name: String,
+ columns: Array[Column],
+ partitioning: Array[Transform],
+ properties: util.Map[String, String])
+ extends InMemoryTableWithV2Filter(name, columns, partitioning, properties) {
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder
= {
+ new InMemoryCatalystRuntimeFilterScanBuilder(schema, options)
+ }
+
+ class InMemoryCatalystRuntimeFilterScanBuilder(
+ tableSchema: StructType,
+ options: CaseInsensitiveStringMap)
+ extends InMemoryScanBuilder(tableSchema, options) {
+ override def build: Scan = InMemoryCatalystRuntimeFilterBatchScan(
+ data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq,
+ schema, tableSchema, options)
+ }
+
+ /**
+ * Scan that receives runtime filters as Catalyst expressions. Pruning comes
from
+ * [[CatalystRuntimeFilteringScan]], so fully-pushed predicates are enforced
when Spark drops
+ * the post-scan [[org.apache.spark.sql.execution.FilterExec]].
+ */
+ case class InMemoryCatalystRuntimeFilterBatchScan(
+ var _data: Seq[InputPartition],
+ readSchema: StructType,
+ tableSchema: StructType,
+ options: CaseInsensitiveStringMap)
+ extends BatchScanBaseClass(_data, readSchema, tableSchema)
+ with CatalystRuntimeFilteringScan {
+
+ private val restrictedFilterAttrs: Option[Set[String]] =
+
Option(InMemoryCatalystRuntimeFilterTable.this.properties.get(FilterAttributesKey))
+ .map(_.split(",").map(_.trim).toSet)
+
+ override def filterAttributes(): Array[NamedReference] = {
+ val scanFields = readSchema.fields.map(_.name).toSet
+ partitioning.flatMap(_.references()).filter { ref =>
+ val name = ref.fieldNames.mkString(".")
+ scanFields.contains(name) &&
+ restrictedFilterAttrs.forall(_.contains(name))
+ }
+ }
+
+ override def fullyPushedFilterAttributes(): Array[NamedReference] = {
+ val fullyPushedFilterAttrs = Option(
+
InMemoryCatalystRuntimeFilterTable.this.properties.get(FullyPushedFilterAttributesKey))
+ .map(_.split(",").map(_.trim).toSet)
+ .getOrElse(Set.empty)
+ filterAttributes().filter { ref =>
+ fullyPushedFilterAttrs.contains(ref.fieldNames.mkString("."))
+ }
+ }
+
+ }
+}
+
+object InMemoryCatalystRuntimeFilterTable {
+ /**
+ * Table property: comma-separated column names to expose from
+ * filterAttributes. Default: all partition columns.
+ */
+ private[catalog] val FilterAttributesKey = "filter-attributes"
+
+ /**
+ * Table property: comma-separated column names to expose from
+ * fullyPushedFilterAttributes. Default: none.
+ */
+ private[catalog] val FullyPushedFilterAttributesKey =
"fully-pushed-filter-attributes"
+}
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala
index fd02f926e814..6cb23d505f78 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala
@@ -25,7 +25,7 @@ import
org.apache.spark.sql.connector.catalog.constraints.Constraint
import org.apache.spark.sql.connector.distributions.{Distribution,
Distributions}
import org.apache.spark.sql.connector.expressions.{FieldReference,
LogicalExpressions, NamedReference, SortDirection, SortOrder, Transform}
import org.apache.spark.sql.connector.expressions.filter.Predicate
-import org.apache.spark.sql.connector.read.{Scan, ScanBuilder}
+import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder}
import org.apache.spark.sql.connector.write.{BatchWrite, DeltaBatchWrite,
DeltaWrite, DeltaWriteBuilder, DeltaWriter, DeltaWriterFactory,
LogicalWriteInfo, PhysicalWriteInfo, RequiresDistributionAndOrdering,
RowLevelOperation, RowLevelOperationBuilder, RowLevelOperationInfo,
SupportsDelta, Write, WriteBuilder, WriterCommitMessage}
import org.apache.spark.sql.connector.write.RowLevelOperation.Command
import org.apache.spark.sql.types.StructType
@@ -78,7 +78,10 @@ class InMemoryRowLevelOperationTable private (
private final val SUPPORTS_DELTAS = "supports-deltas"
private final val SPLIT_UPDATES = "split-updates"
private final val NO_METADATA = "no-metadata"
+ private final val USE_CATALYST_RUNTIME_FILTERING =
"use-catalyst-runtime-filtering"
private final val noMetadata = properties.getOrDefault(NO_METADATA, "false")
== "true"
+ private final val useCatalystRuntimeFiltering =
+ properties.getOrDefault(USE_CATALYST_RUNTIME_FILTERING, "false") == "true"
// used in row-level operation tests to verify replaced partitions
var replacedPartitions: Seq[Seq[Any]] = Seq.empty
@@ -133,7 +136,7 @@ class InMemoryRowLevelOperationTable private (
case class PartitionBasedOperation(command: Command, options:
CaseInsensitiveStringMap)
extends RowLevelOperation with RowLevelOperationWithOptions {
- var configuredScan: InMemoryBatchScan = _
+ var configuredScan: BatchScanBaseClass = _
override def requiredMetadataAttributes(): Array[NamedReference] = {
if (noMetadata) {
@@ -144,12 +147,8 @@ class InMemoryRowLevelOperationTable private (
}
override def newScanBuilder(options: CaseInsensitiveStringMap):
ScanBuilder = {
- new InMemoryScanBuilder(schema, options) {
- override def build: Scan = {
- val scan = super.build()
- configuredScan = scan.asInstanceOf[InMemoryBatchScan]
- scan
- }
+ newRowLevelScanBuilder(options) { scan =>
+ configuredScan = scan
}
}
@@ -186,7 +185,7 @@ class InMemoryRowLevelOperationTable private (
override def description(): String = "InMemoryPartitionReplaceOperation"
}
- private case class PartitionBasedReplaceData(scan: InMemoryBatchScan)
+ private case class PartitionBasedReplaceData(scan: BatchScanBaseClass)
extends TestBatchWrite {
override protected def doCommit(
@@ -216,7 +215,7 @@ class InMemoryRowLevelOperationTable private (
override def rowId(): Array[NamedReference] = Array(PK_COLUMN_REF)
override def newScanBuilder(options: CaseInsensitiveStringMap):
ScanBuilder = {
- new InMemoryScanBuilder(schema, options)
+ newRowLevelScanBuilder(options)(_ => ())
}
override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
@@ -267,6 +266,55 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+
+ /**
+ * Builds a scan for row-level operations. When
`use-catalyst-runtime-filtering` is set, the
+ * scan mixes in [[CatalystRuntimeFilteringScan]] so group filtering goes
through the Catalyst
+ * path.
+ */
+ private def newRowLevelScanBuilder(
+ options: CaseInsensitiveStringMap)(
+ onBuild: BatchScanBaseClass => Unit): ScanBuilder = {
+ new InMemoryScanBuilder(schema, options) {
+ override protected def createScan(
+ partitions: Seq[InputPartition],
+ readSchema: StructType,
+ tableSchema: StructType,
+ options: CaseInsensitiveStringMap): BatchScanBaseClass = {
+ if (useCatalystRuntimeFiltering) {
+ InMemoryCatalystRowLevelBatchScan(partitions, readSchema,
tableSchema, options)
+ } else {
+ super.createScan(partitions, readSchema, tableSchema, options)
+ }
+ }
+
+ override def build: Scan = {
+ val scan = super.build().asInstanceOf[BatchScanBaseClass]
+ onBuild(scan)
+ scan
+ }
+ }
+ }
+
+ /**
+ * Row-level batch scan that receives runtime filters as Catalyst
expressions. Pruning comes
+ * from [[CatalystRuntimeFilteringScan]], so group filtering actually drops
partitions (needed
+ * for `replacedPartitions` assertions).
+ */
+ case class InMemoryCatalystRowLevelBatchScan(
+ var _data: Seq[InputPartition],
+ readSchema: StructType,
+ tableSchema: StructType,
+ options: CaseInsensitiveStringMap)
+ extends BatchScanBaseClass(_data, readSchema, tableSchema)
+ with CatalystRuntimeFilteringScan {
+
+ override def filterAttributes(): Array[NamedReference] = {
+ val scanFields = readSchema.fields.map(_.name).toSet
+ partitioning.flatMap(_.references())
+ .filter(ref => scanFields.contains(ref.fieldNames.mkString(".")))
+ }
+ }
}
private class DeltaBufferedRowsWriterFactory(schema: StructType) extends
DeltaWriterFactory {
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala
new file mode 100644
index 000000000000..b8415eae3e15
--- /dev/null
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala
@@ -0,0 +1,50 @@
+/*
+ * 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.catalog
+
+import java.util
+
+import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException
+import org.apache.spark.sql.connector.expressions.Transform
+
+class InMemoryTableCatalystRuntimeFilterCatalog extends InMemoryTableCatalog {
+ import CatalogV2Implicits._
+
+ override def createTable(
+ ident: Identifier,
+ columns: Array[Column],
+ partitions: Array[Transform],
+ properties: util.Map[String, String]): Table = {
+ if (tables.containsKey(ident)) {
+ throw new TableAlreadyExistsException(ident.asMultipartIdentifier)
+ }
+
+ InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties)
+
+ val tableName = s"$name.${ident.quoted}"
+ val table = new InMemoryCatalystRuntimeFilterTable(
+ tableName, columns, partitions, properties)
+ tables.put(ident, table)
+ namespaces.putIfAbsent(ident.namespace.toList, Map())
+ table
+ }
+
+ override def createTable(ident: Identifier, tableInfo: TableInfo): Table = {
+ createTable(ident, tableInfo.columns(), tableInfo.partitions(),
tableInfo.properties)
+ }
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala
index ce9092ea810b..8805bfe75298 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala
@@ -171,7 +171,9 @@ class DataSourceV2Strategy(session: SparkSession) extends
Strategy with Predicat
// Extract scalar subquery filters on runtime-filterable columns for
runtime pushdown.
// These filters stay in postScanFilters for correctness (FilterExec
above scan),
// but are also routed into runtimeFilters so BatchScanExec can use them
for
- // partition pruning via SupportsRuntimeV2Filtering.filter().
+ // partition pruning via SupportsRuntimeV2Filtering.filter(). The
exceptions are filters
+ // that only reference attributes the scan fully evaluates, which are
dropped from
+ // postScanFilters below.
// Non-deterministic filters are not routed: they would be pushed to the
source for
// pruning while the FilterExec above the scan re-evaluates them, so the
two evaluations
// may disagree and rows the source pruned away could not be recovered.
This is the
@@ -186,6 +188,12 @@ class DataSourceV2Strategy(session: SparkSession) extends
Strategy with Predicat
} else {
Seq.empty
}
+ // Screen with the same test pushdown applies, or a filter dropped here
and rejected there
+ // would be evaluated nowhere.
+ val fullyPushedRuntimeFilters = scalarSubqueryFilters.filter { f =>
+ f.references.subsetOf(relation.fullyPushedRuntimeFilterAttrs) &&
+ PushDownUtils.isPushablePartitionFilter(f, includeSubquery = true)
+ }
// dynamicFilters need no such check: a DynamicPruningSubquery over a
non-deterministic
// filtering plan is itself non-deterministic, so
CleanupDynamicPruningFilters has already
// rewritten it to TrueLiteral by the time we get here.
@@ -194,7 +202,8 @@ class DataSourceV2Strategy(session: SparkSession) extends
Strategy with Predicat
val batchExec = BatchScanExec(relation.output, relation.scan,
runtimeFilters,
relation.ordering, relation.relation.table,
relation.keyGroupedPartitioning)
DataSourceV2Strategy.withProjectAndFilter(
- project, postScanFilters, batchExec, !batchExec.supportsColumnar) ::
Nil
+ project, postScanFilters.diff(fullyPushedRuntimeFilters),
+ batchExec, !batchExec.supportsColumnar) :: Nil
case PhysicalOperation(p, f, r: StreamingDataSourceV2ScanRelation)
if r.startOffset.isDefined && r.endOffset.isDefined =>
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala
index 5e1e22a0c66f..fa330212370f 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala
@@ -22,7 +22,7 @@ import scala.collection.mutable
import org.apache.spark.SparkException
import org.apache.spark.internal.{Logging, LogKeys}
import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
AttributeSet, DynamicPruning, DynamicPruningExpression, Expression,
ExpressionSet, GetStructField, NamedExpression, PythonUDF, SchemaPruning,
SubqueryExpression, V2ExpressionUtils}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
AttributeSet, DynamicPruning, DynamicPruningExpression, Expression,
ExpressionSet, GetStructField, Literal, NamedExpression, PythonUDF,
SchemaPruning, SubqueryExpression, V2ExpressionUtils}
import org.apache.spark.sql.catalyst.plans.logical.SampleMethod
import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning,
Partitioning}
import org.apache.spark.sql.catalyst.types.DataTypeUtils
@@ -35,7 +35,7 @@ import org.apache.spark.sql.connector.read.{HasPartitionKey,
InputPartition, Sam
import org.apache.spark.sql.execution.{InSubqueryExec, ScalarSubquery =>
ExecScalarSubquery}
import org.apache.spark.sql.execution.datasources.{DataSourceStrategy,
DataSourceUtils}
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.internal.connector.{PartitionPredicateField,
PartitionPredicateImpl, SupportsPushDownCatalystFilters}
+import org.apache.spark.sql.internal.connector.{PartitionPredicateField,
PartitionPredicateImpl, SupportsPushDownCatalystFilters,
SupportsRuntimeCatalystFiltering}
import org.apache.spark.sql.sources
import org.apache.spark.sql.types.{StructField, StructType}
import org.apache.spark.util.ArrayImplicits.SparkArrayOps
@@ -167,14 +167,20 @@ object PushDownUtils extends Logging {
* the first pass are used to derive PartitionPredicates in the second pass,
avoiding duplicate
* pushdown.
*
- * Note: Do not call multiple times for the same `scan` instance;
- * [[SupportsRuntimeV2Filtering.filter]] is mutating.
+ * Note: `filter` is mutating, and Spark may call this more than once for
the same `scan`
+ * instance: a plan can hold several scan nodes sharing one scan (e.g. the
two branches of a
+ * group-based UPDATE), each pushing its own copy of the runtime filters.
Successive calls are
+ * additive: a scan ANDs the newly pushed predicates with those it already
holds.
*
* Note: `runtimeFilters` must not contain non-deterministic filters. A
runtime filter is also
* evaluated by the `FilterExec` above the scan, so pushing a
non-deterministic one would
* evaluate it twice with different results. `DataSourceV2Strategy` enforces
this where
* `runtimeFilters` is built (SPARK-58207).
*
+ * A scan implementing [[SupportsRuntimeCatalystFiltering]] takes a separate
path: all
+ * runtime filters are pushed as Catalyst expressions in a single call, with
no translation to
+ * connector predicates and no `filterAttributes` gating. The two paths are
mutually exclusive.
+ *
* @return true if any filters were pushed to the data source
*/
def pushRuntimeFilters(
@@ -183,6 +189,11 @@ object PushDownUtils extends Logging {
table: Table,
output: Seq[AttributeReference]): Boolean = {
scan match {
+ case _: SupportsRuntimeV2Filtering with SupportsRuntimeCatalystFiltering
=>
+ throw SparkException.internalError(
+ "A scan must not implement both SupportsRuntimeV2Filtering and " +
+ s"SupportsRuntimeCatalystFiltering, but ${scan.getClass.getName}
implements both.")
+
case filterableScan: SupportsRuntimeV2Filtering if
runtimeFilters.nonEmpty =>
// Push down translatable runtime filters.
val filtersToTranslated = runtimeFilters.flatMap { f =>
@@ -218,6 +229,41 @@ object PushDownUtils extends Logging {
}
translatedFiltersPushed || partPredicatesPushed
+
+ case catalystScan: SupportsRuntimeCatalystFiltering if
runtimeFilters.nonEmpty =>
+ // A runtime filter is normally evaluated twice: the source prunes
with it, and the
+ // FilterExec above the scan applies it again. The two have to agree,
so this screen
+ // pushes only predicates the source can be trusted to evaluate in
Spark's place.
+ //
+ // But pushing a non-deterministic one would let the source prune on
its own coin flip and
+ // Spark flip again for the rows that survive, so we handle that
specially.
+ //
+ // Not pushing a filter is safe only while its FilterExec still
evaluates it. A fully
+ // pushed filter has none, so the source is its only evaluator. That
is why
+ // DataSourceV2Strategy deletes a FilterExec at planning time only for
a filter we push
+ // below, and we use the same method (isPushablePartitionFilter) to
determine if we
+ // should push it.
+ //
+ // Note: today every filter reaching either site passes this check
+ // (deleted by DataSourceV2Strategy, and pushed by this method), since
runtimeFilters
+ // holds only deterministic filters (SPARK-58207) and
ExtractPythonUDFs has already
+ // lifted any Python UDF out of the post-scan filters. But sharing the
check keeps the two
+ // decisions consistent if non-deterministic filters reach the site.
+ //
+ // A DPP filter degrades to TrueLiteral once its subquery is pruned
away, so it matches
+ // every row. The V2 path above drops these implicitly, since
translateRuntimeFilterV2
+ // returns None; here we push Catalyst expressions directly, so we
remove them explicitly.
+ val catalystFilters = runtimeFilters
+ .filter(isPushablePartitionFilter(_, includeSubquery = true))
+ .flatMap(unwrapRuntimeFilterExpression)
+ .filterNot(_ == Literal.TrueLiteral)
+ if (catalystFilters.nonEmpty) {
+ catalystScan.filter(catalystFilters.toArray)
+ true
+ } else {
+ false
+ }
+
case _ =>
false
}
@@ -230,8 +276,8 @@ object PushDownUtils extends Logging {
* pre-filter partition set.
*
* Notes:
- * - Do not call multiple times for the same `scan` instance;
- * [[SupportsRuntimeV2Filtering.filter]] is mutating.
+ * - `filter` is mutating, and Spark may call this more than once for the
same `scan` instance
+ * (see [[pushRuntimeFilters]]); successive calls are additive.
* - When `outputPartitioning` is a [[KeyedPartitioning]], every split from
* `planInputPartitions()` used on this path must implement
[[HasPartitionKey]].
*
@@ -403,7 +449,7 @@ object PushDownUtils extends Logging {
val partitionAttributes = partitionFields.map(_.attrRef)
val (partFilters, nonPartitionFilters) =
DataSourceUtils.getPartitionFiltersAndDataFilters(partitionAttributes,
flattenedFilters)
- val (pushable, nonPushable) =
partFilters.partition(isPushablePartitionFilter)
+ val (pushable, nonPushable) =
partFilters.partition(isPushablePartitionFilter(_))
val (partitionPredicates, errorPartitionPredicates) =
pushable.partitionMap { e =>
PartitionPredicateImpl(e, partitionFields).toLeft(e)
}
@@ -438,21 +484,40 @@ object PushDownUtils extends Logging {
private[v2] def createRuntimePartitionPredicates(
runtimeFilters: Seq[Expression],
partitionFields: Seq[PartitionPredicateField]):
Seq[PartitionPredicateImpl] = {
- val catalystExprs = runtimeFilters.flatMap {
+ val catalystExprs = runtimeFilters.flatMap(unwrapRuntimeFilterExpression)
+ val flattened = flattenNestedPartitionFilters(catalystExprs,
partitionFields).keys
+ createPartitionPredicates(flattened.toSeq, partitionFields)._1
+ }
+
+ /** Unwraps a runtime filter to the Catalyst predicate for pushdown. */
+ private def unwrapRuntimeFilterExpression(rf: Expression):
Option[Expression] =
+ rf match {
case DynamicPruningExpression(in: InSubqueryExec) if
in.isResultUnavailable =>
None
case DynamicPruningExpression(e) => Some(e)
case _: DynamicPruning => None
case f => Some(f.transform { case s: ExecScalarSubquery => s.toLiteral })
}
- val flattened = flattenNestedPartitionFilters(catalystExprs,
partitionFields).keys
- createPartitionPredicates(flattened.toSeq, partitionFields)._1
- }
- private def isPushablePartitionFilter(f: Expression) =
+ /**
+ * Whether the data source can be trusted to evaluate `f` in place of Spark:
a non-deterministic
+ * expression would not give the same answer twice, and a Python UDF or a
subquery cannot be
+ * evaluated by the source at all. Spark will attempt to push only these
expressions to the
+ * data source.
+ *
+ * Spark will also call this before dropping a fully pushed runtime filter
from the
+ * post-scan filters, so that it never removes the only evaluator of a
filter that Spark
+ * refuses to push.
+ *
+ * @param includeSubquery whether a subquery in `f` should be tolerated.
Only safe for a runtime
+ * filter, which is pushed at execution time, after its subquery has been
evaluated.
+ */
+ private[v2] def isPushablePartitionFilter(
+ f: Expression,
+ includeSubquery: Boolean = false): Boolean =
f.deterministic &&
- !SubqueryExpression.hasSubquery(f) &&
- !f.exists(_.isInstanceOf[PythonUDF])
+ !f.exists(_.isInstanceOf[PythonUDF]) &&
+ (includeSubquery || !SubqueryExpression.hasSubquery(f))
/**
* Replaces all partition column references with canonical
[[AttributeReference]]s
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala
index 0e885b994ede..1bf12a695bc4 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala
@@ -28,6 +28,7 @@ import org.apache.spark.sql.execution.LogicalRDD
import org.apache.spark.sql.execution.columnar.InMemoryRelation
import org.apache.spark.sql.execution.datasources.{HadoopFsRelation,
LogicalRelation}
import org.apache.spark.sql.execution.datasources.v2.ExtractV2Scan
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
/**
* Dynamic partition pruning optimization is performed based on the type and
* selectivity of the join operation. During query optimization, we insert a
@@ -86,6 +87,14 @@ object PartitionPruning extends Rule[LogicalPlan] with
PredicateHelper with Join
} else {
None
}
+ case (resExp, r @ ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering))
=>
+ val filterAttrs = V2ExpressionUtils.resolveAttributeRefs(
+ scan.filterAttributes(), r.output)
+ if (resExp.references.subsetOf(filterAttrs)) {
+ Some(r)
+ } else {
+ None
+ }
case _ => None
}
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala
index 9f8409efa360..87139a3a20e1 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala
@@ -24,9 +24,11 @@ import
org.apache.spark.sql.catalyst.optimizer.RewritePredicateSubquery
import org.apache.spark.sql.catalyst.planning.{DeltaBasedRowLevelOperation,
GroupBasedRowLevelOperation}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter,
LogicalPlan, RowLevelWrite}
import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering
+import org.apache.spark.sql.connector.expressions.NamedReference
+import org.apache.spark.sql.connector.read.{Scan, SupportsRuntimeV2Filtering}
import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE,
MERGE, UPDATE}
import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Implicits,
DataSourceV2Relation, DataSourceV2ScanRelation, ExtractV2Scan}
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
import org.apache.spark.util.ArrayImplicits._
/**
@@ -51,26 +53,39 @@ class
RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla
override def apply(plan: LogicalPlan): LogicalPlan = plan transformDown {
case GroupBasedRowLevelOperation(replaceData, _, Some(cond),
- ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) if
canInjectGroupFilters(cond, scan) =>
- injectGroupFilters(replaceData, cond, scan)
+ ExtractV2Scan(scan: SupportsRuntimeV2Filtering))
+ if canInjectGroupFilters(cond, scan.filterAttributes) =>
+ injectGroupFilters(replaceData, cond, scan, scan.filterAttributes)
+
+ case GroupBasedRowLevelOperation(replaceData, _, Some(cond),
+ ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering))
+ if canInjectGroupFilters(cond, scan.filterAttributes()) =>
+ injectGroupFilters(replaceData, cond, scan, scan.filterAttributes())
+
+ case DeltaBasedRowLevelOperation(writeDelta, _, Some(cond),
+ ExtractV2Scan(scan: SupportsRuntimeV2Filtering))
+ if canInjectGroupFilters(cond, scan.filterAttributes) =>
+ injectGroupFilters(writeDelta, cond, scan, scan.filterAttributes)
case DeltaBasedRowLevelOperation(writeDelta, _, Some(cond),
- ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) if
canInjectGroupFilters(cond, scan) =>
- injectGroupFilters(writeDelta, cond, scan)
+ ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering))
+ if canInjectGroupFilters(cond, scan.filterAttributes()) =>
+ injectGroupFilters(writeDelta, cond, scan, scan.filterAttributes())
}
private def canInjectGroupFilters(
cond: Expression,
- scan: SupportsRuntimeV2Filtering): Boolean = {
+ filterAttrs: Array[NamedReference]): Boolean = {
conf.runtimeRowLevelOperationGroupFilterEnabled &&
cond != TrueLiteral &&
- scan.filterAttributes.nonEmpty
+ filterAttrs.nonEmpty
}
private def injectGroupFilters(
write: RowLevelWrite,
cond: Expression,
- scan: SupportsRuntimeV2Filtering): LogicalPlan = {
+ scan: Scan,
+ filterAttrs: Array[NamedReference]): LogicalPlan = {
// use reference equality on scan to find required scan relations
val newQuery = write.query transformUp {
case r: DataSourceV2ScanRelation if r.scan eq scan =>
@@ -79,9 +94,9 @@ class
RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla
val originalTable = r.relation.table.asRowLevelOperationTable.table
val relation = r.relation.copy(table = originalTable)
val matchingRowsPlan = buildMatchingRowsPlan(write, relation, cond)
- val filterAttrs = scan.filterAttributes.toImmutableArraySeq
- val buildKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrs,
matchingRowsPlan)
- val pruningKeys =
V2ExpressionUtils.resolveRefs[Attribute](filterAttrs, r)
+ val filterAttrsSeq = filterAttrs.toImmutableArraySeq
+ val buildKeys =
V2ExpressionUtils.resolveRefs[Attribute](filterAttrsSeq, matchingRowsPlan)
+ val pruningKeys =
V2ExpressionUtils.resolveRefs[Attribute](filterAttrsSeq, r)
Filter(buildDynamicPruningCond(matchingRowsPlan, buildKeys,
pruningKeys), r)
}
// optimize subqueries to rewrite them as joins and trigger job planning
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala
new file mode 100644
index 000000000000..63709e32ecd5
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala
@@ -0,0 +1,400 @@
+/*
+ * 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.{SparkConf, SparkException}
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.expressions.{Add, AttributeReference,
DynamicPruning, DynamicPruningExpression, EqualTo, Expression, GreaterThan,
Literal, RLike}
+import
org.apache.spark.sql.connector.catalog.{InMemoryCatalystRuntimeFilterTable,
InMemoryTableCatalystRuntimeFilterCatalog}
+import org.apache.spark.sql.connector.expressions.{FieldReference,
NamedReference}
+import org.apache.spark.sql.connector.expressions.filter.Predicate
+import org.apache.spark.sql.connector.read.{Scan, SupportsRuntimeV2Filtering}
+import org.apache.spark.sql.execution.{FilterExec, ScalarSubquery =>
ExecScalarSubquery}
+import org.apache.spark.sql.execution.ExplainUtils.stripAQEPlan
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec,
DataSourceV2ScanRelation, DataSourceV2Strategy}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{IntegerType, StringType, StructType}
+
+/**
+ * Tests for scans that implement
+ *
[[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]],
+ * where runtime filters are pushed once as Catalyst expressions instead of
connector
+ * predicates.
+ */
+class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession {
+
+ protected val v2Source = classOf[FakeV2ProviderWithCustomSchema].getName
+ protected val catalogName = "testcatalystruntimefilter"
+
+ override def sparkConf: SparkConf = super.sparkConf
+ .set(s"spark.sql.catalog.$catalogName",
+ classOf[InMemoryTableCatalystRuntimeFilterCatalog].getName)
+
+ private def withDPPConf(f: => Unit): Unit = {
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10")(f)
+ }
+
+ test("scalar subquery on partition column -> pushed as Catalyst expression")
{
+ val tbl = s"$catalogName.tbl1"
+ val dim = s"$catalogName.dim1"
+ withTable(tbl, dim) {
+ sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED
BY (part)")
+ for (i <- 0 until 5) {
+ sql(s"INSERT INTO $tbl VALUES ($i, $i)")
+ }
+ sql(s"CREATE TABLE $dim (val INT) USING $v2Source")
+ sql(s"INSERT INTO $dim VALUES (3)")
+
+ val df = sql(s"SELECT * FROM $tbl WHERE part = (SELECT max(val) FROM
$dim)")
+ checkAnswer(df, Row(3, 3))
+
+ assertScalarSubqueryRuntimeFilters(df)
+ val part = AttributeReference("part", IntegerType, nullable = false)()
+ assertPushedCatalystPredicatesEqual(df, EqualTo(part, Literal(3)))
+ // `part` is not declared fully pushed, so Spark still evaluates the
filter after the scan.
+ assertScalarSubqueryEvaluatedAfterScan(df, expected = true)
+ }
+ }
+
+ test("predicate on fully pushed filter attributes -> not evaluated after the
scan") {
+ val tbl = s"$catalogName.tbl_fully_pushed"
+ val dim = s"$catalogName.dim_fully_pushed"
+ withTable(tbl, dim) {
+ sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED
BY (part) " +
+ "TBLPROPERTIES('fully-pushed-filter-attributes' = 'part')")
+ // Matching and nonmatching partitions: the scan must prune nonmatching
ones itself
+ // because Spark drops the post-scan FilterExec for fully pushed
attributes.
+ for (i <- 0 until 5) {
+ sql(s"INSERT INTO $tbl VALUES ($i, $i)")
+ }
+ sql(s"CREATE TABLE $dim (val INT) USING $v2Source")
+ sql(s"INSERT INTO $dim VALUES (3)")
+
+ val df = sql(s"SELECT * FROM $tbl WHERE part = (SELECT max(val) FROM
$dim)")
+ checkAnswer(df, Row(3, 3))
+
+ assertScalarSubqueryRuntimeFilters(df)
+ val part = AttributeReference("part", IntegerType, nullable = false)()
+ assertPushedCatalystPredicatesEqual(df, EqualTo(part, Literal(3)))
+ assertScalarSubqueryEvaluatedAfterScan(df, expected = false)
+ }
+ }
+
+ test("non-deterministic predicate on fully pushed attributes -> evaluated
after the scan") {
+ val tbl = s"$catalogName.tbl_nondeterministic"
+ val dim = s"$catalogName.dim_nondeterministic"
+ withTable(tbl, dim) {
+ sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED
BY (part) " +
+ "TBLPROPERTIES('fully-pushed-filter-attributes' = 'part')")
+ for (i <- 0 until 5) {
+ sql(s"INSERT INTO $tbl VALUES ($i, $i)")
+ }
+ sql(s"CREATE TABLE $dim (val INT) USING $v2Source")
+ sql(s"INSERT INTO $dim VALUES (3)")
+
+ // A non-deterministic filter is never pushed, so it must keep its
post-scan FilterExec even
+ // though it only references a fully pushed attribute. Dropping it there
would leave nothing
+ // to evaluate it and the scan would return the nonmatching partitions
too.
+ val df = sql(
+ s"SELECT * FROM $tbl WHERE part = (SELECT max(val) FROM $dim) OR
rand() < 0.5")
+ // The row in the matching partition always qualifies, the others
qualify at random.
+ assert(df.collect().contains(Row(3, 3)))
+
+ assertScalarSubqueryEvaluatedAfterScan(df, expected = true)
+ assertScalarSubqueryRuntimeFilters(df, expectedCount = 0)
+ assertPushedCatalystPredicates(df, 0)
+ }
+ }
+
+ test("predicate on partly fully pushed filter attributes -> evaluated after
the scan") {
+ val tbl = s"$catalogName.tbl_partly_pushed"
+ val dim = s"$catalogName.dim_partly_pushed"
+ withTable(tbl, dim) {
+ sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source " +
+ "PARTITIONED BY (p1, p2) " +
+ "TBLPROPERTIES('fully-pushed-filter-attributes' = 'p1')")
+ for (i <- 0 until 5) {
+ sql(s"INSERT INTO $tbl VALUES ($i, 1, 2)")
+ }
+ sql(s"CREATE TABLE $dim (val INT) USING $v2Source")
+ sql(s"INSERT INTO $dim VALUES (3)")
+
+ // The predicate also references p2, which is not declared fully pushed,
so it is not
+ // considered fully pushed and Spark keeps evaluating it after the scan.
+ val df = sql(s"SELECT * FROM $tbl WHERE p1 + p2 = (SELECT max(val) FROM
$dim)")
+ checkAnswer(df, (0 until 5).map(i => Row(i, 1, 2)))
+
+ assertScalarSubqueryRuntimeFilters(df)
+ val p1 = AttributeReference("p1", IntegerType, nullable = false)()
+ val p2 = AttributeReference("p2", IntegerType, nullable = false)()
+ assertPushedCatalystPredicatesEqual(df, EqualTo(Add(p1, p2), Literal(3)))
+ assertScalarSubqueryEvaluatedAfterScan(df, expected = true)
+ }
+ }
+
+ test("arithmetic around a scalar subquery -> subquery literalized,
expression pushed intact") {
+ val tbl = s"$catalogName.tbl2"
+ val dim = s"$catalogName.dim2"
+ withTable(tbl, dim) {
+ sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED
BY (part)")
+ for (i <- 0 until 5) {
+ sql(s"INSERT INTO $tbl VALUES ($i, $i)")
+ }
+ sql(s"CREATE TABLE $dim (val INT) USING $v2Source")
+ sql(s"INSERT INTO $dim VALUES (2)")
+
+ // The scalar subquery is literalized on the way to the scan, and the
arithmetic around it
+ // survives. This one would also translate to a V2 predicate (`part >
3`, after the
+ // literalized operands are folded); see the next test for one that
would not.
+ val df = sql(s"SELECT * FROM $tbl WHERE part > (SELECT max(val) FROM
$dim) + 1")
+ checkAnswer(df, Row(4, 4))
+
+ assertScalarSubqueryRuntimeFilters(df)
+ val part = AttributeReference("part", IntegerType, nullable = false)()
+ assertPushedCatalystPredicatesEqual(
+ df, GreaterThan(part, Add(Literal(2), Literal(1))))
+ }
+ }
+
+ test("filter with no V2 translation -> pushed instead of dropped") {
+ val tbl = s"$catalogName.tbl_untranslatable"
+ val dim = s"$catalogName.dim_untranslatable"
+ withTable(tbl, dim) {
+ sql(s"CREATE TABLE $tbl (id INT, part STRING) USING $v2Source
PARTITIONED BY (part)")
+ for (i <- 0 until 5) {
+ sql(s"INSERT INTO $tbl VALUES ($i, '$i')")
+ }
+ sql(s"CREATE TABLE $dim (val STRING) USING $v2Source")
+ sql(s"INSERT INTO $dim VALUES ('3')")
+
+ // `V2ExpressionBuilder` has no RLike case, so this filter has no V2
predicate at all and
+ // the V2 interfaces would never see it. The pattern is still a subquery
when the optimizer
+ // runs, so nothing rewrites the RLIKE into a translatable Contains
beforehand either.
+ val df = sql(s"SELECT * FROM $tbl WHERE part RLIKE (SELECT max(val) FROM
$dim)")
+ checkAnswer(df, Row(3, "3"))
+
+ assertScalarSubqueryRuntimeFilters(df)
+ val runtimeFilter = collectBatchScan(df).runtimeFilters.head
+
assert(DataSourceV2Strategy.translateScalarSubqueryFilterV2(runtimeFilter).isEmpty,
+ s"Expected no V2 translation for $runtimeFilter")
+ val part = AttributeReference("part", StringType, nullable = false)()
+ assertPushedCatalystPredicatesEqual(df, RLike(part, Literal("3")))
+ }
+ }
+
+ test("DPP filter -> pushed as InSubqueryExec expression") {
+ val fact = s"$catalogName.fact3"
+ val dim = s"$catalogName.dim3"
+ withTable(fact, dim) {
+ sql(s"CREATE TABLE $fact (id INT, part INT) USING $v2Source PARTITIONED
BY (part)")
+ for (i <- 0 until 5) {
+ sql(s"INSERT INTO $fact VALUES ($i, $i)")
+ }
+ sql(s"CREATE TABLE $dim (dim_id INT, dim_val STRING) USING $v2Source")
+ sql(s"INSERT INTO $dim VALUES (2, 'two')")
+
+ withDPPConf {
+ val df = sql(
+ s"""SELECT f.id, f.part FROM $fact f JOIN $dim d
+ |ON f.part = d.dim_id WHERE d.dim_val = 'two'""".stripMargin)
+ checkAnswer(df, Row(2, 2))
+
+ assertDPPRuntimeFilters(df)
+ val dppPredicate = collectBatchScan(df).runtimeFilters.collectFirst {
+ case DynamicPruningExpression(e) => e
+ }.get
+ assertPushedCatalystPredicatesEqual(df, dppPredicate)
+ }
+ }
+ }
+
+ test("scan implementing both runtime filtering interfaces -> rejected") {
+ val tbl = s"$catalogName.tbl_both_interfaces"
+ withTable(tbl) {
+ sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED
BY (part)")
+ sql(s"INSERT INTO $tbl VALUES (1, 1)")
+
+ val scanRelation = sql(s"SELECT * FROM
$tbl").queryExecution.optimizedPlan.collectFirst {
+ case r: DataSourceV2ScanRelation => r
+ }.getOrElse(fail("Expected a DataSourceV2ScanRelation"))
+
+ // every runtime filtering path starts at runtimeFilterAttrs
+ val e = intercept[SparkException] {
+ scanRelation.copy(scan = new
BothRuntimeFilteringInterfacesScan).runtimeFilterAttrs
+ }
+ assert(e.getMessage.contains("A scan must not implement both
SupportsRuntimeV2Filtering " +
+ "and SupportsRuntimeCatalystFiltering"))
+ }
+ }
+
+ test("filter on column outside filterAttributes -> not pushed") {
+ val tbl = s"$catalogName.tbl4"
+ val dim = s"$catalogName.dim4"
+ withTable(tbl, dim) {
+ sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source " +
+ "PARTITIONED BY (p1, p2) " +
+ "TBLPROPERTIES('filter-attributes' = 'p1')")
+ for (i <- 0 until 5) {
+ sql(s"INSERT INTO $tbl VALUES ($i, $i, 10)")
+ }
+ sql(s"CREATE TABLE $dim (val INT) USING $v2Source")
+ sql(s"INSERT INTO $dim VALUES (10)")
+
+ // p2 is a partition column but is not declared filterable, so no
runtime filter is derived.
+ val df = sql(s"SELECT * FROM $tbl WHERE p2 = (SELECT max(val) FROM
$dim)")
+ checkAnswer(df, (0 until 5).map(i => Row(i, i, 10)))
+
+ assert(collectBatchScan(df).runtimeFilters.isEmpty,
+ "Expected no runtime filters for a column outside filterAttributes")
+ assertPushedCatalystPredicates(df, 0)
+ }
+ }
+
+ test("no runtime filter -> filter() is never called") {
+ val tbl = s"$catalogName.tbl5"
+ withTable(tbl) {
+ sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED
BY (part)")
+ for (i <- 0 until 5) {
+ sql(s"INSERT INTO $tbl VALUES ($i, $i)")
+ }
+
+ val df = sql(s"SELECT * FROM $tbl WHERE part = 3")
+ checkAnswer(df, Row(3, 3))
+
+ assert(collectBatchScan(df).runtimeFilters.isEmpty)
+ assertPushedCatalystPredicates(df, 0)
+ }
+ }
+
+ //
---------------------------------------------------------------------------
+ // Helper methods
+ //
---------------------------------------------------------------------------
+
+ private def assertDPPRuntimeFilters(
+ df: DataFrame, expectedCount: Int = 1): Unit = {
+ val batchScan = collectBatchScan(df)
+ val dppFilters = batchScan.runtimeFilters.collect {
+ case d: DynamicPruningExpression => d
+ }
+ assert(dppFilters.size === expectedCount,
+ s"Expected $expectedCount DynamicPruningExpression(s) " +
+ s"in runtimeFilters, got ${dppFilters.size}")
+ }
+
+ private def assertScalarSubqueryRuntimeFilters(
+ df: DataFrame, expectedCount: Int = 1): Unit = {
+ val batchScan = collectBatchScan(df)
+ val scalarFilters = batchScan.runtimeFilters.collect {
+ case f if !f.isInstanceOf[DynamicPruning] => f
+ }
+ val dppFilters = batchScan.runtimeFilters.collect {
+ case d: DynamicPruning => d
+ }
+ assert(scalarFilters.size === expectedCount,
+ s"Expected $expectedCount scalar subquery runtime filter(s), " +
+ s"got ${scalarFilters.size}")
+ assert(dppFilters.isEmpty,
+ "Expected non-DPP runtime filters (scalar subquery)")
+ }
+
+ /**
+ * Checks whether a scalar subquery runtime filter is still evaluated by a
[[FilterExec]] above
+ * the scan. Filters that only reference `fullyPushedFilterAttributes` are
dropped from it.
+ */
+ private def assertScalarSubqueryEvaluatedAfterScan(
+ df: DataFrame,
+ expected: Boolean): Unit = {
+ val postScanConditions =
stripAQEPlan(df.queryExecution.executedPlan).collect {
+ case f: FilterExec => f.condition
+ }
+ val evaluated =
postScanConditions.exists(_.exists(_.isInstanceOf[ExecScalarSubquery]))
+ assert(evaluated === expected,
+ s"Expected scalar subquery evaluated after scan to be $expected, " +
+ s"post-scan filter conditions: $postScanConditions")
+ }
+
+ private def collectBatchScan(df: DataFrame): BatchScanExec = {
+ stripAQEPlan(df.queryExecution.executedPlan).collectFirst {
+ case b: BatchScanExec => b
+ }.getOrElse(fail("Expected BatchScanExec in plan"))
+ }
+
+ private def getPushedCatalystPredicates(df: DataFrame): Seq[Expression] = {
+ collectBatchScan(df).scan match {
+ case s:
InMemoryCatalystRuntimeFilterTable#InMemoryCatalystRuntimeFilterBatchScan =>
+ s.pushedCatalystPredicates
+ case other =>
+ fail(s"Expected InMemoryCatalystRuntimeFilterBatchScan, got $other")
+ }
+ }
+
+ private def assertPushedCatalystPredicates(df: DataFrame, expected: Int):
Unit = {
+ val preds = getPushedCatalystPredicates(df)
+ assert(preds.size === expected,
+ s"Expected $expected pushed Catalyst runtime predicate(s), got
${preds.size}: $preds")
+ }
+
+ /**
+ * Binds [[AttributeReference]]s in `expected` to the scan output (by name)
and checks that the
+ * pushed Catalyst runtime predicates match exactly via
[[Expression.semanticEquals]].
+ */
+ private def assertPushedCatalystPredicatesEqual(
+ df: DataFrame,
+ expected: Expression*): Unit = {
+ val batchScan = collectBatchScan(df)
+ val actual = getPushedCatalystPredicates(df)
+ val normalizedExpected = expected.map(bindToScanOutput(_,
batchScan.output))
+ assert(actual.size === normalizedExpected.size,
+ s"Expected ${normalizedExpected.size} pushed Catalyst predicate(s), " +
+ s"got ${actual.size}: $actual")
+ actual.zip(normalizedExpected).foreach { case (a, e) =>
+ assert(a.semanticEquals(e),
+ s"Pushed Catalyst predicate mismatch.\nExpected: $e\nActual: $a")
+ }
+ }
+
+ private def bindToScanOutput(
+ expr: Expression,
+ output: Seq[AttributeReference]): Expression = {
+ val resolver = SQLConf.get.resolver
+ expr.transformUp {
+ case a: AttributeReference =>
+ output.find(o => resolver(o.name, a.name))
+ .map(_.withNullability(a.nullable).withQualifier(a.qualifier))
+ .getOrElse(a)
+ }
+ }
+}
+
+/** A scan violating the rule that only one runtime filtering interface may be
implemented. */
+private class BothRuntimeFilteringInterfacesScan
+ extends Scan with SupportsRuntimeV2Filtering with
SupportsRuntimeCatalystFiltering {
+
+ override def readSchema(): StructType = new StructType().add("part",
IntegerType)
+
+ override def filterAttributes(): Array[NamedReference] =
Array(FieldReference("part"))
+
+ override def filter(predicates: Array[Predicate]): Unit = {}
+
+ override def filter(expressions: Array[Expression]): Unit = {}
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala
new file mode 100644
index 000000000000..ce104f5631f1
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala
@@ -0,0 +1,50 @@
+/*
+ * 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.sql.Row
+
+class DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite
+ extends RowLevelOperationCatalystRuntimeFilterSuiteBase {
+
+ override protected def extraTableProps: java.util.Map[String, String] = {
+ val props = super.extraTableProps
+ props.put("supports-deltas", "true")
+ props
+ }
+
+ test("delete does not use group filtering when the group key is not
scanned") {
+ // the table is partitioned by dep, so hr and software are the two groups
+ createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING",
+ """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" }
+ |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" }
+ |""".stripMargin)
+
+ val executedPlan = executeAndKeepPlan {
+ sql(s"DELETE FROM $tableNameAsString WHERE salary IN (300, 400, 500)")
+ }
+ // a delta-based delete scans the row ID, the condition columns and the
metadata columns, so
+ // `dep` is not read and the scan cannot declare it as a filter attribute
+ assertNoCatalystGroupFilter(executedPlan)
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil)
+ }
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala
new file mode 100644
index 000000000000..34b43e2354ab
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala
@@ -0,0 +1,60 @@
+/*
+ * 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.sql.Row
+import org.apache.spark.sql.connector.catalog.InMemoryTable
+import org.apache.spark.sql.connector.write.DeleteSummary
+
+class GroupBasedRowLevelOperationCatalystRuntimeFilterSuite
+ extends RowLevelOperationCatalystRuntimeFilterSuiteBase {
+
+ test("delete runtime group filtering with SupportsRuntimeCatalystFiltering")
{
+ // the table is partitioned by dep, so hr and software are the two groups
+ createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING",
+ """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" }
+ |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" }
+ |""".stripMargin)
+
+ // only pk 1 matches, so hr is rewritten and its other row (pk 3) is
copied over
+ val executedPlan = executeAndKeepPlan {
+ sql(s"DELETE FROM $tableNameAsString WHERE salary IN (300, 400, 500)")
+ }
+ assertCatalystGroupFilter(
+ executedPlan,
+ expectedFilterAttrs = Seq("dep"),
+ expectedFilter = GroupFilter(scanSchema = "salary INT, dep STRING",
groups = Seq("hr")))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil)
+
+ checkReplacedPartitions(Seq("hr"))
+ checkDeleteMetrics(numDeletedRows = 1, numCopiedRows = 1)
+ }
+
+ private def checkDeleteMetrics(numDeletedRows: Long, numCopiedRows: Long):
Unit = {
+ val t = catalog.loadTable(ident).asInstanceOf[InMemoryTable]
+ val summary = t.commits.last.writeSummary.get.asInstanceOf[DeleteSummary]
+ assert(summary.numDeletedRows() === numDeletedRows,
+ s"Expected numDeletedRows=$numDeletedRows, got
${summary.numDeletedRows()}")
+ assert(summary.numCopiedRows() === numCopiedRows,
+ s"Expected numCopiedRows=$numCopiedRows, got ${summary.numCopiedRows()}")
+ }
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala
b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala
new file mode 100644
index 000000000000..810131eb54ce
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala
@@ -0,0 +1,229 @@
+/*
+ * 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.sql.Row
+import org.apache.spark.sql.catalyst.expressions.DynamicPruningExpression
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.connector.catalog.{BufferedRows,
InMemoryRowLevelOperationTable}
+import org.apache.spark.sql.execution.InSubqueryExec
+import org.apache.spark.sql.execution.ReusedSubqueryExec
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.unsafe.types.UTF8String
+
+/**
+ * Verifies that row-level runtime group filtering injects filters for scans
that implement
+ *
[[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]],
where the filter
+ * reaches the connector as a Catalyst expression instead of a connector
predicate.
+ *
+ * The tests here apply to both group-based and delta-based row-level
operations. DELETE is left
+ * to the concrete suites because delta-based deletes only scan the row ID and
the condition
+ * columns, so the group key is not available to filter on.
+ */
+abstract class RowLevelOperationCatalystRuntimeFilterSuiteBase
+ extends RowLevelOperationSuiteBase {
+
+ import testImplicits._
+
+ override protected def extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("use-catalyst-runtime-filtering", "true")
+ props
+ }
+
+ test("update runtime group filtering with SupportsRuntimeCatalystFiltering")
{
+ withTempView("updated_id") {
+ // the table is partitioned by dep, so hr and software are the two groups
+ createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING",
+ """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" }
+ |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" }
+ |""".stripMargin)
+
+ // the subquery blocks planning-time pushdown, leaving group filtering
to do the pruning;
+ // only id 1 matches, and it lives in hr
+ val updatedIdDF = Seq(Some(1), None).toDF()
+ updatedIdDF.createOrReplaceTempView("updated_id")
+
+ val executedPlan = executeAndKeepPlan {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT *
FROM updated_id)")
+ }
+ assertCatalystGroupFilter(
+ executedPlan,
+ expectedFilterAttrs = Seq("dep"),
+ expectedFilter = GroupFilter(scanSchema = "id INT, dep STRING", groups
= Seq("hr")))
+
+ // software was never read, so its rows must come back untouched
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 1, -1, "hr") :: Row(2, 2, 150, "software") :: Row(3, 3, 120,
"hr") :: Nil)
+ }
+ }
+
+ test("merge runtime group filtering with SupportsRuntimeCatalystFiltering") {
+ withTempView("source") {
+ createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING",
+ """{ "pk": 1, "id": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "salary": 200, "dep": "hr" }
+ |{ "pk": 3, "id": 3, "salary": 300, "dep": "hr" }
+ |{ "pk": 4, "id": 4, "salary": 400, "dep": "software" }
+ |{ "pk": 5, "id": 5, "salary": 500, "dep": "software" }
+ |""".stripMargin)
+
+ // pk 1 to 3 match rows in hr, pk 6 matches nothing and becomes an
insert, so hr is the
+ // only group that has to be rewritten
+ val sourceDF = Seq(1, 2, 3, 6).toDF("pk")
+ sourceDF.createOrReplaceTempView("source")
+
+ val executedPlan = executeAndKeepPlan {
+ sql(
+ s"""MERGE INTO $tableNameAsString t
+ |USING source s
+ |ON t.pk = s.pk
+ |WHEN MATCHED THEN
+ | UPDATE SET t.salary = t.salary + 1
+ |WHEN NOT MATCHED THEN
+ | INSERT (pk, id, salary, dep) VALUES (s.pk, 0, 0, 'hr')
+ |""".stripMargin)
+ }
+ assertCatalystGroupFilter(
+ executedPlan,
+ expectedFilterAttrs = Seq("dep"),
+ expectedFilter = GroupFilter(scanSchema = "pk INT, dep STRING", groups
= Seq("hr")))
+
+ // software was never read, so its rows must come back untouched
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Seq(
+ Row(1, 1, 101, "hr"),
+ Row(2, 2, 201, "hr"),
+ Row(3, 3, 301, "hr"),
+ Row(4, 4, 400, "software"),
+ Row(5, 5, 500, "software"),
+ Row(6, 0, 0, "hr")))
+ }
+ }
+
+ /**
+ * Asserts the injected group filter down to its contents: the scan declares
+ * `expectedFilterAttrs` in `filterAttributes`, every scan node carries one
dynamic pruning
+ * filter matching `expectedFilter`, the connector received that same filter
as a Catalyst
+ * expression, and the scan then read only `expectedFilter.groups`.
+ *
+ * A group-based UPDATE is rewritten as a union of two branches sharing one
scan, so the plan
+ * can hold more than one scan node. Each carries its own copy of the
filter, keyed on that
+ * branch's own attributes and with its own expr IDs, and pushes it
separately (see the UPDATE
+ * case in RowLevelOperationRuntimeGroupFiltering.buildMatchingRowsPlan).
Every copy is checked
+ * against `expectedFilter` rather than against one another, so the
expectation stays explicit
+ * and does not depend on how many copies the rewrite happens to produce.
+ */
+ protected def assertCatalystGroupFilter(
+ executedPlan: SparkPlan,
+ expectedFilterAttrs: Seq[String],
+ expectedFilter: GroupFilter): Unit = {
+ val batchScans = collect(executedPlan) { case s: BatchScanExec => s }
+ assert(batchScans.nonEmpty, "expected a batch scan for the row-level
operation")
+ val scan = catalystScan(batchScans.head)
+ assert(batchScans.forall(_.scan eq scan),
+ s"expected all ${batchScans.size} scan nodes to share one scan")
+
+ val filterAttrs =
scan.filterAttributes().map(_.fieldNames.mkString(".")).toSeq
+ assert(filterAttrs === expectedFilterAttrs,
+ s"expected the scan to declare $expectedFilterAttrs as filter
attributes, got $filterAttrs")
+
+ batchScans.foreach { batchScan =>
+ batchScan.runtimeFilters match {
+ case Seq(DynamicPruningExpression(inSubquery: InSubqueryExec)) =>
+ assertGroupFilter(inSubquery, expectedFilterAttrs, expectedFilter)
+ case other => fail(s"expected a single dynamic pruning group filter,
got $other")
+ }
+ }
+
+ // the scan must receive the Catalyst subquery Spark planned, not a
translated connector
+ // predicate, once per scan node
+ val pushed = scan.pushedCatalystPredicates
+ assert(pushed.size === batchScans.size,
+ s"expected each of the ${batchScans.size} scan node(s) to push the
filter once, got $pushed")
+ pushed.foreach {
+ case inSubquery: InSubqueryExec =>
+ assertGroupFilter(inSubquery, expectedFilterAttrs, expectedFilter)
+ case other =>
+ fail(s"expected the group filter pushed as an InSubqueryExec, got
$other")
+ }
+
+ val scannedGroups =
scan.data.map(_.asInstanceOf[BufferedRows].keyString()).distinct
+ assert(scannedGroups.sorted === expectedFilter.groups.sorted,
+ s"scan must read only the filtered groups, got
${scannedGroups.mkString(", ")}")
+ }
+
+ /**
+ * The expected shape of a group filter subquery: the columns it reads,
which must be only those
+ * needed to evaluate the row-level condition, and the groups it resolves to
at runtime.
+ */
+ protected case class GroupFilter(scanSchema: String, groups: Seq[String])
+
+ private def assertGroupFilter(
+ filter: InSubqueryExec,
+ expectedFilterAttrs: Seq[String],
+ expectedFilter: GroupFilter): Unit = {
+ assert(filter.child.references.toSeq.map(_.name) === expectedFilterAttrs,
+ s"expected the group filter keyed on $expectedFilterAttrs, got
${filter.child}")
+
+ // the second branch of a group-based UPDATE reuses the first branch's
subquery, and
+ // ReusedSubqueryExec is a leaf node, so unwrap it to reach the plan
underneath
+ val subqueryPlan = filter.plan match {
+ case reused: ReusedSubqueryExec => reused.child
+ case plan => plan
+ }
+ val subqueryScan = find(subqueryPlan) { case _: BatchScanExec => true;
case _ => false }
+ .getOrElse(fail(s"could not find the scan of group filter subquery
${filter.plan.name}"))
+ assert(
+ DataTypeUtils.sameType(subqueryScan.schema,
StructType.fromDDL(expectedFilter.scanSchema)),
+ s"unexpected group filter subquery scan schema
${subqueryScan.schema.sql}")
+
+ val groups = filter.values()
+ .getOrElse(fail("group filter subquery produced no values"))
+ .map(_.asInstanceOf[UTF8String].toString)
+ assert(groups.toSeq.sorted === expectedFilter.groups.sorted,
+ s"group filter must select the groups holding matching rows, got
${groups.mkString(", ")}")
+ }
+
+ /** Asserts no group filter was injected, e.g. because the scan does not
read the group key. */
+ protected def assertNoCatalystGroupFilter(executedPlan: SparkPlan): Unit = {
+ val batchScan = collect(executedPlan) { case s: BatchScanExec => s }.head
+ val scan = catalystScan(batchScan)
+ assert(scan.filterAttributes().isEmpty,
+ s"expected no filter attributes, got
${scan.filterAttributes().mkString(", ")}")
+ assert(batchScan.runtimeFilters.isEmpty,
+ s"expected no runtime filters, got ${batchScan.runtimeFilters}")
+ assert(scan.pushedCatalystPredicates.isEmpty,
+ s"expected no pushed predicates, got ${scan.pushedCatalystPredicates}")
+ }
+
+ private type CatalystRowLevelScan =
+ InMemoryRowLevelOperationTable#InMemoryCatalystRowLevelBatchScan
+
+ private def catalystScan(batchScan: BatchScanExec): CatalystRowLevelScan = {
+ batchScan.scan match {
+ case s: InMemoryRowLevelOperationTable#InMemoryCatalystRowLevelBatchScan
=> s
+ case other => fail(s"expected InMemoryCatalystRowLevelBatchScan, got
${other.getClass}")
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
