szehon-ho commented on code in PR #58296:
URL: https://github.com/apache/spark/pull/58296#discussion_r3869814425
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala:
##########
@@ -48,3 +50,41 @@ class InMemoryTableCatalystRuntimeFilterCatalog extends
InMemoryTableCatalog {
createTable(ident, tableInfo.columns(), tableInfo.partitions(),
tableInfo.properties)
}
}
+
+/**
+ * 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 {
Review Comment:
Done. `createTable` and `alterTable` now both go through a shared,
overridable `newInMemoryTable` factory on `BasicInMemoryTableCatalog`. The
Catalyst catalogs mix in `InMemoryCatalystRuntimeFilterTableFactory`, so ALTER
TABLE reconstructs `InMemoryCatalystRuntimeFilterTable` instead of matching the
subclass as `InMemoryTableWithV2Filter` and rebuilding the predicate-filtering
fixture. `InMemoryTableWithV2FilterCatalog` overrides the same factory so its
ALTER path stays on the V2-filter table. Added a test (`ALTER TABLE keeps the
Catalyst runtime-filter table type`) that asserts the reconstructed table type
after `ADD COLUMNS`.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -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) {
+
+ // storage-partitioned join should kick in and fill the missing
partitions & splits
Review Comment:
Done. Scoped the comment to the `pushDownValues=true` case, since the
`false` branch below requires a shuffle and no `GroupPartitionsExec`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
