This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 6e6564ed44 [spark] Make bucket-preserving scans opt-in and retain
parallel read units (#10076)
6e6564ed44 is described below
commit 6e6564ed445fe3139fd89f6ba7e7f2836e573662
Author: Zouxxyy <[email protected]>
AuthorDate: Tue Sep 22 12:14:28 2026 +0800
[spark] Make bucket-preserving scans opt-in and retain parallel read units
(#10076)
---
docs/docs/spark/configuration.mdx | 3 +
docs/docs/spark/sql-query.md | 38 +++
docs/generated/spark_connector_configuration.html | 6 +
.../scala/org/apache/paimon/spark/PaimonScan.scala | 2 +-
.../scala/org/apache/paimon/spark/PaimonScan.scala | 28 +-
.../spark/sql/BucketedScanCompatibilityTest.scala | 21 ++
.../spark/sql/BucketedScanPlanningTest.scala | 21 ++
.../apache/paimon/spark/SparkConnectorOptions.java | 13 +
.../scala/org/apache/paimon/spark/PaimonScan.scala | 28 +-
.../apache/paimon/spark/PaimonScanBuilder.scala | 10 +-
.../DisableUnnecessaryPaimonBucketedScan.scala | 178 -----------
.../extensions/PaimonSparkSessionExtensions.scala | 3 +-
.../paimon/spark/read/BinPackingSplits.scala | 34 ++-
.../apache/paimon/spark/BinPackingSplitsTest.scala | 45 ++-
.../spark/sql/BucketedScanPlanningTestBase.scala | 338 +++++++++++++++++++++
...canSuite.scala => BucketedScanPolicyTest.scala} | 127 ++++----
.../paimon/spark/sql/BucketedTableQueryTest.scala | 1 +
17 files changed, 599 insertions(+), 297 deletions(-)
diff --git a/docs/docs/spark/configuration.mdx
b/docs/docs/spark/configuration.mdx
index 23c2efb6db..daed24caa9 100644
--- a/docs/docs/spark/configuration.mdx
+++ b/docs/docs/spark/configuration.mdx
@@ -81,6 +81,9 @@ SELECT * FROM default.T1 JOIN default.T2 ON xxxx;
## Spark Connector Options
+Batch scans use regular split packing by default. To opt into bucket
distribution for storage
+partition joins, see [Scan Layout and Storage Partition
Joins](./sql-query#scan-layout-and-storage-partition-joins).
+
Streaming admission limits operate on whole splits. See [Triggers and Read
Limits](./structured-streaming#triggers-and-read-limits)
for how byte, row, and file thresholds are applied.
diff --git a/docs/docs/spark/sql-query.md b/docs/docs/spark/sql-query.md
index ffedd7e5ef..71883eeae0 100644
--- a/docs/docs/spark/sql-query.md
+++ b/docs/docs/spark/sql-query.md
@@ -61,6 +61,44 @@ For example:
SELECT *, __paimon_file_path, __paimon_partition, __paimon_bucket,
__paimon_row_index FROM t;
```
+### Scan Layout and Storage Partition Joins
+
+By default, Paimon plans batch read tasks using its regular split packing. It
does not report
+bucket distribution or scan ordering to Spark, so scan parallelism is not
limited by the number
+of selected buckets. Spark still adds the exchanges and sorts required by the
query.
+
+To let Spark use a fixed-bucket table's layout for storage partition joins or
grouped aggregates,
+enable both options before planning the query:
+
+```sql
+SET spark.paimon.scan.preserve-data-grouping=true;
+SET spark.sql.sources.v2.bucketing.enabled=true;
+
+SELECT * FROM t1 JOIN t2 ON t1.bucket_key = t2.bucket_key;
+```
+
+`scan.preserve-data-grouping` defaults to `false`. It can also be set as a
table property or as a
+DataFrame read option. Session and read options follow the precedence
described in
+[Configuration](./configuration). The effective choice is fixed when a scan is
created;
+changing a session option later affects newly created scans.
+
+In grouped mode, Paimon preserves complete splits and can provide multiple
read units for the
+same bucket. Spark may group those units into one task per bucket. When a
supported join uses
+`spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled=true`,
Spark can use
+multiple tasks for a bucket. This is a join optimization; enabling grouped
mode can still reduce
+the parallelism of a plain scan or TopN query.
+
+Grouped mode requires Spark 3.3 or later and a supported fixed-bucket layout.
If Spark V2
+bucketing is disabled, or the scan cannot report a supported bucket layout,
Paimon uses regular
+split packing. Grouping does not guarantee that every join can avoid a shuffle.
+
+**Migration:** Enabling Spark V2 bucketing alone, including its default of
`true` in Spark 4.1,
+no longer opts Paimon into grouped scans. Workloads that depend on bucket
distribution to avoid
+shuffles must also enable `scan.preserve-data-grouping`. The former Paimon
adaptive rule that
+disabled bucket scans after physical planning has been removed;
+`spark.sql.sources.bucketing.autoBucketedScan.enabled` no longer changes a
Paimon scan's layout.
+AQE and non-AQE queries use the same scan policy.
+
### Batch Time Travel
Paimon batch reads with time travel can specify a snapshot or a tag and read
the corresponding data.
diff --git a/docs/generated/spark_connector_configuration.html
b/docs/generated/spark_connector_configuration.html
index d80d14f258..61684d295f 100644
--- a/docs/generated/spark_connector_configuration.html
+++ b/docs/generated/spark_connector_configuration.html
@@ -92,6 +92,12 @@ under the License.
<td>Boolean</td>
<td>Whether to verify SparkSession is initialized with required
configurations.</td>
</tr>
+ <tr>
+ <td><h5>scan.preserve-data-grouping</h5></td>
+ <td style="word-wrap: break-word;">false</td>
+ <td>Boolean</td>
+ <td>Whether batch scans preserve bucket grouping for Spark to use
the table's distribution and ordering. Requires Spark V2 bucketing to be
enabled. If false, scans use regular split packing and report no bucket
distribution or ordering. If true, Spark may group multiple read units into one
task per bucket, or use them separately for a partially clustered join. Set as
a table/read option or with spark.paimon.scan.preserve-data-grouping. The
choice is fixed when the scan is cre [...]
+ </tr>
<tr>
<td><h5>source.split.target-size-with-column-pruning</h5></td>
<td style="word-wrap: break-word;">false</td>
diff --git
a/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
index fab2cfe861..d655c483e3 100644
---
a/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
+++
b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
@@ -37,5 +37,5 @@ case class PaimonScan(
override val pushedFullTextSearch: Option[FullTextSearch] = None,
override val pushedVariantExtractions: Map[Seq[String],
Seq[VariantExtractionInfo]] = Map.empty,
override val pushedMapSelectedKeys: Map[String, Seq[String]] = Map.empty,
- bucketedScanDisabled: Boolean = true)
+ preserveDataGrouping: Boolean = false)
extends PaimonBaseScan(table) {}
diff --git
a/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
index 9cc7e409c7..f0476f6c49 100644
---
a/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
+++
b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
@@ -20,7 +20,7 @@ package org.apache.paimon.spark
import org.apache.paimon.partition.PartitionPredicate
import org.apache.paimon.predicate.{FullTextSearch, HybridSearch, Predicate,
TopN, VectorSearch}
-import org.apache.paimon.spark.read.VariantExtractionInfo
+import org.apache.paimon.spark.read.{BinPackingSplits, VariantExtractionInfo}
import org.apache.paimon.table.{BucketMode, FileStoreTable, InnerTable}
import org.apache.paimon.table.source.{DataSplit, Split}
@@ -41,14 +41,10 @@ case class PaimonScan(
override val pushedFullTextSearch: Option[FullTextSearch] = None,
override val pushedVariantExtractions: Map[Seq[String],
Seq[VariantExtractionInfo]] = Map.empty,
override val pushedMapSelectedKeys: Map[String, Seq[String]] = Map.empty,
- bucketedScanDisabled: Boolean = false)
+ preserveDataGrouping: Boolean = false)
extends PaimonBaseScan(table)
with SupportsReportPartitioning {
- def disableBucketedScan(): PaimonScan = {
- copy(bucketedScanDisabled = true)
- }
-
@transient
private lazy val extractBucketTransform: Option[Transform] = {
table match {
@@ -97,14 +93,16 @@ case class PaimonScan(
}
private def shouldDoBucketedScan: Boolean = {
- !bucketedScanDisabled && conf.v2BucketingEnabled &&
extractBucketTransform.isDefined
+ preserveDataGrouping && extractBucketTransform.isDefined
}
// Since Spark 3.3
override def outputPartitioning: Partitioning = {
- extractBucketTransform
- .map(bucket => new KeyGroupedPartitioning(Array(bucket),
inputPartitions.size))
- .getOrElse(new UnknownPartitioning(0))
+ if (shouldDoBucketedScan) {
+ new KeyGroupedPartitioning(Array(extractBucketTransform.get),
inputPartitions.size)
+ } else {
+ new UnknownPartitioning(0)
+ }
}
override def getInputPartitions(splits: Array[Split]):
Seq[PaimonInputPartition] = {
@@ -112,13 +110,7 @@ case class PaimonScan(
return super.getInputPartitions(splits)
}
- splits
- .map(_.asInstanceOf[DataSplit])
- .groupBy(_.bucket())
- .map {
- case (bucket, groupedSplits) =>
- PaimonBucketedInputPartition(groupedSplits, bucket)
- }
- .toSeq
+ BinPackingSplits(coreOptions, readRowSizeRatio)
+ .packByBucket(splits.map(_.asInstanceOf[DataSplit]))
}
}
diff --git
a/paimon-spark/paimon-spark-4.0/src/test/scala/org/apache/paimon/spark/sql/BucketedScanCompatibilityTest.scala
b/paimon-spark/paimon-spark-4.0/src/test/scala/org/apache/paimon/spark/sql/BucketedScanCompatibilityTest.scala
new file mode 100644
index 0000000000..987775b30e
--- /dev/null
+++
b/paimon-spark/paimon-spark-4.0/src/test/scala/org/apache/paimon/spark/sql/BucketedScanCompatibilityTest.scala
@@ -0,0 +1,21 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.paimon.spark.sql
+
+class BucketedScanCompatibilityTest extends BucketedScanPlanningTestBase
diff --git
a/paimon-spark/paimon-spark-4.1/src/test/scala/org/apache/paimon/spark/sql/BucketedScanPlanningTest.scala
b/paimon-spark/paimon-spark-4.1/src/test/scala/org/apache/paimon/spark/sql/BucketedScanPlanningTest.scala
new file mode 100644
index 0000000000..5fca180f38
--- /dev/null
+++
b/paimon-spark/paimon-spark-4.1/src/test/scala/org/apache/paimon/spark/sql/BucketedScanPlanningTest.scala
@@ -0,0 +1,21 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.paimon.spark.sql
+
+class BucketedScanPlanningTest extends BucketedScanPlanningTestBase
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java
index 4dd9329d1c..82d6855589 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkConnectorOptions.java
@@ -25,6 +25,19 @@ import static org.apache.paimon.options.ConfigOptions.key;
/** Options for spark connector. */
public class SparkConnectorOptions {
+ public static final ConfigOption<Boolean> SCAN_PRESERVE_DATA_GROUPING =
+ key("scan.preserve-data-grouping")
+ .booleanType()
+ .defaultValue(false)
+ .withDescription(
+ "Whether batch scans preserve bucket grouping for
Spark to use the table's "
+ + "distribution and ordering. Requires
Spark V2 bucketing to be enabled. "
+ + "If false, scans use regular split
packing and report no bucket distribution "
+ + "or ordering. If true, Spark may group
multiple read units into one task "
+ + "per bucket, or use them separately for
a partially clustered join. "
+ + "Set as a table/read option or with
spark.paimon.scan.preserve-data-grouping. "
+ + "The choice is fixed when the scan is
created.");
+
public static final ConfigOption<Boolean>
REQUIRED_SPARK_CONFS_CHECK_ENABLED =
key("requiredSparkConfsCheck.enabled")
.booleanType()
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
index b6d89cf6e8..a2e1ede366 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala
@@ -23,7 +23,7 @@ import org.apache.paimon.partition.PartitionPredicate
import org.apache.paimon.predicate.{FullTextSearch, HybridSearch, Predicate,
TopN, VectorSearch}
import org.apache.paimon.spark.catalog.functions.BucketFunction
import org.apache.paimon.spark.commands.BucketExpression.quote
-import org.apache.paimon.spark.read.VariantExtractionInfo
+import org.apache.paimon.spark.read.{BinPackingSplits, VariantExtractionInfo}
import org.apache.paimon.table.{BucketMode, FileStoreTable, InnerTable}
import org.apache.paimon.table.source.{DataSplit, Split}
@@ -48,15 +48,11 @@ case class PaimonScan(
override val pushedFullTextSearch: Option[FullTextSearch] = None,
override val pushedVariantExtractions: Map[Seq[String],
Seq[VariantExtractionInfo]] = Map.empty,
override val pushedMapSelectedKeys: Map[String, Seq[String]] = Map.empty,
- bucketedScanDisabled: Boolean = false)
+ preserveDataGrouping: Boolean = false)
extends PaimonBaseScan(table)
with SupportsReportPartitioning
with SupportsReportOrdering {
- def disableBucketedScan(): PaimonScan = {
- copy(bucketedScanDisabled = true)
- }
-
@transient
private lazy val extractBucketTransform: Option[Transform] = {
table match {
@@ -118,14 +114,16 @@ case class PaimonScan(
}
private def shouldDoBucketedScan: Boolean = {
- !bucketedScanDisabled && conf.v2BucketingEnabled &&
extractBucketTransform.isDefined
+ preserveDataGrouping && extractBucketTransform.isDefined
}
// Since Spark 3.3
override def outputPartitioning: Partitioning = {
- extractBucketTransform
- .map(bucket => new KeyGroupedPartitioning(Array(bucket),
inputPartitions.size))
- .getOrElse(new UnknownPartitioning(0))
+ if (shouldDoBucketedScan) {
+ new KeyGroupedPartitioning(Array(extractBucketTransform.get),
inputPartitions.size)
+ } else {
+ new UnknownPartitioning(0)
+ }
}
// Since Spark 3.4
@@ -178,13 +176,7 @@ case class PaimonScan(
return super.getInputPartitions(splits)
}
- splits
- .map(_.asInstanceOf[DataSplit])
- .groupBy(_.bucket())
- .map {
- case (bucket, groupedSplits) =>
- PaimonBucketedInputPartition(groupedSplits, bucket)
- }
- .toSeq
+ BinPackingSplits(coreOptions, readRowSizeRatio)
+ .packByBucket(splits.map(_.asInstanceOf[DataSplit]))
}
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
index b32a07e37b..03e4ad8c69 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala
@@ -19,6 +19,7 @@
package org.apache.paimon.spark
import org.apache.paimon.CoreOptions
+import org.apache.paimon.options.Options
import org.apache.paimon.partition.PartitionPredicate
import org.apache.paimon.predicate._
import org.apache.paimon.predicate.SortValue.{NullOrdering, SortDirection}
@@ -30,6 +31,7 @@ import org.apache.spark.sql.connector.expressions
import org.apache.spark.sql.connector.expressions.{NamedReference, SortOrder}
import org.apache.spark.sql.connector.expressions.aggregate.Aggregation
import org.apache.spark.sql.connector.read._
+import org.apache.spark.sql.internal.SQLConf
import scala.collection.JavaConverters._
@@ -164,6 +166,11 @@ class PaimonScanBuilder(val table: InnerTable)
pushedPartitionFilters)
}
+ // Capture the effective layout in the scan's value state, so copies
and query reuse
+ // cannot lose it or recompute it from a later session configuration.
+ val preserveDataGrouping = Options
+ .fromMap(actualTable.options())
+ .get(SparkConnectorOptions.SCAN_PRESERVE_DATA_GROUPING) &&
SQLConf.get.v2BucketingEnabled
PaimonScan(
actualTable,
requiredSchema,
@@ -174,7 +181,8 @@ class PaimonScanBuilder(val table: InnerTable)
vectorSearch,
hybridSearch,
fullTextSearch,
- acceptedVariantExtractions
+ acceptedVariantExtractions,
+ preserveDataGrouping = preserveDataGrouping
)
}
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/adaptive/DisableUnnecessaryPaimonBucketedScan.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/adaptive/DisableUnnecessaryPaimonBucketedScan.scala
deleted file mode 100644
index b0101ded21..0000000000
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/adaptive/DisableUnnecessaryPaimonBucketedScan.scala
+++ /dev/null
@@ -1,178 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.paimon.spark.execution.adaptive
-
-import org.apache.paimon.spark.PaimonScan
-
-import org.apache.spark.sql.catalyst.plans.physical.{AllTuples,
ClusteredDistribution}
-import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.execution._
-import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
-import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
-import org.apache.spark.sql.execution.exchange.{Exchange, ShuffleExchangeLike}
-
-// spotless:off
-/**
- * This rule is inspired from Spark [[DisableUnnecessaryBucketedScan]] but
work for v2 scan.
- *
- * Disable unnecessary bucketed table scan based on actual physical query plan.
- * NOTE: this rule is designed to be applied right after
[[EnsureRequirements]],
- * where all [[ShuffleExchangeLike]] and [[SortExec]] have been added to plan
properly.
- *
- * When BUCKETING_ENABLED and AUTO_BUCKETED_SCAN_ENABLED are set to true, go
through
- * query plan to check where bucketed table scan is unnecessary, and disable
bucketed table
- * scan if:
- *
- * 1. The sub-plan from root to bucketed table scan, does not contain
- * [[hasInterestingPartitionOrOrder]] operator.
- *
- * 2. The sub-plan from the nearest downstream
[[hasInterestingPartitionOrOrder]] operator
- * to the bucketed table scan and at least one [[ShuffleExchangeLike]].
- *
- * Examples:
- * 1. no [[hasInterestingPartitionOrOrder]] operator:
- * Project
- * |
- * Filter
- * |
- * Scan(t1: i, j)
- * (bucketed on column j, DISABLE bucketed scan)
- *
- * 2. join:
- * SortMergeJoin(t1.i = t2.j)
- * / \
- * Sort(i) Sort(j)
- * / \
- * Shuffle(i) Scan(t2: i, j)
- * / (bucketed on column j, enable bucketed scan)
- * Scan(t1: i, j)
- * (bucketed on column j, DISABLE bucketed scan)
- *
- * 3. aggregate:
- * HashAggregate(i, ..., Final)
- * |
- * Shuffle(i)
- * |
- * HashAggregate(i, ..., Partial)
- * |
- * Filter
- * |
- * Scan(t1: i, j)
- * (bucketed on column j, DISABLE bucketed scan)
- *
- * The idea of [[hasInterestingPartitionOrOrder]] is inspired from
"interesting order" in
- * the paper "Access Path Selection in a Relational Database Management System"
- * (https://dl.acm.org/doi/10.1145/582095.582099).
- */
-// spotless:on
-object DisableUnnecessaryPaimonBucketedScan extends Rule[SparkPlan] {
-
- /**
- * Disable bucketed table scan with pre-order traversal of plan.
- *
- * @param hashInterestingPartitionOrOrder
- * The traversed plan has operator with interesting partition and order.
- * @param hasExchange
- * The traversed plan has [[Exchange]] operator.
- */
- private def disableBucketScan(
- plan: SparkPlan,
- hashInterestingPartitionOrOrder: Boolean,
- hasExchange: Boolean): SparkPlan = {
- plan match {
- case p if hasInterestingPartitionOrOrder(p) =>
- // Operator with interesting partition, propagates
`hashInterestingPartitionOrOrder` as true
- // to its children, and resets `hasExchange`.
- p.mapChildren(
- disableBucketScan(_, hashInterestingPartitionOrOrder = true,
hasExchange = false))
- case exchange: ShuffleExchangeLike =>
- // Exchange operator propagates `hasExchange` as true to its child.
- exchange.mapChildren(
- disableBucketScan(_, hashInterestingPartitionOrOrder, hasExchange =
true))
- case batch: BatchScanExec =>
- val paimonBucketedScan = extractPaimonBucketedScan(batch)
- if (paimonBucketedScan.isDefined && (!hashInterestingPartitionOrOrder
|| hasExchange)) {
- val (batch, paimonScan) = paimonBucketedScan.get
- val newBatch = batch.copy(scan = paimonScan.disableBucketedScan())
- newBatch.copyTagsFrom(batch)
- newBatch
- } else {
- batch
- }
- case p if canPassThrough(p) =>
- p.mapChildren(disableBucketScan(_, hashInterestingPartitionOrOrder,
hasExchange))
- case other =>
- other.mapChildren(
- disableBucketScan(_, hashInterestingPartitionOrOrder = false,
hasExchange = false))
- }
- }
-
- private def hasInterestingPartitionOrOrder(plan: SparkPlan): Boolean = {
- val hashPartition = plan.requiredChildDistribution.exists {
- case _: ClusteredDistribution | AllTuples => true
- case _ => false
- }
- // Some operators may only require local sort without distribution,
- // so we do not disable bucketed scan for these queries.
- val hashOrder = plan.requiredChildOrdering.exists(_.nonEmpty)
- hashPartition || hashOrder
- }
-
- /**
- * Check if the operator is allowed single-child operator. We may revisit
this method later as we
- * probably can remove this restriction to allow arbitrary operator between
bucketed table scan
- * and operator with interesting partition.
- */
- private def canPassThrough(plan: SparkPlan): Boolean = {
- plan match {
- case _: ProjectExec | _: FilterExec => true
- case s: SortExec if !s.global => true
- case partialAgg: BaseAggregateExec =>
- partialAgg.requiredChildDistributionExpressions.isEmpty
- case _ => false
- }
- }
-
- def extractPaimonBucketedScan(plan: SparkPlan): Option[(BatchScanExec,
PaimonScan)] =
- plan match {
- case batch: BatchScanExec =>
- batch.scan match {
- case scan: PaimonScan if scan.inputPartitions.forall(_.bucketed) =>
- Some((batch, scan))
- case _ => None
- }
- case _ => None
- }
-
- def apply(plan: SparkPlan): SparkPlan = {
- lazy val hasBucketedScan = plan.exists {
- case p if extractPaimonBucketedScan(p).isDefined => true
- case _ => false
- }
-
- // TODO: replace it with `conf.v2BucketingEnabled` after dropping Spark3.1
- val v2BucketingEnabled =
- conf.getConfString("spark.sql.sources.v2.bucketing.enabled",
"false").toBoolean
- if (!v2BucketingEnabled || !conf.autoBucketedScanEnabled ||
!hasBucketedScan) {
- plan
- } else {
- disableBucketScan(plan, hashInterestingPartitionOrOrder = false,
hasExchange = false)
- }
- }
-}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala
index b216454ba4..504c5ccae7 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala
@@ -23,7 +23,7 @@ import
org.apache.paimon.spark.catalyst.optimizer.{MergePaimonScalarSubqueries,
import
org.apache.paimon.spark.catalyst.plans.logical.PaimonTableValuedFunctions
import org.apache.paimon.spark.commands.BucketExpression
import org.apache.paimon.spark.execution.{OldCompatibleStrategy,
PaimonStrategy}
-import
org.apache.paimon.spark.execution.adaptive.{DisablePostponeCarrierShuffleCoalescing,
DisableUnnecessaryPaimonBucketedScan}
+import
org.apache.paimon.spark.execution.adaptive.DisablePostponeCarrierShuffleCoalescing
import org.apache.spark.sql.SparkSessionExtensions
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
@@ -114,7 +114,6 @@ class PaimonSparkSessionExtensions extends
(SparkSessionExtensions => Unit) {
extensions.injectPlannerStrategy(spark => OldCompatibleStrategy(spark))
// query stage preparation
- extensions.injectQueryStagePrepRule(_ =>
DisableUnnecessaryPaimonBucketedScan)
extensions.injectQueryStagePrepRule(_ =>
DisablePostponeCarrierShuffleCoalescing)
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BinPackingSplits.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BinPackingSplits.scala
index ed331ff617..1d0901564c 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BinPackingSplits.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BinPackingSplits.scala
@@ -21,7 +21,7 @@ package org.apache.paimon.spark.read
import org.apache.paimon.CoreOptions
import org.apache.paimon.CoreOptions._
import org.apache.paimon.io.DataFileMeta
-import org.apache.paimon.spark.PaimonInputPartition
+import org.apache.paimon.spark.{PaimonBucketedInputPartition,
PaimonInputPartition}
import org.apache.paimon.spark.util.SplitUtils
import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplit
import org.apache.paimon.table.format.FormatDataSplit
@@ -82,6 +82,24 @@ case class BinPackingSplits(coreOptions: CoreOptions,
readRowSizeRatio: Double =
}
}
+ /**
+ * Preserve bucket keys without collapsing a whole bucket into one input
partition. Spark can
+ * group these partitions when a distribution is required, or schedule them
separately for a
+ * partially clustered join. Keep each DataSplit intact to preserve
merge-on-read and
+ * data-evolution file groups.
+ */
+ def packByBucket(splits: Array[DataSplit]):
Seq[PaimonBucketedInputPartition] = {
+ if (splits.isEmpty) {
+ return Seq.empty
+ }
+ val maxSplitBytes = computeMaxSplitBytes(splits)
+ splits.groupBy(_.bucket()).toSeq.sortBy(_._1).flatMap {
+ case (bucket, bucketSplits) =>
+ packWholeDataSplits(bucketSplits, maxSplitBytes)
+ .map(group => PaimonBucketedInputPartition(group, bucket))
+ }
+ }
+
private def packDataSplit(splits: Array[DataSplit]):
Array[PaimonInputPartition] = {
val maxSplitBytes = computeMaxSplitBytes(splits)
@@ -142,15 +160,21 @@ case class BinPackingSplits(coreOptions: CoreOptions,
readRowSizeRatio: Double =
}
private def packDataEvolutionSplit(splits: Array[DataSplit]):
Array[PaimonInputPartition] = {
- val maxSplitBytes = computeMaxSplitBytes(splits)
+ packWholeDataSplits(splits, computeMaxSplitBytes(splits))
+ .map(group => PaimonInputPartition(group))
+ .toArray
+ }
+ private def packWholeDataSplits(
+ splits: Array[DataSplit],
+ maxSplitBytes: Long): Seq[Seq[DataSplit]] = {
var currentSize = 0L
val currentSplits = new ArrayBuffer[DataSplit]
- val partitions = new ArrayBuffer[PaimonInputPartition]
+ val partitions = new ArrayBuffer[Seq[DataSplit]]
def closeInputPartition(): Unit = {
if (currentSplits.nonEmpty) {
- partitions += PaimonInputPartition(currentSplits.toArray)
+ partitions += currentSplits.toVector
currentSplits.clear()
currentSize = 0L
}
@@ -173,7 +197,7 @@ case class BinPackingSplits(coreOptions: CoreOptions,
readRowSizeRatio: Double =
}
closeInputPartition()
- partitions.toArray
+ partitions.toSeq
}
private def copyDataSplit(
diff --git
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/BinPackingSplitsTest.scala
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/BinPackingSplitsTest.scala
index 9173681953..7bc8be0864 100644
---
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/BinPackingSplitsTest.scala
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/BinPackingSplitsTest.scala
@@ -150,6 +150,46 @@ class BinPackingSplitsTest extends PaimonSparkTestBase {
}
}
+ test("Paimon: bucketed packing preserves independent split groups and bucket
keys") {
+ withSparkSQLConf("spark.sql.files.minPartitionNum" -> "1") {
+ val splits = Seq(0, 0, 0, 1, 1).zipWithIndex.map {
+ case (bucket, i) =>
+ newDataSplitFromFiles(
+ Seq(newDataFile(s"bucket-$bucket-$i.parquet", 60L)),
+ rawConvertible = false,
+ bucket = bucket)
+ }
+ val packing = BinPackingSplits(CoreOptions.fromMap(
+ Map("source.split.target-size" -> "130 B",
"source.split.open-file-cost" -> "0 B").asJava))
+ val partitions = packing.packByBucket(splits.toArray)
+ assert(partitions.map(_.bucket) == Seq(0, 0, 1))
+ assert(partitions.map(_.splits.size) == Seq(2, 1, 2))
+ partitions.foreach {
+ partition =>
+ assert(partition.splits.forall(_.asInstanceOf[DataSplit].bucket() ==
partition.bucket))
+ }
+ partitions.flatMap(_.splits).zip(splits).foreach {
+ case (actual, original) => Assertions.assertSame(original, actual)
+ }
+ }
+ }
+
+ test("Paimon: bucketed packing keeps oversized merge splits and deletion
files intact") {
+ val split = newDataSplit("merge", Seq(100L, 100L), deletionFileLength =
Some(10L))
+ val packing = BinPackingSplits(
+ CoreOptions.fromMap(
+ Map(
+ "deletion-vectors.enabled" -> "true",
+ "source.split.target-size" -> "50 B",
+ "source.split.open-file-cost" -> "0 B").asJava))
+ val partitions = packing.packByBucket(Array(split))
+ assert(partitions.size == 1)
+ Assertions.assertSame(split, partitions.head.splits.head)
+ assert(split.dataFiles().size() == 2)
+ assert(split.deletionFiles().get().size() == 2)
+ assert(packing.packByBucket(Array.empty[DataSplit]).isEmpty)
+ }
+
test("Paimon: get read splits with column pruning") {
withTable("t") {
sql(
@@ -196,11 +236,12 @@ class BinPackingSplitsTest extends PaimonSparkTestBase {
files: Seq[DataFileMeta],
rawConvertible: Boolean,
deletionFileLength: Option[Long] = None,
- deletionFilePrefix: String = "delete"): DataSplit = {
+ deletionFilePrefix: String = "delete",
+ bucket: Int = 0): DataSplit = {
val builder = DataSplit
.builder()
.withSnapshot(1)
- .withBucket(0)
+ .withBucket(bucket)
.withPartition(BinaryRow.EMPTY_ROW)
.withDataFiles(files.asJava)
.rawConvertible(rawConvertible)
diff --git
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/BucketedScanPlanningTestBase.scala
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/BucketedScanPlanningTestBase.scala
new file mode 100644
index 0000000000..03141a358d
--- /dev/null
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/BucketedScanPlanningTestBase.scala
@@ -0,0 +1,338 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.paimon.spark.sql
+
+import org.apache.paimon.spark.{PaimonScan, PaimonSparkTestBase}
+
+import org.apache.spark.sql.{DataFrame, Row}
+import
org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning,
UnknownPartitioning}
+import org.apache.spark.sql.execution.SortExec
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec,
DataSourceV2ScanRelation}
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
+
+/** The same scan contract is exercised on each Spark 4 minor version. */
+abstract class BucketedScanPlanningTestBase
+ extends PaimonSparkTestBase
+ with AdaptiveSparkPlanHelper {
+
+ private val preserveGrouping = "spark.paimon.scan.preserve-data-grouping"
+ private val v2Bucketing = "spark.sql.sources.v2.bucketing.enabled"
+ private val autoBucketedScan =
"spark.sql.sources.bucketing.autoBucketedScan.enabled"
+ private val partialClustering =
+ "spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled"
+
+ private def createEvents(
+ primaryKey: Boolean = false,
+ partitions: Int = 8,
+ rows: Int = 80,
+ multipleKeys: Boolean = false): Unit = {
+ val pk = if (primaryKey) ", 'primary-key'='ds,id,seq'" else ""
+ val id = if (multipleKeys) "id % 4" else "42"
+ sql(s"""CREATE TABLE t_events (id BIGINT, seq BIGINT, payload STRING, ds
INT)
+ |PARTITIONED BY (ds) TBLPROPERTIES (
+ |'bucket'='16', 'bucket-key'='id', 'source.split.target-size'='1 B',
+ |'source.split.open-file-cost'='0 B' $pk)""".stripMargin)
+ sql(s"""INSERT INTO t_events
+ |SELECT CAST($id AS BIGINT), id, 'old', CAST(id % $partitions AS
INT)
+ |FROM range($rows)""".stripMargin)
+ }
+
+ private def createDimension(multipleKeys: Boolean = false): Unit = {
+ sql("CREATE TABLE t_dim (id BIGINT) TBLPROPERTIES ('bucket'='16',
'bucket-key'='id')")
+ sql(
+ if (multipleKeys) "INSERT INTO t_dim SELECT id FROM range(4)"
+ else "INSERT INTO t_dim VALUES (42)")
+ }
+
+ private def createOrderedTable(): Unit = {
+ sql("""CREATE TABLE t_ordered (id BIGINT, payload STRING)
+ |TBLPROPERTIES ('primary-key'='id', 'bucket'='4',
+ |'source.split.target-size'='128 mb',
'source.split.open-file-cost'='4 mb')
+ |""".stripMargin)
+ sql("INSERT INTO t_ordered SELECT id, concat('v', cast(id AS STRING)) FROM
range(100)")
+ }
+
+ private def eventScan(df: DataFrame): BatchScanExec = {
+ collect(df.queryExecution.executedPlan) {
+ case batch: BatchScanExec
+ if batch.scan.isInstanceOf[PaimonScan] && batch.output.exists(_.name
== "seq") =>
+ batch
+ }.head
+ }
+
+ private def numShuffles(df: DataFrame): Int = {
+ collect(df.queryExecution.executedPlan) { case exchange:
ShuffleExchangeLike => exchange }.size
+ }
+
+ private def checkLayout(scan: PaimonScan, grouped: Boolean): Unit = {
+ assert(scan.inputPartitions.nonEmpty)
+ assert(scan.inputPartitions.forall(_.bucketed == grouped))
+ assert(scan.outputPartitioning.isInstanceOf[KeyGroupedPartitioning] ==
grouped)
+ if (!grouped) {
+ assert(scan.outputPartitioning.isInstanceOf[UnknownPartitioning])
+ assert(scan.outputOrdering().isEmpty)
+ }
+ }
+
+ for ((aqe, force) <- Seq((false, false), (true, false), (true, true)); pk <-
Seq(false, true)) {
+ test(s"Default scanning retains 45 read tasks with aqe=$aqe force=$force
primaryKey=$pk") {
+ withTable("t_events") {
+ createEvents(primaryKey = pk, partitions = 45, rows = 90)
+ withSparkSQLConf(
+ preserveGrouping -> "false",
+ v2Bucketing -> "true",
+ autoBucketedScan -> "true",
+ "spark.sql.adaptive.enabled" -> aqe.toString,
+ "spark.sql.adaptive.forceApply" -> force.toString
+ ) {
+ spark.conf.unset(preserveGrouping)
+ for (explicit <- Seq(None, Some("true"), Some("false"))) {
+ spark.conf.unset(v2Bucketing)
+ explicit.foreach(value => spark.conf.set(v2Bucketing, value))
+ val df = sql("""SELECT id, seq FROM t_events
+ |WHERE id = 42 AND length(payload) > 0
+ |ORDER BY seq DESC LIMIT 20""".stripMargin)
+ checkAnswer(df, (70L until 90L).reverse.map(seq => Row(42L, seq)))
+ val batch = eventScan(df)
+ val scan = batch.scan.asInstanceOf[PaimonScan]
+ checkLayout(scan, grouped = false)
+ assert(scan.inputSplits.length == 45)
+ assert(batch.inputRDD.getNumPartitions == 45)
+ }
+ }
+ }
+ }
+ }
+
+ test("Table, session and read options select the scan layout") {
+ withTable("t_events") {
+ createEvents()
+ sql("ALTER TABLE t_events SET TBLPROPERTIES
('scan.preserve-data-grouping'='true')")
+ withSparkSQLConf(v2Bucketing -> "true", preserveGrouping -> "false") {
+ spark.conf.unset(preserveGrouping)
+ val tableDefault = sql("SELECT id, seq FROM t_events")
+ tableDefault.collect()
+ checkLayout(eventScan(tableDefault).scan.asInstanceOf[PaimonScan],
grouped = true)
+ assert(eventScan(tableDefault).inputRDD.getNumPartitions == 1)
+
+ spark.conf.set(preserveGrouping, "false")
+ val sessionOverride = sql("SELECT id, seq FROM t_events")
+ sessionOverride.collect()
+ checkLayout(eventScan(sessionOverride).scan.asInstanceOf[PaimonScan],
grouped = false)
+ assert(eventScan(sessionOverride).inputRDD.getNumPartitions == 8)
+
+ val readOverride = spark.read
+ .option("scan.preserve-data-grouping", "true")
+ .table("t_events")
+ .select("id", "seq")
+ readOverride.collect()
+ checkLayout(eventScan(readOverride).scan.asInstanceOf[PaimonScan],
grouped = true)
+ assert(eventScan(readOverride).inputRDD.getNumPartitions == 1)
+ }
+ }
+ }
+
+ test("A scan fixes its layout before physical planning and ignores later
configuration changes") {
+ withTable("t_ordered") {
+ createOrderedTable()
+ for (preserve <- Seq(false, true); v2 <- Seq(false, true)) {
+ withSparkSQLConf(preserveGrouping -> preserve.toString, v2Bucketing ->
v2.toString) {
+ val df = sql("SELECT id FROM t_ordered")
+ val scan = df.queryExecution.optimizedPlan.collectFirst {
+ case relation: DataSourceV2ScanRelation if
relation.scan.isInstanceOf[PaimonScan] =>
+ relation.scan.asInstanceOf[PaimonScan]
+ }.get
+ withSparkSQLConf(
+ preserveGrouping -> (!preserve).toString,
+ v2Bucketing -> (!v2).toString) {
+ checkLayout(scan, preserve && v2)
+ assert(scan.outputOrdering().nonEmpty == (preserve && v2))
+ checkLayout(scan.copy(), preserve && v2)
+ assert(scan.copy().outputOrdering().nonEmpty == (preserve && v2))
+ }
+ }
+ }
+ }
+ }
+
+ test("Scans with different layouts cannot be considered equal for query
reuse") {
+ withTable("t_ordered") {
+ createOrderedTable()
+ withSparkSQLConf(preserveGrouping -> "true", v2Bucketing -> "false") {
+ def plannedScan(): PaimonScan = {
+ sql("SELECT id FROM
t_ordered").queryExecution.optimizedPlan.collectFirst {
+ case relation: DataSourceV2ScanRelation if
relation.scan.isInstanceOf[PaimonScan] =>
+ relation.scan.asInstanceOf[PaimonScan]
+ }.get
+ }
+ val ungrouped = plannedScan()
+ spark.conf.set(v2Bucketing, "true")
+ val grouped = plannedScan()
+ checkLayout(ungrouped, grouped = false)
+ checkLayout(grouped, grouped = true)
+ assert(ungrouped != grouped)
+ }
+ }
+ }
+
+ for (aqe <- Seq(false, true)) {
+ test(s"TopN and SORT BY stay correct in both scan modes with aqe=$aqe") {
+ withTable("t_ordered") {
+ createOrderedTable()
+ withSparkSQLConf(
+ v2Bucketing -> "true",
+ "spark.sql.adaptive.enabled" -> aqe.toString,
+ "spark.sql.adaptive.forceApply" -> aqe.toString,
+ "spark.sql.files.minPartitionNum" -> "1") {
+ for (preserve <- Seq(false, true); auto <- Seq(false, true)) {
+ withSparkSQLConf(
+ preserveGrouping -> preserve.toString,
+ autoBucketedScan -> auto.toString) {
+ val topN = sql("SELECT id FROM t_ordered ORDER BY id LIMIT 5")
+ assert(topN.collect().map(_.getLong(0)).toSeq == (0L until 5L))
+ val sorted = sql("SELECT id FROM t_ordered SORT BY id")
+ val parts = sorted.rdd
+ .mapPartitions(rows =>
Iterator(rows.map(_.getLong(0)).toVector))
+ .collect()
+ assert(parts.flatten.sorted.toSeq == (0L until 100L))
+ assert(parts.forall(p => p == p.sorted))
+ assert(collect(sorted.queryExecution.executedPlan) {
+ case sort: SortExec => sort
+ }.isEmpty == preserve)
+ }
+ }
+ }
+ }
+ }
+ }
+
+ for (aqe <- Seq(false, true); preserve <- Seq(false, true)) {
+ test(
+ s"Generate and Sample retain aggregation and join semantics with
aqe=$aqe grouped=$preserve") {
+ withTable("t_events", "t_dim") {
+ createEvents()
+ createDimension()
+ withSparkSQLConf(
+ preserveGrouping -> preserve.toString,
+ v2Bucketing -> "true",
+ autoBucketedScan -> "true",
+ partialClustering -> "false",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ "spark.sql.shuffle.partitions" -> "4",
+ "spark.sql.adaptive.enabled" -> aqe.toString,
+ "spark.sql.adaptive.forceApply" -> aqe.toString
+ ) {
+ val generated = sql("""SELECT id, count(*) FROM t_events
+ |LATERAL VIEW explode(array(1, 2)) v AS x
GROUP BY id
+ |""".stripMargin)
+ checkAnswer(generated, Seq(Row(42L, 160L)))
+ assert((numShuffles(generated) == 0) == preserve)
+ val sampled =
+ sql("SELECT id, count(*) FROM t_events TABLESAMPLE (100 PERCENT)
GROUP BY id")
+ checkAnswer(sampled, Seq(Row(42L, 80L)))
+ assert((numShuffles(sampled) == 0) == preserve)
+ val joined = sql("""SELECT /*+ MERGE(e, d) */ e.id, e.seq, e.x
+ |FROM (SELECT id, seq, explode(array(1, 2)) AS x
FROM t_events) e
+ |JOIN t_dim d ON e.id = d.id""".stripMargin)
+ checkAnswer(joined, for (seq <- 0L until 80L; x <- 1 to 2) yield
Row(42L, seq, x))
+ assert((numShuffles(joined) == 0) == preserve)
+ }
+ }
+ }
+ }
+
+ for (aqe <- Seq(false, true); partial <- Seq(false, true)) {
+ test(s"Explicit grouping retains SPJ read units with aqe=$aqe
partial=$partial") {
+ withTable("t_events", "t_dim") {
+ createEvents()
+ createDimension()
+ withSparkSQLConf(
+ preserveGrouping -> "true",
+ v2Bucketing -> "true",
+ autoBucketedScan -> "true",
+ partialClustering -> partial.toString,
+ "spark.sql.sources.v2.bucketing.pushPartValues.enabled" -> "true",
+ "spark.sql.requireAllClusterKeysForCoPartition" -> "false",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ "spark.sql.adaptive.enabled" -> aqe.toString
+ ) {
+ val df = sql("""SELECT /*+ MERGE(e, d) */ e.id, e.seq
+ |FROM t_events e JOIN t_dim d ON e.id =
d.id""".stripMargin)
+ checkAnswer(df, (0L until 80L).map(seq => Row(42L, seq)))
+ val batch = eventScan(df)
+ assert(batch.scan.asInstanceOf[PaimonScan].inputPartitions.size == 8)
+ assert(batch.inputRDD.getNumPartitions == (if (partial) 8 else 1))
+ assert(numShuffles(df) == 0)
+ }
+ }
+ }
+ }
+
+ test("Partially clustered joins retain rows across multiple bucket keys") {
+ withTable("t_events", "t_dim") {
+ createEvents(multipleKeys = true)
+ createDimension(multipleKeys = true)
+ withSparkSQLConf(
+ preserveGrouping -> "true",
+ v2Bucketing -> "true",
+ partialClustering -> "true",
+ "spark.sql.sources.v2.bucketing.pushPartValues.enabled" -> "true",
+ "spark.sql.requireAllClusterKeysForCoPartition" -> "false",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ "spark.sql.adaptive.enabled" -> "false"
+ ) {
+ val df = sql("""SELECT /*+ MERGE(e, d) */ e.id, e.seq
+ |FROM t_events e JOIN t_dim d ON e.id =
d.id""".stripMargin)
+ checkAnswer(df, (0L until 80L).map(seq => Row(seq % 4, seq)))
+ assert(eventScan(df).inputRDD.getNumPartitions > 1)
+ assert(numShuffles(df) == 0)
+ }
+ }
+ }
+
+ test("Both scan modes preserve primary-key updates and deletes") {
+ withTable("t_events", "t_dim") {
+ createEvents(primaryKey = true)
+ createDimension()
+ sql(
+ "INSERT INTO t_events SELECT CAST(42 AS BIGINT), id, 'new', CAST(id %
8 AS INT) FROM range(8)")
+ sql("DELETE FROM t_events WHERE seq = 3")
+ val expected =
+ (0L until 80L).filter(_ != 3L).map(seq => Row(42L, seq, if (seq < 8L)
"new" else "old"))
+ for (preserve <- Seq(false, true)) {
+ withSparkSQLConf(
+ preserveGrouping -> preserve.toString,
+ v2Bucketing -> "true",
+ partialClustering -> "true",
+ "spark.sql.sources.v2.bucketing.pushPartValues.enabled" -> "true",
+ "spark.sql.requireAllClusterKeysForCoPartition" -> "false",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ "spark.sql.adaptive.enabled" -> "false"
+ ) {
+ val df = sql("""SELECT /*+ MERGE(e, d) */ e.id, e.seq, e.payload
+ |FROM t_events e JOIN t_dim d ON e.id =
d.id""".stripMargin)
+ checkAnswer(df, expected)
+ checkLayout(eventScan(df).scan.asInstanceOf[PaimonScan], preserve)
+ assert((numShuffles(df) == 0) == preserve)
+ }
+ }
+ }
+ }
+}
diff --git
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/DisableUnnecessaryPaimonBucketedScanSuite.scala
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/BucketedScanPolicyTest.scala
similarity index 69%
rename from
paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/DisableUnnecessaryPaimonBucketedScanSuite.scala
rename to
paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/BucketedScanPolicyTest.scala
index f47d40285a..41a2b0c735 100644
---
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/DisableUnnecessaryPaimonBucketedScanSuite.scala
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/BucketedScanPolicyTest.scala
@@ -19,45 +19,41 @@
package org.apache.paimon.spark.sql
import org.apache.paimon.spark.PaimonSparkTestBase
-import
org.apache.paimon.spark.execution.adaptive.DisableUnnecessaryPaimonBucketedScan
import org.apache.spark.SparkConf
-import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
-class DisableUnnecessaryPaimonBucketedScanSuite
- extends PaimonSparkTestBase
- with AdaptiveSparkPlanHelper {
+class BucketedScanPolicyTest extends PaimonSparkTestBase with
AdaptiveSparkPlanHelper {
override protected def sparkConf: SparkConf = {
- // Make non-shuffle query work with stage preparation rule
+ // Exercise adaptive planning even for queries without a shuffle.
super.sparkConf
.set("spark.sql.adaptive.forceApply", "true")
}
- private def checkDisableBucketedScan(
- query: String,
- expectedNumScanWithAutoScanEnabled: Int,
- expectedNumScanWithAutoScanDisabled: Int): Unit = {
-
- def checkNumBucketedScan(df: DataFrame, expectedNumBucketedScan: Int):
Unit = {
- val plan = df.queryExecution.executedPlan
- val bucketedScan = collect(plan) {
- case p if
DisableUnnecessaryPaimonBucketedScan.extractPaimonBucketedScan(p).isDefined => p
- }
- assert(bucketedScan.length == expectedNumBucketedScan, query)
- }
-
+ private def checkScanModes(query: String, expectedNumBucketedScan: Int):
Unit = {
withSparkSQLConf("spark.sql.sources.v2.bucketing.enabled" -> "true") {
- withSparkSQLConf("spark.sql.sources.bucketing.autoBucketedScan.enabled"
-> "true") {
- val df = sql(query)
- val result = df.collect()
- checkNumBucketedScan(df, expectedNumScanWithAutoScanEnabled)
-
-
withSparkSQLConf("spark.sql.sources.bucketing.autoBucketedScan.enabled" ->
"false") {
- val expected = sql(query)
- checkAnswer(expected, result)
- checkNumBucketedScan(expected, expectedNumScanWithAutoScanDisabled)
+ var expected: Array[org.apache.spark.sql.Row] = null
+ for (preserve <- Seq(false, true)) {
+ withSparkSQLConf("spark.paimon.scan.preserve-data-grouping" ->
preserve.toString) {
+ val df = sql(query)
+ val result = df.collect()
+ if (!preserve) {
+ expected = result
+ } else {
+ checkAnswer(df, expected.toSeq)
+ }
+ val scans = collect(df.queryExecution.executedPlan) {
+ case scan: BatchScanExec
+ if scan.scan.isInstanceOf[org.apache.paimon.spark.PaimonScan]
&&
+ scan.scan
+ .asInstanceOf[org.apache.paimon.spark.PaimonScan]
+ .inputPartitions
+ .forall(_.bucketed) =>
+ scan
+ }
+ assert(scans.size == (if (preserve) expectedNumBucketedScan else 0),
query)
}
}
}
@@ -78,7 +74,7 @@ class DisableUnnecessaryPaimonBucketedScanSuite
"INSERT INTO t3 VALUES (1, 1, 'x1'), (2, 2, 'x3'), (3, 3, 'x3'), (4, 4,
'x4'), (5, 5, 'x5')")
}
- test("Disable unnecessary bucketed table scan - basic test") {
+ test("Scan modes preserve primary-key results - basic test") {
assume(gteqSpark3_3)
withTable("t1", "t2", "t3") {
@@ -86,41 +82,41 @@ class DisableUnnecessaryPaimonBucketedScanSuite
Seq(
// Read bucketed table
- ("SELECT * FROM t1", 0, 1),
- ("SELECT i FROM t1", 0, 1),
- ("SELECT j FROM t1", 0, 0),
+ ("SELECT * FROM t1", 1),
+ ("SELECT i FROM t1", 1),
+ ("SELECT j FROM t1", 0),
// Filter on bucketed column
- ("SELECT * FROM t1 WHERE i = 1", 0, 1),
+ ("SELECT * FROM t1 WHERE i = 1", 1),
// Filter on non-bucketed column
- ("SELECT * FROM t1 WHERE j = 1", 0, 1),
+ ("SELECT * FROM t1 WHERE j = 1", 1),
// Join with same buckets
- ("SELECT /*+ broadcast(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.i", 0, 2),
- ("SELECT /*+ shuffle_hash(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.i", 2,
2),
- ("SELECT /*+ merge(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.i", 2, 2),
+ ("SELECT /*+ broadcast(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.i", 2),
+ ("SELECT /*+ shuffle_hash(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.i", 2),
+ ("SELECT /*+ merge(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.i", 2),
// Join with different buckets
- ("SELECT /*+ broadcast(t1)*/ * FROM t1 JOIN t3 ON t1.i = t3.i", 0, 2),
- ("SELECT /*+ shuffle_hash(t1)*/ * FROM t1 JOIN t3 ON t1.i = t3.i", 0,
2),
- ("SELECT /*+ merge(t1)*/ * FROM t1 JOIN t3 ON t1.i = t3.i", 0, 2),
+ ("SELECT /*+ broadcast(t1)*/ * FROM t1 JOIN t3 ON t1.i = t3.i", 2),
+ ("SELECT /*+ shuffle_hash(t1)*/ * FROM t1 JOIN t3 ON t1.i = t3.i", 2),
+ ("SELECT /*+ merge(t1)*/ * FROM t1 JOIN t3 ON t1.i = t3.i", 2),
// Join on non-bucketed column
- ("SELECT /*+ broadcast(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.j", 0, 2),
- ("SELECT /*+ shuffle_hash(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.j", 0,
2),
- ("SELECT /*+ merge(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.j", 0, 2),
- ("SELECT /*+ broadcast(t1)*/ * FROM t1 JOIN t2 ON t1.j = t2.j", 0, 2),
- ("SELECT /*+ shuffle_hash(t1)*/ * FROM t1 JOIN t2 ON t1.j = t2.j", 0,
2),
- ("SELECT /*+ merge(t1)*/ * FROM t1 JOIN t2 ON t1.j = t2.j", 0, 2),
+ ("SELECT /*+ broadcast(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.j", 2),
+ ("SELECT /*+ shuffle_hash(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.j", 2),
+ ("SELECT /*+ merge(t1)*/ * FROM t1 JOIN t2 ON t1.i = t2.j", 2),
+ ("SELECT /*+ broadcast(t1)*/ * FROM t1 JOIN t2 ON t1.j = t2.j", 2),
+ ("SELECT /*+ shuffle_hash(t1)*/ * FROM t1 JOIN t2 ON t1.j = t2.j", 2),
+ ("SELECT /*+ merge(t1)*/ * FROM t1 JOIN t2 ON t1.j = t2.j", 2),
// Aggregate on bucketed column
- ("SELECT SUM(i) FROM t1 GROUP BY i", 1, 1),
+ ("SELECT SUM(i) FROM t1 GROUP BY i", 1),
// Aggregate on non-bucketed column
- ("SELECT SUM(i) FROM t1 GROUP BY j", 0, 1),
- ("SELECT j, SUM(i), COUNT(j) FROM t1 GROUP BY j", 0, 1)
+ ("SELECT SUM(i) FROM t1 GROUP BY j", 1),
+ ("SELECT j, SUM(i), COUNT(j) FROM t1 GROUP BY j", 1)
).foreach {
- case (query, numScanWithAutoScanEnabled, numScanWithAutoScanDisabled)
=>
- checkDisableBucketedScan(query, numScanWithAutoScanEnabled,
numScanWithAutoScanDisabled)
+ case (query, numBucketedScans) =>
+ checkScanModes(query, numBucketedScans)
}
}
}
- test("Disable unnecessary bucketed table scan - multiple joins test") {
+ test("Scan modes preserve primary-key results - multiple joins test") {
assume(gteqSpark3_3)
withTable("t1", "t2", "t3") {
@@ -133,28 +129,24 @@ class DisableUnnecessaryPaimonBucketedScanSuite
SELECT /*+ broadcast(t1, t3)*/ * FROM t1 JOIN t2 JOIN t3
ON t1.i = t2.i AND t2.i = t3.i
""".stripMargin,
- 0,
3),
(
"""
SELECT /*+ broadcast(t1) merge(t3)*/ * FROM t1 JOIN t2 JOIN t3
ON t1.i = t2.i AND t2.i = t3.i
""".stripMargin,
- 0,
3),
(
"""
SELECT /*+ merge(t1) broadcast(t3)*/ * FROM t1 JOIN t2 JOIN t3
ON t1.i = t2.i AND t2.i = t3.i
""".stripMargin,
- 2,
3),
(
"""
SELECT /*+ merge(t1, t3)*/ * FROM t1 LEFT JOIN t2 LEFT JOIN t3
ON t1.i = t2.i AND t2.i = t3.i
""".stripMargin,
- 0,
3),
// Multiple joins on non-bucketed columns
(
@@ -162,30 +154,27 @@ class DisableUnnecessaryPaimonBucketedScanSuite
SELECT /*+ broadcast(t1, t3)*/ * FROM t1 JOIN t2 JOIN t3
ON t1.i = t2.j AND t2.j = t3.i
""".stripMargin,
- 0,
3),
(
"""
SELECT /*+ merge(t1, t3)*/ * FROM t1 JOIN t2 JOIN t3
ON t1.i = t2.j AND t2.j = t3.i
""".stripMargin,
- 0,
3),
(
"""
SELECT /*+ merge(t1, t3)*/ * FROM t1 JOIN t2 JOIN t3
ON t1.j = t2.j AND t2.j = t3.j
""".stripMargin,
- 0,
3)
).foreach {
- case (query, numScanWithAutoScanEnabled, numScanWithAutoScanDisabled)
=>
- checkDisableBucketedScan(query, numScanWithAutoScanEnabled,
numScanWithAutoScanDisabled)
+ case (query, numBucketedScans) =>
+ checkScanModes(query, numBucketedScans)
}
}
}
- test("Disable unnecessary bucketed table scan - other operators test") {
+ test("Scan modes preserve primary-key results - other operators test") {
assume(gteqSpark3_3)
withTable("t1", "t2", "t3") {
@@ -199,7 +188,6 @@ class DisableUnnecessaryPaimonBucketedScanSuite
UNION ALL
(SELECT t2.i FROM t2 GROUP BY t2.i)
""".stripMargin,
- 1,
2),
// Non-allowed operator in sub-plan
(
@@ -208,7 +196,6 @@ class DisableUnnecessaryPaimonBucketedScanSuite
FROM (SELECT t1.i FROM t1 UNION ALL SELECT t2.i FROM t2)
GROUP BY i
""".stripMargin,
- 0,
2),
// Multiple [[Exchange]] in sub-plan
(
@@ -216,7 +203,6 @@ class DisableUnnecessaryPaimonBucketedScanSuite
SELECT j, SUM(i), COUNT(*) FROM t1 GROUP BY j
DISTRIBUTE BY j
""".stripMargin,
- 0,
1),
(
"""
@@ -224,7 +210,6 @@ class DisableUnnecessaryPaimonBucketedScanSuite
FROM (SELECT i, j FROM t1 DISTRIBUTE BY i, j)
GROUP BY j
""".stripMargin,
- 0,
1),
// No bucketed table scan in plan
(
@@ -233,11 +218,10 @@ class DisableUnnecessaryPaimonBucketedScanSuite
FROM (SELECT t1.j FROM t1 JOIN t3 ON t1.j = t3.j)
GROUP BY j
""".stripMargin,
- 0,
0)
).foreach {
- case (query, numScanWithAutoScanEnabled, numScanWithAutoScanDisabled)
=>
- checkDisableBucketedScan(query, numScanWithAutoScanEnabled,
numScanWithAutoScanDisabled)
+ case (query, numBucketedScans) =>
+ checkScanModes(query, numBucketedScans)
}
}
}
@@ -258,10 +242,9 @@ class DisableUnnecessaryPaimonBucketedScanSuite
|""".stripMargin)
val df = spark.sql("select sum(id) from t1 where id is not null")
assert(df.count() == 1)
- checkDisableBucketedScan(
+ checkScanModes(
query = "SELECT SUM(id) FROM t1 WHERE id is not null",
- expectedNumScanWithAutoScanEnabled = 1,
- expectedNumScanWithAutoScanDisabled = 1)
+ expectedNumBucketedScan = 1)
}
}
}
diff --git
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/BucketedTableQueryTest.scala
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/BucketedTableQueryTest.scala
index 745d498b6f..7eaf27f957 100644
---
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/BucketedTableQueryTest.scala
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/BucketedTableQueryTest.scala
@@ -39,6 +39,7 @@ class BucketedTableQueryTest extends PaimonSparkTestBase with
AdaptiveSparkPlanH
}
withSparkSQLConf(
"spark.sql.sources.v2.bucketing.enabled" -> "true",
+ "spark.paimon.scan.preserve-data-grouping" -> "true",
"spark.sql.requireAllClusterKeysForCoPartition" -> "false",
"spark.sql.autoBroadcastJoinThreshold" -> "-1"
) {
