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 da2b3a54d74c3..7b2cfacb652fe 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 @@ -210,8 +210,7 @@ case class DataSourceV2ScanRelation( case s: SupportsRuntimeCatalystFiltering => s.filterAttributes() case _ => Array.empty[NamedReference] } - AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( - filterAttrs.toImmutableArraySeq, this)) + resolveTopLevelFilterAttrs(filterAttrs) } /** @@ -224,6 +223,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 64c0d0f2bc901..cf213d1bd7900 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._ @@ -47,8 +47,285 @@ import org.apache.spark.sql.internal.SQLConf._ import org.apache.spark.sql.types._ import org.apache.spark.tags.ExtendedSQLTest +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) + } + } + } + } + } + } +} + @ExtendedSQLTest -class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with ExplainSuiteHelper { +class KeyGroupedPartitioningSuite + extends KeyGroupedPartitioningSuiteBase with ExplainSuiteHelper { private val functions = Seq( UnboundYearsFunction, UnboundDaysFunction, @@ -70,9 +347,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( @@ -258,17 +532,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), @@ -312,89 +575,36 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with 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 - - 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")) @@ -417,19 +627,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), @@ -1382,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) @@ -2625,48 +2679,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")) @@ -4689,3 +4701,33 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with checkAnswer(df, Seq(Row(1, "aa", 40.0, 42.0), Row(2, "bb", 10.0, 19.5))) } } + +/** + * 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() + } +}