From 90f35f3f5d8ea6bbc0601f4ef8c6e1bd58035527 Mon Sep 17 00:00:00 2001 From: Szehon Ho Date: Thu, 27 Aug 2026 11:35:20 -0700 Subject: [PATCH] [SPARK-59012][SQL] Improve SupportsRuntimeCatalystFiltering test coverage and reject nested filter attributes ### What changes were proposed in this pull request? `SupportsRuntimeCatalystFiltering` was added by SPARK-58523, which covered row-level operations (`RowLevelOperationCatalystRuntimeFilterSuiteBase` and its group- and delta-based subclasses) and a core set of pushdown tests. It was not exercised in the two other places the predicate-based runtime filtering interfaces are: dynamic partition pruning and storage-partitioned joins. This PR closes that gap by reusing the existing suites rather than cloning them, and hardens the `filterAttributes()` contract. **Behavior change** - `DataSourceV2ScanRelation.runtimeFilterAttrs` (and `fullyPushedRuntimeFilterAttrs`) now reject a nested (multi-part) runtime-filter reference, including one over a struct column. Both runtime filtering interfaces require each `filterAttributes()` reference to be a top-level read-schema attribute; a nested reference such as `s.tz` previously resolved and widened to the enclosing struct column `s`, making runtime filters over every field of `s` eligible. **Test fixtures** - `InMemoryBaseTable`: `CatalystRuntimeFilteringScan` now prunes on nested partition keys. It previously looked up a partition attribute by joining `fieldNames` into a single top-level name, so a `GetStructField` chain such as `s.part` never matched and no pruning happened. Partition attributes now keep their `fieldNames` as name-part sequences, and `partitionAttrFor` matches the expression's path component-wise, so a quoted top-level column `` `a.b` `` stays distinct from a nested `a`.`b`. - `InMemoryCatalystRuntimeFilterTable`: threads the full table-creation metadata (constraints, distribution, ordering, partition counts, advisory size, strictness, and `numRowsPerSplit`) through to `InMemoryBaseTable`, and derives `filterAttributes()` / `fullyPushedFilterAttributes()` from a shared helper. - `InMemoryTableCatalystRuntimeFilterCatalog`: adds `InMemoryCatalystRuntimeFilterCatalog`, the `InMemoryCatalog` counterpart, so the Catalyst fixture can be used where functions and procedures are needed. - `InMemoryTableWithV2Filter`: threads the same table-creation metadata through to `InMemoryBaseTable`. **Reused suites** - `DynamicPartitionPruningSuite`: adds `DynamicPartitionPruningV2CatalystFilterSuiteAEOff` / `AEOn`, mirroring the existing `DynamicPartitionPruningV2FilterSuite` pair. - `KeyGroupedPartitioningSuite`: the shared SPJ fixtures move to a new `KeyGroupedPartitioningSuiteBase`, and the three runtime-filtering tests move to a `KeyGroupedPartitioningRuntimeFilterTests` trait. Two small suites then run those three tests once per interface: `KeyGroupedPartitioningRuntimeFilterSuite` (predicate-based, the default `InMemoryCatalog`) and `KeyGroupedPartitioningCatalystRuntimeFilterSuite`. - `DistributionAndOrderingSuiteBase`: `catalogClassName` becomes overridable so a subclass can vary the catalog behind `testcat`. Note for reviewers: `SPARK-42038: partially clustered: with dynamic partition filtering` and `SPARK-45652: SPJ should handle empty partition after dynamic filtering` are unchanged, but now report under `KeyGroupedPartitioningRuntimeFilterSuite` instead of `KeyGroupedPartitioningSuite`. Most of the diff in that file is this movement. **New tests in `DataSourceV2CatalystRuntimeFilterSuite`** Covering behavior specific to pushing Catalyst expressions: a DPP filter on a nested partition field arriving with the nested access intact, multiple predicates pushed in a single `filter()` call, a filter with no V2 translation being pushed instead of dropped, a scan implementing both runtime filtering interfaces being rejected, two partition columns whose dotted names collide (a quoted top-level `` `x.y` `` and a nested `x`.`y`) binding to the correct partition slot, and the `filterAttributes()` contract when a reported attribute is not a top-level scan attribute -- a missing attribute and a nested reference (over both an int and a struct column) are all rejected. One test calls `PushDownUtils.replanWithRuntimeFilters` directly, to reach the SPJ partitioning-preservation checks that a well-behaved source cannot trigger: dropping `HasPartitionKey`, reporting a partition key that was not in the original partitioning, or growing a key's split count. ### Why are the changes needed? `SupportsRuntimeCatalystFiltering` is the path a scan takes when runtime filters are pushed as Catalyst expressions instead of connector predicates. Dynamic partition pruning and storage-partitioned joins are the two features that produce those filters, and neither was tested against this interface, so regressions in the Catalyst path would not have been caught by the suites that cover the equivalent predicate-based path. ### Does this PR introduce _any_ user-facing change? No end-user-facing behavior change. For connector authors, a scan that declares a nested runtime-filter attribute -- a violation of the `filterAttributes()` contract -- is now rejected with an internal error instead of silently widening to the enclosing top-level column. ### How was this patch tested? Existing and new unit tests. Locally, on the rebased branch: - `DataSourceV2CatalystRuntimeFilterSuite`: 16 tests passed, including the new dotted/nested collision test and the nested-reference rejection. - `DynamicPartitionPruningV2CatalystFilterSuiteAEOff` / `AEOn` and `DynamicPartitionPruningV2FilterSuiteAEOff` / `AEOn`: 154 tests passed. - `KeyGroupedPartitioningSuite`: 96 tests passed. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Cursor with Claude Opus 5 Closes #58296 from szehon-ho/catalyst-runtime-filter-coverage. Authored-by: Szehon Ho Signed-off-by: Szehon Ho (cherry picked from commit d7e22ce8f28ac61634b0535201c15dc8f3e7c6ed) Signed-off-by: Szehon Ho --- .../datasources/v2/DataSourceV2Relation.scala | 18 +- .../SupportsRuntimeCatalystFiltering.scala | 8 +- .../connector/catalog/InMemoryBaseTable.scala | 71 +- .../InMemoryCatalystRuntimeFilterTable.scala | 55 +- .../catalog/InMemoryTableCatalog.scala | 52 +- ...oryTableCatalystRuntimeFilterCatalog.scala | 61 +- .../catalog/InMemoryTableWithV2Filter.scala | 17 +- .../InMemoryTableWithV2FilterCatalog.scala | 43 +- .../sql/DynamicPartitionPruningSuite.scala | 28 +- ...taSourceV2CatalystRuntimeFilterSuite.scala | 342 +++++++++- .../DistributionAndOrderingSuiteBase.scala | 5 +- .../KeyGroupedPartitioningSuite.scala | 627 ++++++++++-------- 12 files changed, 912 insertions(+), 415 deletions(-) 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 1c7a7abe4a8ad..1cbef86b96d36 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 @@ -207,8 +207,7 @@ case class DataSourceV2ScanRelation( case s: SupportsRuntimeCatalystFiltering => s.filterAttributes() case _ => Array.empty[NamedReference] } - AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( - filterAttrs.toImmutableArraySeq, this)) + resolveTopLevelFilterAttrs(filterAttrs) } /** @@ -221,6 +220,21 @@ case class DataSourceV2ScanRelation( case s: SupportsRuntimeCatalystFiltering => s.fullyPushedFilterAttributes() case _ => Array.empty[NamedReference] } + resolveTopLevelFilterAttrs(filterAttrs) + } + + /** + * Resolves the given runtime-filter references against this relation's output. Both runtime + * filtering interfaces require each reference to be a top-level attribute of the read schema, + * so a nested reference is rejected. + */ + private def resolveTopLevelFilterAttrs(filterAttrs: Array[NamedReference]): AttributeSet = { + filterAttrs.find(_.fieldNames.length > 1).foreach { ref => + throw SparkException.internalError( + s"Runtime filter attribute '${ref.fieldNames.mkString(".")}' declared by " + + s"${scan.getClass.getName} must be a top-level attribute of the scan read schema, " + + "but it is a nested reference.") + } AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( filterAttrs.toImmutableArraySeq, this)) } 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 index 5e5900a211961..eaa4857f1abe3 100644 --- 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 @@ -42,8 +42,8 @@ trait SupportsRuntimeCatalystFiltering extends Scan { * * 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. + * references are rejected, and attributes pruned out of the read schema fail to resolve, when + * Spark builds the scan relation. */ def filterAttributes(): Array[NamedReference] @@ -64,8 +64,8 @@ trait SupportsRuntimeCatalystFiltering extends Scan { * 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. + * references are rejected, and attributes pruned out of the read schema fail to resolve, when + * Spark builds the scan relation. */ def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty 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 1b5d7bb75fe66..d9540b8fff854 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 @@ -29,7 +29,7 @@ import scala.collection.mutable.{ArrayBuffer, ListBuffer} import scala.jdk.CollectionConverters._ import org.apache.spark.sql.catalyst.InternalRow -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.expressions.{AttributeReference, BindReferences, Cast, EvalMode, Expression => CatalystExpression, GenericInternalRow, GetStructField, 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} @@ -718,22 +718,24 @@ abstract class InMemoryBaseTable( protected def tableSchema: StructType private val catalystPredicates = ArrayBuffer.empty[CatalystExpression] + private var filterCalls = 0 override def filter(expressions: Array[CatalystExpression]): Unit = { catalystPredicates ++= expressions + filterCalls += 1 val partAttrs = partitionAttributes if (partAttrs.isEmpty) return + val partAttrRefs = partAttrs.map(_._2) - 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) + // Top down, so `s.part` is rewritten before its `s` child is considered. + val remapped = expr.transformDown { + case e => partitionAttrFor(e, partAttrs).getOrElse(e) } // 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) + if (remapped.references.forall(r => partAttrRefs.exists(_.exprId == r.exprId))) { + val bound = BindReferences.bindReference(remapped, partAttrRefs) val pred = CatalystPredicate.createInterpreted(bound) self.data = self.data.filter { p => try { @@ -755,15 +757,60 @@ abstract class InMemoryBaseTable( /** 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] = { + def filterCallCount: Int = filterCalls + + /** + * The `AttributeReference`s standing for the partition key InternalRow fields, in its field + * order, each paired with the name-part sequence of its partition column. The parts are kept + * unflattened so a quoted top-level column `a.b` (parts `Seq("a.b")`) stays distinct from a + * nested column `a`.`b` (parts `Seq("a", "b")`). Example: + * - `PARTITIONED BY (part, s.nested)` -> `(Seq("part"), AttributeReference(part))`, then + * `(Seq("s", "nested"), AttributeReference(s.nested))` + */ + private def partitionAttributes: Seq[(Seq[String], 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)() + val path = ref.fieldNames.toImmutableArraySeq + readSchema.findNestedField(path).orElse(tableSchema.findNestedField(path)).map { + case (_, f) => + path -> AttributeReference(ref.fieldNames.mkString("."), f.dataType, f.nullable)() } }.toSeq } + + /** + * The partition key `AttributeReference` that `e` reads, or None if `e` reads no partition + * column. The path `e` reads is compared to each partition column's name parts component-wise + * with the resolver, so a quoted top-level column `a.b` cannot collide with a nested column + * `a`.`b`. Examples, under `PARTITIONED BY (part, s.nested)` where `nested` is field 0 of `s`: + * - `AttributeReference(part)` -> `AttributeReference(part)` + * - `GetStructField(AttributeReference(s), 0)` -> `AttributeReference(s.nested)` + * - `AttributeReference(s)` -> None if `s` itself is not a partition column, only `s.nested` + */ + private def partitionAttrFor( + e: CatalystExpression, + partAttrs: Seq[(Seq[String], AttributeReference)]): Option[AttributeReference] = { + val resolver = SQLConf.get.resolver + partitionKeyPath(e).flatMap { path => + partAttrs.collectFirst { + case (parts, attr) if parts.length == path.length && + parts.lazyZip(path).forall((part, name) => resolver(part, name)) => attr + } + } + } + + /** + * The name parts `e` reads, or None if it reads neither a column nor a struct field. Each + * `GetStructField` ordinal is the field's position in its parent struct. Examples: + * - `AttributeReference(a)` -> `Seq("a")`, the top level column a + * - `GetStructField(AttributeReference(a), 0)` -> `Seq("a", "b")`, the nested column a.b + * - `GetStructField(GetStructField(AttributeReference(a), 0), 0)` -> `Seq("a", "b", "c")` + */ + private def partitionKeyPath(e: CatalystExpression): Option[Seq[String]] = e match { + case a: AttributeReference => Some(Seq(a.name)) + case g: GetStructField => + partitionKeyPath(g.child).map(parent => parent :+ g.childSchema(g.ordinal).name) + case _ => None + } } case class InMemoryBatchScan( 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 index 0f32028f20a94..369e8123f1d7e 100644 --- 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 @@ -21,7 +21,9 @@ import java.util import InMemoryCatalystRuntimeFilterTable._ -import org.apache.spark.sql.connector.expressions.{NamedReference, Transform} +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, NamedReference, SortOrder, 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 @@ -41,8 +43,17 @@ class InMemoryCatalystRuntimeFilterTable( name: String, columns: Array[Column], partitioning: Array[Transform], - properties: util.Map[String, String]) - extends InMemoryTableWithV2Filter(name, columns, partitioning, properties) { + properties: util.Map[String, String], + constraints: Array[Constraint] = Array.empty, + distribution: Distribution = Distributions.unspecified(), + ordering: Array[SortOrder] = Array.empty, + numPartitions: Option[Int] = None, + advisoryPartitionSize: Option[Long] = None, + isDistributionStrictlyRequired: Boolean = true, + numRowsPerSplit: Int = Int.MaxValue) + extends InMemoryTableWithV2Filter(name, columns, partitioning, properties, constraints, + distribution, ordering, numPartitions, advisoryPartitionSize, isDistributionStrictlyRequired, + numRowsPerSplit) { override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { new InMemoryCatalystRuntimeFilterScanBuilder(schema, options) @@ -74,25 +85,35 @@ class InMemoryCatalystRuntimeFilterTable( Option(InMemoryCatalystRuntimeFilterTable.this.properties.get(FilterAttributesKey)) .map(_.split(",").map(_.trim).toSet) - override def filterAttributes(): Array[NamedReference] = { + private val fullyPushedFilterAttrs: Set[String] = Option( + InMemoryCatalystRuntimeFilterTable.this.properties.get(FullyPushedFilterAttributesKey)) + .map(_.split(",").map(_.trim).toSet) + .getOrElse(Set.empty) + + /** + * The partition columns, each named by the top level read schema column it lives under, the + * form both interface methods require. Columns pruned out of the read schema are dropped, + * since neither method may name one. Examples: + * - `PARTITIONED BY (part)` -> `"part"` + * - `PARTITIONED BY (s.nested)` -> `"s"`, the struct column holding the partition field + */ + private def partitionAttrNames: Array[String] = { val scanFields = readSchema.fields.map(_.name).toSet - partitioning.flatMap(_.references()).filter { ref => - val name = ref.fieldNames.mkString(".") - scanFields.contains(name) && - restrictedFilterAttrs.forall(_.contains(name)) - } + partitioning.flatMap(_.references()).map(_.fieldNames.head).distinct + .filter(scanFields.contains) } - 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(".")) - } + override def filterAttributes(): Array[NamedReference] = { + partitionAttrNames + .filter(name => restrictedFilterAttrs.forall(_.contains(name))) + .map(FieldReference.column) } + // Not intersected with `filterAttributes()`, so a table can declare a fully pushed attribute + // that is not a filter attribute, a combination the interface forbids. + override def fullyPushedFilterAttributes(): Array[NamedReference] = { + partitionAttrNames.filter(fullyPushedFilterAttrs.contains).map(FieldReference.column) + } } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala index 96f18e8b901c7..cc5c75dd779be 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala @@ -193,14 +193,39 @@ class BasicInMemoryTableCatalog extends TableCatalog { InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties) val tableName = s"$name.${ident.quoted}" - val table = new InMemoryTable(tableName, columns, partitions, properties, constraints, - distribution, ordering, requiredNumPartitions, advisoryPartitionSize, - distributionStrictlyRequired, numRowsPerSplit) + val table = newInMemoryTable( + tableName, columns, partitions, properties, constraints, distribution, ordering, + requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired, numRowsPerSplit, + util.UUID.randomUUID().toString) tables.put(ident, table) namespaces.putIfAbsent(ident.namespace.toList, Map()) table } + /** + * Builds the in-memory table this catalog serves. Subclasses that expose a specialized table + * type must override this so both CREATE and ALTER reconstruct the same class. + */ + // scalastyle:off argcount + protected def newInMemoryTable( + name: String, + columns: Array[Column], + partitioning: Array[Transform], + properties: util.Map[String, String], + constraints: Array[Constraint], + distribution: Distribution, + ordering: Array[SortOrder], + requiredNumPartitions: Option[Int], + advisoryPartitionSize: Option[Long], + distributionStrictlyRequired: Boolean, + numRowsPerSplit: Int, + id: String): InMemoryBaseTable = { + // scalastyle:on argcount + new InMemoryTable(name, columns, partitioning, properties, constraints, distribution, + ordering, requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired, + numRowsPerSplit, id) + } + override def alterTable(ident: Identifier, changes: TableChange*): Table = { val table = loadTable(ident).asInstanceOf[InMemoryBaseTable] val properties = CatalogV2Util.applyPropertiesChanges(table.properties, changes) @@ -238,22 +263,13 @@ class BasicInMemoryTableCatalog extends TableCatalog { val currentVersion = table.version() val columnsWithIds = InMemoryBaseTable.assignMissingIds( CatalogV2Util.structTypeToV2Columns(schema)) + val reconstructedId = Option(table.id()).getOrElse(util.UUID.randomUUID().toString) val newTable = table match { - case _: InMemoryTable => - new InMemoryTable( - name = table.name, - columns = columnsWithIds, - partitioning = finalPartitioning, - properties = properties, - constraints = constraints, - id = table.id) - .alterTableWithData(table.data, schemaAfterDrops) - case _: InMemoryTableWithV2Filter => - new InMemoryTableWithV2Filter( - name = table.name, - columns = columnsWithIds, - partitioning = finalPartitioning, - properties = properties) + case _: InMemoryTable | _: InMemoryTableWithV2Filter => + newInMemoryTable( + table.name, columnsWithIds, finalPartitioning, properties, constraints, + table.distribution, table.ordering, table.numPartitions, table.advisoryPartitionSize, + table.isDistributionStrictlyRequired, table.numRowsPerSplit, reconstructedId) .alterTableWithData(table.data, schemaAfterDrops) case other => throw new UnsupportedOperationException( 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 index b8415eae3e15b..e699f439223f9 100644 --- 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 @@ -19,32 +19,43 @@ 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 +import org.apache.spark.sql.connector.catalog.constraints.Constraint +import org.apache.spark.sql.connector.distributions.Distribution +import org.apache.spark.sql.connector.expressions.{SortOrder, Transform} -class InMemoryTableCatalystRuntimeFilterCatalog extends InMemoryTableCatalog { - import CatalogV2Implicits._ - - override def createTable( - ident: Identifier, +/** + * Mix-in that constructs [[InMemoryCatalystRuntimeFilterTable]] from the shared in-memory + * catalog factory used by both CREATE TABLE and ALTER TABLE. + */ +trait InMemoryCatalystRuntimeFilterTableFactory { self: BasicInMemoryTableCatalog => + // scalastyle:off argcount + override protected def newInMemoryTable( + name: String, 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) + partitioning: Array[Transform], + properties: util.Map[String, String], + constraints: Array[Constraint], + distribution: Distribution, + ordering: Array[SortOrder], + requiredNumPartitions: Option[Int], + advisoryPartitionSize: Option[Long], + distributionStrictlyRequired: Boolean, + numRowsPerSplit: Int, + id: String): InMemoryBaseTable = { + // scalastyle:on argcount + new InMemoryCatalystRuntimeFilterTable( + name, columns, partitioning, properties, constraints, distribution, ordering, + requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired, numRowsPerSplit) } } + +class InMemoryTableCatalystRuntimeFilterCatalog extends InMemoryTableCatalog + with InMemoryCatalystRuntimeFilterTableFactory + +/** + * The [[InMemoryCatalog]] counterpart of [[InMemoryTableCatalystRuntimeFilterCatalog]]: it hands + * out tables whose scans take runtime filters as Catalyst expressions, and honors + * `numRowsPerSplit` so that a partition key can have several splits. + */ +class InMemoryCatalystRuntimeFilterCatalog extends InMemoryCatalog + with InMemoryCatalystRuntimeFilterTableFactory diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala index e9d73d0f9fe1e..b59a81cc29092 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala @@ -21,7 +21,9 @@ import java.util import org.scalatest.Assertions.assert -import org.apache.spark.sql.connector.expressions.{FieldReference, LiteralValue, NamedReference, Transform} +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, LiteralValue, NamedReference, SortOrder, Transform} import org.apache.spark.sql.connector.expressions.filter.{And, Predicate} import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder, SupportsRuntimeV2Filtering} import org.apache.spark.sql.connector.write.{LogicalWriteInfo, SupportsOverwriteV2, WriteBuilder, WriterCommitMessage} @@ -33,8 +35,17 @@ class InMemoryTableWithV2Filter( name: String, columns: Array[Column], partitioning: Array[Transform], - properties: util.Map[String, String]) - extends InMemoryBaseTable(name, columns, partitioning, properties) with SupportsDeleteV2 { + properties: util.Map[String, String], + constraints: Array[Constraint] = Array.empty, + distribution: Distribution = Distributions.unspecified(), + ordering: Array[SortOrder] = Array.empty, + numPartitions: Option[Int] = None, + advisoryPartitionSize: Option[Long] = None, + isDistributionStrictlyRequired: Boolean = true, + numRowsPerSplit: Int = Int.MaxValue) + extends InMemoryBaseTable(name, columns, partitioning, properties, constraints, distribution, + ordering, numPartitions, advisoryPartitionSize, isDistributionStrictlyRequired, + numRowsPerSplit) with SupportsDeleteV2 { override def canDeleteWhere(predicates: Array[Predicate]): Boolean = { InMemoryTableWithV2Filter.supportsPredicates(predicates) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2FilterCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2FilterCatalog.scala index ef2f5e26f0029..be5a907f52125 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2FilterCatalog.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2FilterCatalog.scala @@ -19,31 +19,28 @@ 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 +import org.apache.spark.sql.connector.catalog.constraints.Constraint +import org.apache.spark.sql.connector.distributions.Distribution +import org.apache.spark.sql.connector.expressions.{SortOrder, Transform} class InMemoryTableWithV2FilterCatalog extends InMemoryTableCatalog { - import CatalogV2Implicits._ - - override def createTable( - ident: Identifier, + // scalastyle:off argcount + override protected def newInMemoryTable( + name: String, 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 InMemoryTableWithV2Filter(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) + partitioning: Array[Transform], + properties: util.Map[String, String], + constraints: Array[Constraint], + distribution: Distribution, + ordering: Array[SortOrder], + requiredNumPartitions: Option[Int], + advisoryPartitionSize: Option[Long], + distributionStrictlyRequired: Boolean, + numRowsPerSplit: Int, + id: String): InMemoryBaseTable = { + // scalastyle:on argcount + new InMemoryTableWithV2Filter( + name, columns, partitioning, properties, constraints, distribution, ordering, + requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired, numRowsPerSplit) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala index 77754ef5952bc..9f497680b7fdb 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.optimizer.BuildRight import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, Inner, LeftAnti} import org.apache.spark.sql.catalyst.plans.logical.LocalRelation -import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} +import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableCatalystRuntimeFilterCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ import org.apache.spark.sql.execution.datasources.v2.BatchScanExec @@ -2690,6 +2690,32 @@ class DynamicPartitionPruningV2FilterSuiteAEOn extends DynamicPartitionPruningV2FilterSuite with EnableAdaptiveExecutionSuite +/** + * Runs the DSv2 dynamic partition pruning tests against scans that receive runtime filters as + * Catalyst expressions, via + * [[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]], rather than as + * connector predicates. This is the counterpart of [[DynamicPartitionPruningV2FilterSuite]], + * which covers the same tests for + * [[org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering]]. + */ +abstract class DynamicPartitionPruningV2CatalystFilterSuite + extends DynamicPartitionPruningV2Suite { + + override protected def initState(): Unit = { + super.initState() + spark.conf.set("spark.sql.catalog.testcat", + classOf[InMemoryTableCatalystRuntimeFilterCatalog].getName) + } +} + +class DynamicPartitionPruningV2CatalystFilterSuiteAEOff + extends DynamicPartitionPruningV2CatalystFilterSuite + with DisableAdaptiveExecutionSuite + +class DynamicPartitionPruningV2CatalystFilterSuiteAEOn + extends DynamicPartitionPruningV2CatalystFilterSuite + with EnableAdaptiveExecutionSuite + private object DppMaterializedInputTestState { private val counters = TrieMap.empty[String, AtomicInteger] 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 index 63709e32ecd5b..c0a1faa032635 100644 --- 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 @@ -18,15 +18,24 @@ 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.{AnalysisException, DataFrame, Row} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Add, AttributeReference, DynamicPruning, DynamicPruningExpression, EqualTo, Expression, GetStructField, GreaterThan, Literal, RLike} +import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning +import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper +import org.apache.spark.sql.connector.catalog.{ + Column, + Identifier, + InMemoryCatalystRuntimeFilterTable, + InMemoryTable, + InMemoryTableCatalystRuntimeFilterCatalog, + TableCatalog} +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference, Transform} import org.apache.spark.sql.connector.expressions.filter.Predicate -import org.apache.spark.sql.connector.read.{Scan, SupportsRuntimeV2Filtering} +import org.apache.spark.sql.connector.read.{Batch, HasPartitionKey, InputPartition, PartitionReaderFactory, 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.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, DataSourceV2Strategy, PushDownUtils} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering import org.apache.spark.sql.test.SharedSparkSession @@ -73,6 +82,13 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { 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) + + // The answer alone would be right even without pruning, since that post-scan filter drops + // the extra rows. + val batchScan = collectBatchScan(df) + assert(batchScan.inputPartitions.size === 5) + assert(batchScan.filteredPartitions.flatten.size === 1, + s"expected 1 partition after pruning, got ${batchScan.filteredPartitions.flatten.size}") } } @@ -228,6 +244,40 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { } } + test("DPP filter on a nested partition field -> pushed with the nested access intact") { + val fact = s"$catalogName.fact_nested_dpp" + val dim = s"$catalogName.dim_nested_dpp" + withTable(fact, dim) { + sql(s"CREATE TABLE $fact (id INT, s STRUCT) USING $v2Source " + + "PARTITIONED BY (s.part)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $fact VALUES ($i, named_struct('part', $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 FROM $fact f JOIN $dim d + |ON f.s.part = d.dim_id WHERE d.dim_val = 'two'""".stripMargin) + checkAnswer(df, Row(2)) + + assertDPPRuntimeFilters(df) + // The scan only reports `s`, the struct holding the partition field, but the predicate it + // receives keeps the nested access, so it can still tell which partition to keep. + val pushed = getPushedCatalystPredicates(df) + assert(pushed.size === 1, s"expected a single pushed predicate, got $pushed") + assert(pushed.head.exists(_.isInstanceOf[GetStructField]), + s"expected the pushed predicate to keep the nested access, got ${pushed.head}") + + val batchScan = collectBatchScan(df) + assert(batchScan.inputPartitions.size === 5) + assert(batchScan.filteredPartitions.flatten.size === 1, + s"expected 1 partition after pruning, got ${batchScan.filteredPartitions.flatten.size}") + } + } + } + test("scan implementing both runtime filtering interfaces -> rejected") { val tbl = s"$catalogName.tbl_both_interfaces" withTable(tbl) { @@ -247,26 +297,160 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { } } - test("filter on column outside filterAttributes -> not pushed") { + test("filter on column outside filterAttributes -> not pushed, even if declared fully pushed") { val tbl = s"$catalogName.tbl4" val dim = s"$catalogName.dim4" withTable(tbl, dim) { + // p2 is a partition column but is not declared filterable, so no runtime filter is derived + // for it. Declaring it fully pushed as well, which the interface forbids for an attribute + // that is not filterable, must not cost it the post-scan filter: nothing was pushed, so the + // scan prunes nothing and the nonmatching rows would come back. sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source " + "PARTITIONED BY (p1, p2) " + - "TBLPROPERTIES('filter-attributes' = 'p1')") + "TBLPROPERTIES('filter-attributes' = 'p1', 'fully-pushed-filter-attributes' = 'p2')") for (i <- 0 until 5) { - sql(s"INSERT INTO $tbl VALUES ($i, $i, 10)") + sql(s"INSERT INTO $tbl VALUES ($i, $i, $i)") } sql(s"CREATE TABLE $dim (val INT) USING $v2Source") - sql(s"INSERT INTO $dim VALUES (10)") + sql(s"INSERT INTO $dim VALUES (3)") - // 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))) + checkAnswer(df, Row(3, 3, 3)) assert(collectBatchScan(df).runtimeFilters.isEmpty, "Expected no runtime filters for a column outside filterAttributes") assertPushedCatalystPredicates(df, 0) + assertScalarSubqueryEvaluatedAfterScan(df, expected = true) + } + } + + test("two predicates on filter attributes -> pushed together in a single filter() call") { + val tbl = s"$catalogName.tbl_two_predicates" + val dim1 = s"$catalogName.dim_two_predicates1" + val dim2 = s"$catalogName.dim_two_predicates2" + withTable(tbl, dim1, dim2) { + sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source PARTITIONED BY (p1, p2)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i, ${i * 10})") + } + sql(s"CREATE TABLE $dim1 (val INT) USING $v2Source") + sql(s"INSERT INTO $dim1 VALUES (3)") + sql(s"CREATE TABLE $dim2 (val INT) USING $v2Source") + sql(s"INSERT INTO $dim2 VALUES (30)") + + val df = sql(s"SELECT * FROM $tbl WHERE p1 = (SELECT max(val) FROM $dim1) " + + s"AND p2 = (SELECT max(val) FROM $dim2)") + checkAnswer(df, Row(3, 3, 30)) + + assertScalarSubqueryRuntimeFilters(df, expectedCount = 2) + val p1 = AttributeReference("p1", IntegerType, nullable = false)() + val p2 = AttributeReference("p2", IntegerType, nullable = false)() + assertPushedCatalystPredicatesEqual( + df, EqualTo(p1, Literal(3)), EqualTo(p2, Literal(30))) + assert(getCatalystScan(df).filterCallCount === 1, + "expected both predicates pushed in a single filter() call") + } + } + + test("nested field of a filter attribute -> pushed with the nested access intact") { + val tbl = s"$catalogName.tbl_nested" + val dim = s"$catalogName.dim_nested" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, s STRUCT) USING $v2Source " + + "PARTITIONED BY (s.tz)") + for (i <- 0 until 3) { + sql(s"INSERT INTO $tbl VALUES ($i, named_struct('tz', 'tz$i'))") + } + sql(s"CREATE TABLE $dim (val STRING) USING $v2Source") + sql(s"INSERT INTO $dim VALUES ('tz1')") + + // The scan declares the top-level struct column `s` as its filter attribute, so the + // predicate qualifies for pushdown even though it reaches into `s.tz`. Matching the nested + // access against the partition layout is left to the scan, which this fixture does. + val df = sql(s"SELECT * FROM $tbl WHERE s.tz = (SELECT max(val) FROM $dim)") + checkAnswer(df, Row(1, Row("tz1"))) + + assertScalarSubqueryRuntimeFilters(df) + val pushed = getPushedCatalystPredicates(df) + assert(pushed.size === 1, s"expected a single pushed predicate, got $pushed") + val nestedAccesses = pushed.head.collect { case g: GetStructField => g } + assert(nestedAccesses.size === 1, + s"expected the pushed predicate to keep the nested access, got ${pushed.head}") + assert(nestedAccesses.head.childSchema.fieldNames.contains("tz")) + } + } + + test("dotted top-level and nested partition columns -> bound to the correct partition slot") { + val tbl = s"$catalogName.tbl_dotted_collision" + val dim = s"$catalogName.dim_dotted_collision" + withTable(tbl, dim) { + // Two partition columns whose dotted names collide: the quoted top-level column `x.y` and + // the nested field `x`.`y`. They carry different values in each row, so a predicate bound + // to the wrong slot would prune the wrong partitions. `x.y` is 3 exactly where `x`.`y` is + // 30, so binding a filter on the nested field to the top-level slot would find nothing. + sql(s"CREATE TABLE $tbl (id INT, `x.y` INT, x STRUCT) USING $v2Source " + + "PARTITIONED BY (`x.y`, x.y)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i, named_struct('y', ${i * 10}))") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (30)") + + // Alias the table so `f.x.y` unambiguously reads the nested field, not the column `x.y`. + val df = sql(s"SELECT * FROM $tbl f WHERE f.x.y = (SELECT max(val) FROM $dim)") + checkAnswer(df, Row(3, 3, Row(30))) + + assertScalarSubqueryRuntimeFilters(df) + // The pushed predicate keeps the nested access, and the scan prunes to the single partition + // whose nested `x`.`y` is 30 rather than binding to the colliding top-level `x.y` slot. + val pushed = getPushedCatalystPredicates(df) + assert(pushed.size === 1, s"expected a single pushed predicate, got $pushed") + assert(pushed.head.exists(_.isInstanceOf[GetStructField]), + s"expected the pushed predicate to keep the nested access, got ${pushed.head}") + val batchScan = collectBatchScan(df) + assert(batchScan.inputPartitions.size === 5) + assert(batchScan.filteredPartitions.flatten.size === 1, + s"expected 1 partition after pruning, got ${batchScan.filteredPartitions.flatten.size}") + } + } + + test("filterAttributes that is not a top-level scan attribute") { + val tbl = s"$catalogName.tbl_unresolvable_attr" + withTable(tbl) { + sql(s"CREATE TABLE $tbl (id INT, part INT, s STRUCT) USING $v2Source " + + "PARTITIONED BY (part)") + sql(s"INSERT INTO $tbl VALUES (1, 1, named_struct('tz', 'a'))") + + val scanRelation = sql(s"SELECT * FROM $tbl").queryExecution.optimizedPlan.collectFirst { + case r: DataSourceV2ScanRelation => r + }.getOrElse(fail("Expected a DataSourceV2ScanRelation")) + + // An attribute the read schema does not carry, such as one pruned out of the projection. + val missing = intercept[AnalysisException] { + scanRelation.copy(scan = new MissingFilterAttributeScan).runtimeFilterAttrs + } + checkError( + exception = missing, + condition = "_LEGACY_ERROR_TEMP_1137", + parameters = Map("name" -> "missing", "outputStr" -> "id,part,s")) + + // A nested reference is rejected up front, since `filterAttributes()` must return + // top-level read-schema attributes. This holds even over an int column that could never + // carry a nested field. + val nested = intercept[SparkException] { + scanRelation.copy(scan = new NestedFilterAttributeScan).runtimeFilterAttrs + } + assert(nested.getMessage.contains("must be a top-level attribute"), + s"expected the nested reference to be rejected, got ${nested.getMessage}") + + // Over a struct it is rejected the same way, rather than widening to the struct column: + // accepting `s.tz` would make runtime filters over every field of `s` eligible, not just + // `s.tz`. + val struct = intercept[SparkException] { + scanRelation.copy(scan = new StructNestedFilterAttributeScan).runtimeFilterAttrs + } + assert(struct.getMessage.contains("must be a top-level attribute"), + s"expected the nested struct reference to be rejected, got ${struct.getMessage}") } } @@ -286,6 +470,55 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { } } + test("ALTER TABLE keeps the Catalyst runtime-filter table type") { + val tbl = s"$catalogName.tbl_alter" + withTable(tbl) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + sql(s"ALTER TABLE $tbl ADD COLUMNS (extra INT)") + + val table = spark.sessionState.catalogManager.catalog(catalogName) + .asInstanceOf[TableCatalog] + .loadTable(Identifier.of(Array.empty, "tbl_alter")) + assert(table.isInstanceOf[InMemoryCatalystRuntimeFilterTable], + s"ALTER TABLE reconstructed ${table.getClass.getName}") + + sql(s"INSERT INTO $tbl VALUES (0, 0, 10), (1, 1, 11)") + checkAnswer(sql(s"SELECT id, part FROM $tbl WHERE extra = 11"), Row(1, 1)) + } + } + + /** + * While SPJ is active the scan's partitioning has to survive runtime filtering, so the + * post-filter partitions still line up with the other side of the join: splits may be pruned, + * but the source may not drop a partition key, invent one, or grow a key's split count. + */ + test("data source that breaks the partitioning it reported -> rejected") { + val partAttr = AttributeReference("part", IntegerType)() + val table = new InMemoryTable("t", Array(Column.create("part", IntegerType)), + Array.empty[Transform], java.util.Collections.emptyMap[String, String]) + val partitioning = KeyedPartitioning( + Seq(partAttr), + Seq(InternalRowComparableWrapper(InternalRow(1), Seq(partAttr))), + isGrouped = false) + + def replanAfterFiltering(afterFilter: Seq[InputPartition]): Unit = { + val scan = new PartitioningBreakingScan(Seq(KeyedInputPartition(1)), afterFilter) + PushDownUtils.replanWithRuntimeFilters(scan, Seq(EqualTo(partAttr, Literal(1))), table, + Seq(partAttr), partitioning, originalPartitions = Seq.empty) + } + + val keyDropped = intercept[SparkException](replanAfterFiltering(Seq(new InputPartition {}))) + assert(keyDropped.getMessage.contains("must have preserved the original partitioning")) + + val keyInvented = intercept[SparkException](replanAfterFiltering(Seq(KeyedInputPartition(99)))) + assert(keyInvented.getMessage.contains("must not report new partition keys")) + + val splitsGrown = intercept[SparkException] { + replanAfterFiltering(Seq(KeyedInputPartition(1), KeyedInputPartition(1))) + } + assert(splitsGrown.getMessage.contains("must not report new partitions for a given key")) + } + // --------------------------------------------------------------------------- // Helper methods // --------------------------------------------------------------------------- @@ -339,15 +572,20 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { }.getOrElse(fail("Expected BatchScanExec in plan")) } - private def getPushedCatalystPredicates(df: DataFrame): Seq[Expression] = { + private type CatalystScan = + InMemoryCatalystRuntimeFilterTable#InMemoryCatalystRuntimeFilterBatchScan + + private def getCatalystScan(df: DataFrame): CatalystScan = { collectBatchScan(df).scan match { - case s: InMemoryCatalystRuntimeFilterTable#InMemoryCatalystRuntimeFilterBatchScan => - s.pushedCatalystPredicates - case other => - fail(s"Expected InMemoryCatalystRuntimeFilterBatchScan, got $other") + case s: CatalystScan => s + case other => fail(s"Expected InMemoryCatalystRuntimeFilterBatchScan, got $other") } } + private def getPushedCatalystPredicates(df: DataFrame): Seq[Expression] = { + getCatalystScan(df).pushedCatalystPredicates + } + private def assertPushedCatalystPredicates(df: DataFrame, expected: Int): Unit = { val preds = getPushedCatalystPredicates(df) assert(preds.size === expected, @@ -398,3 +636,73 @@ private class BothRuntimeFilteringInterfacesScan override def filter(expressions: Array[Expression]): Unit = {} } + +/** A scan declaring a filter attribute the read schema does not carry. */ +private class MissingFilterAttributeScan extends Scan with SupportsRuntimeCatalystFiltering { + + override def readSchema(): StructType = new StructType().add("part", IntegerType) + + override def filterAttributes(): Array[NamedReference] = Array(FieldReference("missing")) + + override def filter(expressions: Array[Expression]): Unit = {} +} + +/** + * A scan breaking the rule that a filter attribute must be a top level read schema column: it + * reports `part.nested` over the int column `part`, so it is rejected as a nested reference. + */ +private class NestedFilterAttributeScan extends Scan with SupportsRuntimeCatalystFiltering { + + override def readSchema(): StructType = new StructType().add("part", IntegerType) + + override def filterAttributes(): Array[NamedReference] = + Array(FieldReference(Seq("part", "nested"))) + + override def filter(expressions: Array[Expression]): Unit = {} +} + +/** + * A scan breaking the same rule over a struct column: it reports `s.tz` where `s` is a struct. + * The nested reference is rejected rather than widening to the struct column `s`. + */ +private class StructNestedFilterAttributeScan extends Scan with SupportsRuntimeCatalystFiltering { + + override def readSchema(): StructType = + new StructType().add("s", new StructType().add("tz", StringType)) + + override def filterAttributes(): Array[NamedReference] = Array(FieldReference(Seq("s", "tz"))) + + override def filter(expressions: Array[Expression]): Unit = {} +} + +private case class KeyedInputPartition(key: Int) extends InputPartition with HasPartitionKey { + override def partitionKey(): InternalRow = InternalRow(key) +} + +/** + * A scan reporting one set of partitions before filtering and another after, so it can break the + * requirement to preserve the partitioning it originally reported. + */ +private class PartitioningBreakingScan( + initialPartitions: Seq[InputPartition], + afterFilter: Seq[InputPartition]) + extends Scan with Batch with SupportsRuntimeCatalystFiltering { + + private var filtered = false + + override def readSchema(): StructType = new StructType().add("part", IntegerType) + + override def toBatch: Batch = this + + override def planInputPartitions(): Array[InputPartition] = + if (filtered) afterFilter.toArray else initialPartitions.toArray + + override def createReaderFactory(): PartitionReaderFactory = + throw new UnsupportedOperationException() + + override def filterAttributes(): Array[NamedReference] = Array(FieldReference("part")) + + override def filter(expressions: Array[Expression]): Unit = { + filtered = true + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala index c1741cac8ad3c..98ed8bf0b3d0d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala @@ -33,9 +33,12 @@ abstract class DistributionAndOrderingSuiteBase extends SharedSparkSession with BeforeAndAfter with AdaptiveSparkPlanHelper { import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._ + /** The catalog implementation `testcat` is registered with. */ + protected def catalogClassName: String = classOf[InMemoryCatalog].getName + override def beforeAll(): Unit = { super.beforeAll() - spark.conf.set("spark.sql.catalog.testcat", classOf[InMemoryCatalog].getName) + spark.conf.set("spark.sql.catalog.testcat", catalogClassName) } override def afterAll(): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala index f0c676d2d74f8..a4cb930aab5a3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala @@ -25,7 +25,7 @@ import org.apache.spark.sql.{DataFrame, ExplainSuiteHelper, Row} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, Literal, TransformExpression} import org.apache.spark.sql.catalyst.plans.physical -import org.apache.spark.sql.connector.catalog.{Column, Identifier, InMemoryTableCatalog} +import org.apache.spark.sql.connector.catalog.{Column, Identifier, InMemoryCatalystRuntimeFilterCatalog, InMemoryTableCatalog} import org.apache.spark.sql.connector.catalog.functions._ import org.apache.spark.sql.connector.distributions.Distributions import org.apache.spark.sql.connector.expressions._ @@ -45,8 +45,286 @@ import org.apache.spark.sql.functions.{col, max} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.SQLConf._ import org.apache.spark.sql.types._ +import org.apache.spark.tags.ExtendedSQLTest -class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with ExplainSuiteHelper { +abstract class KeyGroupedPartitioningSuiteBase extends DistributionAndOrderingSuiteBase { + + protected val emptyProps: java.util.Map[String, String] = { + Collections.emptyMap[String, String] + } + + protected val items: String = "items" + protected val itemsColumns: Array[Column] = Array( + Column.create("id", LongType), + Column.create("name", StringType), + Column.create("price", FloatType), + Column.create("arrive_time", TimestampType)) + + protected val purchases: String = "purchases" + protected val purchasesColumns: Array[Column] = Array( + Column.create("item_id", LongType), + Column.create("price", FloatType), + Column.create("time", TimestampType)) + + protected def createTable( + table: String, + columns: Array[Column], + partitions: Array[Transform], + ordering: Array[SortOrder] = Array.empty, + catalog: InMemoryTableCatalog = catalog): Unit = { + catalog.createTable(Identifier.of(Array("ns"), table), + columns, partitions, emptyProps, Distributions.unspecified(), ordering, None, None, + numRowsPerSplit = 1) + } + + protected def collectShuffles(plan: SparkPlan): Seq[ShuffleExchangeLike] = { + // here we skip collecting shuffle operators that are not associated with SMJ + collect(plan) { + case s: SortMergeJoinExec => s + }.flatMap(smj => + collect(smj) { + case s: ShuffleExchangeExec => s + }) + }.toSet.toSeq + + protected def collectGroupPartitions(plan: SparkPlan): Seq[GroupPartitionsExec] = { + // here we skip collecting group-partition operators that are not associated with SMJ + collect(plan) { + case s: SortMergeJoinExec => s + }.flatMap(smj => + collect(smj) { + case g: GroupPartitionsExec => g + }) + }.toSet.toSeq + + protected def collectScans(plan: SparkPlan): Seq[BatchScanExec] = { + collect(plan) { case s: BatchScanExec => s } + } + +} + +/** + * Tests for runtime filtering under a storage-partitioned join, whose outcome depends on how the + * scan takes runtime filters. + */ +trait KeyGroupedPartitioningRuntimeFilterTests extends KeyGroupedPartitioningSuiteBase { + + /** + * Helper method to verify that filteredPartitions contains the expected number of + * Some and None values. This is used to verify that dynamic partition filtering + * properly fills filtered-out partitions with None. + */ + private def assertFilteredPartitions( + scans: Seq[BatchScanExec], + expectedTotalPartitions: Seq[Int], + expectedFilteredOutPartitions: Seq[Int]): Unit = { + assert(scans.size === expectedTotalPartitions.size, + s"Expected ${expectedTotalPartitions.size} scans but got ${scans.size}") + + scans.zip(expectedTotalPartitions).zip(expectedFilteredOutPartitions).foreach { + case ((scan, expectedTotal), expectedFiltered) => + val filtered = scan.filteredPartitions + assert(filtered.size === expectedTotal, + s"Expected $expectedTotal total partitions but got ${filtered.size}") + + val noneCount = filtered.count(_.isEmpty) + assert(noneCount === expectedFiltered, + s"Expected $expectedFiltered None values but got $noneCount") + + val someCount = filtered.count(_.isDefined) + assert(someCount === (expectedTotal - expectedFiltered), + s"Expected ${expectedTotal - expectedFiltered} Some values but got $someCount") + } + } + + test("data source partitioning + dynamic partition filtering") { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + 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") { + val items_partitions = Array(identity("id")) + createTable(items, itemsColumns, items_partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + + s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + + s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") + + val purchases_partitions = Array(identity("item_id")) + createTable(purchases, purchasesColumns, purchases_partitions) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(1, 42.0, cast('2020-01-01' as timestamp)), " + + s"(1, 44.0, cast('2020-01-15' as timestamp)), " + + s"(1, 45.0, cast('2020-01-15' as timestamp)), " + + s"(2, 11.0, cast('2020-01-01' as timestamp)), " + + s"(3, 19.5, cast('2020-02-01' as timestamp))") + + Seq(true, false).foreach { pushDownValues => + withSQLConf(SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString) { + // number of unique partitions changed after dynamic filtering - the gap should be filled + // with empty partitions and the job should still succeed + var df = sql(s"SELECT sum(p.price) from testcat.ns.$items i, testcat.ns.$purchases p " + + "WHERE i.id = p.item_id AND i.price > 40.0") + + var shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") + var scans = collectScans(df.queryExecution.executedPlan) + assert(scans.forall(_.outputPartitioning.numPartitions === 5)) + var groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) + assert(groupPartitions.forall(_.outputPartitioning.numPartitions === 3)) + + checkAnswer(df, Seq(Row(131))) + + // Verify that filteredPartitions contains None for filtered-out partitions. + // After DPF with filter i.price > 40.0, only id=1 survives on items side. + // The purchases side should be pruned to only item_id=1. + // purchases: 5 total partitions (3 for id=1, 1 for id=2, 1 for id=3) + // After DPF: 3 Some (id=1), 2 None (id=2, id=3) + assertFilteredPartitions(scans, Seq(5, 5), Seq(0, 2)) + + // dynamic filtering doesn't change partitioning so storage-partitioned join should kick + // in + df = sql(s"SELECT sum(p.price) from testcat.ns.$items i, testcat.ns.$purchases p " + + "WHERE i.id = p.item_id AND i.price >= 10.0") + + shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") + scans = collectScans(df.queryExecution.executedPlan) + assert(scans.forall(_.outputPartitioning.numPartitions === 5)) + groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) + assert(groupPartitions.forall(_.outputPartitioning.numPartitions === 3)) + + checkAnswer(df, Seq(Row(303.5))) + + // With filter i.price >= 10.0, all ids (1, 2, 3) survive, + // so no partitions should be filtered out + assertFilteredPartitions(scans, Seq(5, 5), Seq(0, 0)) + } + } + } + } + + test("SPARK-42038: partially clustered: with dynamic partition filtering") { + val items_partitions = Array(identity("id")) + createTable(items, itemsColumns, items_partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + + s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + + s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp)), " + + s"(4, 'dd', 18.0, cast('2023-01-01' as timestamp))") + + val purchases_partitions = Array(identity("item_id")) + createTable(purchases, purchasesColumns, purchases_partitions) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(1, 42.0, cast('2020-01-01' as timestamp)), " + + s"(1, 44.0, cast('2020-01-15' as timestamp)), " + + s"(1, 45.0, cast('2020-01-15' as timestamp)), " + + s"(1, 50.0, cast('2020-01-15' as timestamp)), " + + s"(1, 55.0, cast('2020-01-15' as timestamp)), " + + s"(1, 60.0, cast('2020-01-15' as timestamp)), " + + s"(1, 65.0, cast('2020-01-15' as timestamp)), " + + s"(2, 11.0, cast('2020-01-01' as timestamp)), " + + s"(3, 19.5, cast('2020-02-01' as timestamp)), " + + s"(5, 25.0, cast('2023-01-01' as timestamp)), " + + s"(5, 26.0, cast('2023-01-01' as timestamp)), " + + s"(5, 28.0, cast('2023-01-01' as timestamp)), " + + s"(6, 50.0, cast('2023-02-01' as timestamp)), " + + s"(6, 50.0, cast('2023-02-01' as timestamp))") + + Seq(true, false).foreach { pushDownValues => + Seq(("true", 15), ("false", 6)).foreach { + case (enable, expected) => + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + 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", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> enable) { + + // When partition values are pushed down, storage-partitioned join fills the missing + // partitions & splits after dynamic filtering with empty partitions & splits. + val df = sql(s"SELECT sum(p.price) from " + + s"testcat.ns.$purchases p, testcat.ns.$items i WHERE " + + s"p.item_id = i.id AND p.price < 45.0") + + checkAnswer(df, Seq(Row(213.5))) + val shuffles = collectShuffles(df.queryExecution.executedPlan) + val scans = collectScans(df.queryExecution.executedPlan) + val groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) + assert(scans.map(_.outputPartitioning.numPartitions) === Seq(14, 6)) + if (pushDownValues) { + assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") + assert(groupPartitions.forall(_.outputPartitioning.numPartitions === expected)) + } else { + assert(shuffles.nonEmpty, + "should contain shuffle when not pushing down partition values") + assert(groupPartitions.isEmpty) + } + + // Verify filteredPartitions for DPF. + // After filter p.price < 45.0, purchases has item_ids {1, 2, 3, 5}. + // Items side should be pruned to these ids. Since items has {1, 2, 3, 4}, + // id=4 should be filtered out. + // purchases: 14 total, all kept (0 None) - no DPF on probe side + // items: 6 total, id=4 filtered (1 None) + assertFilteredPartitions(scans, Seq(14, 6), Seq(0, 1)) + } + } + } + } + + test("SPARK-45652: SPJ should handle empty partition after dynamic filtering") { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + 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") { + val items_partitions = Array(identity("id")) + createTable(items, itemsColumns, items_partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + + s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + + s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") + + val purchases_partitions = Array(identity("item_id")) + createTable(purchases, purchasesColumns, purchases_partitions) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(1, 42.0, cast('2020-01-01' as timestamp)), " + + s"(1, 44.0, cast('2020-01-15' as timestamp)), " + + s"(1, 45.0, cast('2020-01-15' as timestamp)), " + + s"(2, 11.0, cast('2020-01-01' as timestamp)), " + + s"(3, 19.5, cast('2020-02-01' as timestamp))") + + Seq(true, false).foreach { pushDownValues => + Seq(true, false).foreach { partiallyClustered => { + withSQLConf( + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> + partiallyClustered.toString, + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString) { + // The dynamic filtering effectively filtered out all the partitions + val df = sql(s"SELECT p.price from testcat.ns.$items i, testcat.ns.$purchases p " + + "WHERE i.id = p.item_id AND i.price > 50.0") + checkAnswer(df, Seq.empty) + } + } + } + } + } + } +} + +class KeyGroupedPartitioningSuite + extends KeyGroupedPartitioningSuiteBase with ExplainSuiteHelper { private val functions = Seq( UnboundYearsFunction, UnboundDaysFunction, @@ -68,9 +346,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with catalog.clearFunctions() } - private val emptyProps: java.util.Map[String, String] = { - Collections.emptyMap[String, String] - } private val table: String = "tbl" private val columns: Array[Column] = Array( @@ -256,17 +531,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with assert(expectedPartitioning == scan.outputPartitioning) } - private def createTable( - table: String, - columns: Array[Column], - partitions: Array[Transform], - ordering: Array[SortOrder] = Array.empty, - catalog: InMemoryTableCatalog = catalog): Unit = { - catalog.createTable(Identifier.of(Array("ns"), table), - columns, partitions, emptyProps, Distributions.unspecified(), ordering, None, None, - numRowsPerSplit = 1) - } - private val customers: String = "customers" private val customersColumns: Array[Column] = Array( Column.create("customer_name", StringType), @@ -307,92 +571,39 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with s"('aaa', 10, 1), ('bbb', 20, 2), ('ccc', 30, 3)") createTable(orders, ordersColumns, orders_partitions) - sql(s"INSERT INTO testcat.ns.$orders VALUES " + - s"(100.0, 1), (200.0, 1), (150.0, 2), (250.0, 2), (350.0, 2), (400.50, 3)") - - val df = sql( - s""" - |${selectWithMergeJoinHint("c", "o")} - |customer_name, customer_age, order_amount - |FROM testcat.ns.$customers c JOIN testcat.ns.$orders o - |ON c.customer_id = o.customer_id ORDER BY c.customer_id, order_amount - |""".stripMargin) - - val shuffles = collectShuffles(df.queryExecution.executedPlan) - assert(shuffles.length == expectedNumOfShuffleExecs) - - val groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) - assert(groupPartitions.length == expectedGroupPartitionsExecs) - - checkAnswer(df, - Seq(Row("aaa", 10, 100.0), Row("aaa", 10, 200.0), Row("bbb", 20, 150.0), - Row("bbb", 20, 250.0), Row("bbb", 20, 350.0), Row("ccc", 30, 400.50))) - } - - protected def collectAllShuffles(plan: SparkPlan): Seq[ShuffleExchangeLike] = { - collect(plan) { - case s: ShuffleExchangeExec => s - } - } - - protected def collectAllGroupPartitions(plan: SparkPlan): Seq[GroupPartitionsExec] = { - collect(plan) { - case g: GroupPartitionsExec => g - } - } - - protected def collectShuffles(plan: SparkPlan): Seq[ShuffleExchangeLike] = { - // here we skip collecting shuffle operators that are not associated with SMJ - collect(plan) { - case s: SortMergeJoinExec => s - }.flatMap(smj => - collect(smj) { - case s: ShuffleExchangeExec => s - }) - }.toSet.toSeq - - protected def collectGroupPartitions(plan: SparkPlan): Seq[GroupPartitionsExec] = { - // here we skip collecting shuffle operators that are not associated with SMJ - collect(plan) { - case s: SortMergeJoinExec => s - }.flatMap(smj => - collect(smj) { - case g: GroupPartitionsExec => g - }) - }.toSet.toSeq + sql(s"INSERT INTO testcat.ns.$orders VALUES " + + s"(100.0, 1), (200.0, 1), (150.0, 2), (250.0, 2), (350.0, 2), (400.50, 3)") - private def collectScans(plan: SparkPlan): Seq[BatchScanExec] = { - collect(plan) { case s: BatchScanExec => s } - } + val df = sql( + s""" + |${selectWithMergeJoinHint("c", "o")} + |customer_name, customer_age, order_amount + |FROM testcat.ns.$customers c JOIN testcat.ns.$orders o + |ON c.customer_id = o.customer_id ORDER BY c.customer_id, order_amount + |""".stripMargin) - /** - * Helper method to verify that filteredPartitions contains the expected number of - * Some and None values. This is used to verify that dynamic partition filtering - * properly fills filtered-out partitions with None. - */ - private def assertFilteredPartitions( - scans: Seq[BatchScanExec], - expectedTotalPartitions: Seq[Int], - expectedFilteredOutPartitions: Seq[Int]): Unit = { - assert(scans.size === expectedTotalPartitions.size, - s"Expected ${expectedTotalPartitions.size} scans but got ${scans.size}") + val shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.length == expectedNumOfShuffleExecs) - scans.zip(expectedTotalPartitions).zip(expectedFilteredOutPartitions).foreach { - case ((scan, expectedTotal), expectedFiltered) => - val filtered = scan.filteredPartitions - assert(filtered.size === expectedTotal, - s"Expected $expectedTotal total partitions but got ${filtered.size}") + val groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) + assert(groupPartitions.length == expectedGroupPartitionsExecs) - val noneCount = filtered.count(_.isEmpty) - assert(noneCount === expectedFiltered, - s"Expected $expectedFiltered None values but got $noneCount") + checkAnswer(df, + Seq(Row("aaa", 10, 100.0), Row("aaa", 10, 200.0), Row("bbb", 20, 150.0), + Row("bbb", 20, 250.0), Row("bbb", 20, 350.0), Row("ccc", 30, 400.50))) + } - val someCount = filtered.count(_.isDefined) - assert(someCount === (expectedTotal - expectedFiltered), - s"Expected ${expectedTotal - expectedFiltered} Some values but got $someCount") + protected def collectAllShuffles(plan: SparkPlan): Seq[ShuffleExchangeLike] = { + collect(plan) { + case s: ShuffleExchangeExec => s } } + protected def collectAllGroupPartitions(plan: SparkPlan): Seq[GroupPartitionsExec] = { + collect(plan) { + case g: GroupPartitionsExec => g + } + } test("partitioned join: exact distribution (same number of buckets) from both sides") { val customers_partitions = Array(bucket(4, "customer_id")) @@ -415,19 +626,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with testWithCustomersAndOrders(customers_partitions, Array.empty, 2, 0) } - private val items: String = "items" - private val itemsColumns: Array[Column] = Array( - Column.create("id", LongType), - Column.create("name", StringType), - Column.create("price", FloatType), - Column.create("arrive_time", TimestampType)) - - private val purchases: String = "purchases" - private val purchasesColumns: Array[Column] = Array( - Column.create("item_id", LongType), - Column.create("price", FloatType), - Column.create("time", TimestampType)) - private val details: String = "details" private val detailsColumns: Array[Column] = Array( Column.create("item_id", LongType), @@ -1381,149 +1579,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with } } - test("data source partitioning + dynamic partition filtering") { - withSQLConf( - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - 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") { - val items_partitions = Array(identity("id")) - createTable(items, itemsColumns, items_partitions) - sql(s"INSERT INTO testcat.ns.$items VALUES " + - s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + - s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + - s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + - s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + - s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") - - val purchases_partitions = Array(identity("item_id")) - createTable(purchases, purchasesColumns, purchases_partitions) - sql(s"INSERT INTO testcat.ns.$purchases VALUES " + - s"(1, 42.0, cast('2020-01-01' as timestamp)), " + - s"(1, 44.0, cast('2020-01-15' as timestamp)), " + - s"(1, 45.0, cast('2020-01-15' as timestamp)), " + - s"(2, 11.0, cast('2020-01-01' as timestamp)), " + - s"(3, 19.5, cast('2020-02-01' as timestamp))") - - Seq(true, false).foreach { pushDownValues => - withSQLConf(SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString) { - // number of unique partitions changed after dynamic filtering - the gap should be filled - // with empty partitions and the job should still succeed - var df = sql(s"SELECT sum(p.price) from testcat.ns.$items i, testcat.ns.$purchases p " + - "WHERE i.id = p.item_id AND i.price > 40.0") - - var shuffles = collectShuffles(df.queryExecution.executedPlan) - assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") - var scans = collectScans(df.queryExecution.executedPlan) - assert(scans.forall(_.outputPartitioning.numPartitions === 5)) - var groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) - assert(groupPartitions.forall(_.outputPartitioning.numPartitions === 3)) - - checkAnswer(df, Seq(Row(131))) - - // Verify that filteredPartitions contains None for filtered-out partitions. - // After DPF with filter i.price > 40.0, only id=1 survives on items side. - // The purchases side should be pruned to only item_id=1. - // purchases: 5 total partitions (3 for id=1, 1 for id=2, 1 for id=3) - // After DPF: 3 Some (id=1), 2 None (id=2, id=3) - assertFilteredPartitions(scans, Seq(5, 5), Seq(0, 2)) - - // dynamic filtering doesn't change partitioning so storage-partitioned join should kick - // in - df = sql(s"SELECT sum(p.price) from testcat.ns.$items i, testcat.ns.$purchases p " + - "WHERE i.id = p.item_id AND i.price >= 10.0") - - shuffles = collectShuffles(df.queryExecution.executedPlan) - assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") - scans = collectScans(df.queryExecution.executedPlan) - assert(scans.forall(_.outputPartitioning.numPartitions === 5)) - groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) - assert(groupPartitions.forall(_.outputPartitioning.numPartitions === 3)) - - checkAnswer(df, Seq(Row(303.5))) - - // With filter i.price >= 10.0, all ids (1, 2, 3) survive, - // so no partitions should be filtered out - assertFilteredPartitions(scans, Seq(5, 5), Seq(0, 0)) - } - } - } - } - - test("SPARK-42038: partially clustered: with dynamic partition filtering") { - val items_partitions = Array(identity("id")) - createTable(items, itemsColumns, items_partitions) - sql(s"INSERT INTO testcat.ns.$items VALUES " + - s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + - s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + - s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + - s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + - s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp)), " + - s"(4, 'dd', 18.0, cast('2023-01-01' as timestamp))") - - val purchases_partitions = Array(identity("item_id")) - createTable(purchases, purchasesColumns, purchases_partitions) - sql(s"INSERT INTO testcat.ns.$purchases VALUES " + - s"(1, 42.0, cast('2020-01-01' as timestamp)), " + - s"(1, 44.0, cast('2020-01-15' as timestamp)), " + - s"(1, 45.0, cast('2020-01-15' as timestamp)), " + - s"(1, 50.0, cast('2020-01-15' as timestamp)), " + - s"(1, 55.0, cast('2020-01-15' as timestamp)), " + - s"(1, 60.0, cast('2020-01-15' as timestamp)), " + - s"(1, 65.0, cast('2020-01-15' as timestamp)), " + - s"(2, 11.0, cast('2020-01-01' as timestamp)), " + - s"(3, 19.5, cast('2020-02-01' as timestamp)), " + - s"(5, 25.0, cast('2023-01-01' as timestamp)), " + - s"(5, 26.0, cast('2023-01-01' as timestamp)), " + - s"(5, 28.0, cast('2023-01-01' as timestamp)), " + - s"(6, 50.0, cast('2023-02-01' as timestamp)), " + - s"(6, 50.0, cast('2023-02-01' as timestamp))") - - Seq(true, false).foreach { pushDownValues => - Seq(("true", 15), ("false", 6)).foreach { - case (enable, expected) => - withSQLConf( - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - 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", - SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, - SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> enable) { - - // storage-partitioned join should kick in and fill the missing partitions & splits - // after dynamic filtering with empty partitions & splits, respectively. - val df = sql(s"SELECT sum(p.price) from " + - s"testcat.ns.$purchases p, testcat.ns.$items i WHERE " + - s"p.item_id = i.id AND p.price < 45.0") - - checkAnswer(df, Seq(Row(213.5))) - val shuffles = collectShuffles(df.queryExecution.executedPlan) - val scans = collectScans(df.queryExecution.executedPlan) - val groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) - assert(scans.map(_.outputPartitioning.numPartitions) === Seq(14, 6)) - if (pushDownValues) { - assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") - assert(groupPartitions.forall(_.outputPartitioning.numPartitions === expected)) - } else { - assert(shuffles.nonEmpty, - "should contain shuffle when not pushing down partition values") - assert(groupPartitions.isEmpty) - } - - // Verify filteredPartitions for DPF. - // After filter p.price < 45.0, purchases has item_ids {1, 2, 3, 5}. - // Items side should be pruned to these ids. Since items has {1, 2, 3, 4}, - // id=4 should be filtered out. - // purchases: 14 total, all kept (0 None) - no DPF on probe side - // items: 6 total, id=4 filtered (1 None) - assertFilteredPartitions(scans, Seq(14, 6), Seq(0, 1)) - } - } - } - } - test("SPARK-41471: shuffle one side: only one side reports partitioning") { val items_partitions = Array(identity("id")) createTable(items, itemsColumns, items_partitions) @@ -2631,48 +2686,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with } } - test("SPARK-45652: SPJ should handle empty partition after dynamic filtering") { - withSQLConf( - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - 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") { - val items_partitions = Array(identity("id")) - createTable(items, itemsColumns, items_partitions) - sql(s"INSERT INTO testcat.ns.$items VALUES " + - s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + - s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + - s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + - s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + - s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") - - val purchases_partitions = Array(identity("item_id")) - createTable(purchases, purchasesColumns, purchases_partitions) - sql(s"INSERT INTO testcat.ns.$purchases VALUES " + - s"(1, 42.0, cast('2020-01-01' as timestamp)), " + - s"(1, 44.0, cast('2020-01-15' as timestamp)), " + - s"(1, 45.0, cast('2020-01-15' as timestamp)), " + - s"(2, 11.0, cast('2020-01-01' as timestamp)), " + - s"(3, 19.5, cast('2020-02-01' as timestamp))") - - Seq(true, false).foreach { pushDownValues => - Seq(true, false).foreach { partiallyClustered => { - withSQLConf( - SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> - partiallyClustered.toString, - SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString) { - // The dynamic filtering effectively filtered out all the partitions - val df = sql(s"SELECT p.price from testcat.ns.$items i, testcat.ns.$purchases p " + - "WHERE i.id = p.item_id AND i.price > 50.0") - checkAnswer(df, Seq.empty) - } - } - } - } - } - } - test("SPARK-48012: one-side shuffle with partition transforms") { val items_partitions = Array(bucket(2, "id"), identity("arrive_time")) val items_partitions2 = Array(identity("arrive_time"), bucket(2, "id")) @@ -4716,3 +4729,33 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with } } } + +/** + * Runs the runtime filtering tests against a catalog whose scans take runtime filters as connector + * predicates, via [[org.apache.spark.sql.connector.read.SupportsRuntimeFiltering]]. + */ +@ExtendedSQLTest +class KeyGroupedPartitioningRuntimeFilterSuite + extends KeyGroupedPartitioningSuiteBase with KeyGroupedPartitioningRuntimeFilterTests { + + after { + catalog.clearTables() + } +} + +/** + * Runs the runtime filtering tests against a catalog whose scans take runtime filters as Catalyst + * expressions, via + * [[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]]. + */ +@ExtendedSQLTest +class KeyGroupedPartitioningCatalystRuntimeFilterSuite + extends KeyGroupedPartitioningSuiteBase with KeyGroupedPartitioningRuntimeFilterTests { + + override protected def catalogClassName: String = + classOf[InMemoryCatalystRuntimeFilterCatalog].getName + + after { + catalog.clearTables() + } +}