cloud-fan commented on code in PR #58296:
URL: https://github.com/apache/spark/pull/58296#discussion_r3868500259
##########
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:
**Non-blocking:**
Preserve the Catalyst-filtering table type through ALTER TABLE as well as
creation. The inherited alterTable path matches this subclass as
InMemoryTableWithV2Filter and rebuilds the predicate-filtering fixture, so
later DPP or SPJ tests exercise the wrong interface. Please route both create
and alter through a shared overridable table factory that forwards the full
metadata.
##########
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:
**Nit:**
Scope this comment to pushDownValues=true. In the false branch immediately
below, the test requires a shuffle and no GroupPartitionsExec, so SPJ does not
kick in for every iteration described here.
--
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]
