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 92404efb76 [spark] Add opt-in repartitioning for oversized Blob table
scans (#10457)
92404efb76 is described below
commit 92404efb769dabd7d53e1bfe59df7041975ab9d2
Author: Zhou <[email protected]>
AuthorDate: Wed Oct 14 10:03:40 2026 +0800
[spark] Add opt-in repartitioning for oversized Blob table scans (#10457)
---
docs/docs/spark/sql-query.md | 49 +++
docs/generated/spark_connector_configuration.html | 6 +
.../apache/paimon/spark/SparkConnectorOptions.java | 12 +
.../optimizer/RepartitionLargeBlobScan.scala | 131 +++++++
.../extensions/PaimonSparkSessionExtensions.scala | 8 +-
.../spark/RepartitionLargeBlobScanTest.scala | 407 +++++++++++++++++++++
6 files changed, 612 insertions(+), 1 deletion(-)
diff --git a/docs/docs/spark/sql-query.md b/docs/docs/spark/sql-query.md
index 71883eeae0..bd3a12f368 100644
--- a/docs/docs/spark/sql-query.md
+++ b/docs/docs/spark/sql-query.md
@@ -167,6 +167,55 @@ To inspect row kinds, including deletes, query the
[`audit_log` system table](..
## Query Optimization
+### Repartition Oversized Scan Output
+
+This optimization is disabled by default. Enable it for the Spark session with:
+
+```sql
+SET spark.paimon.read.repartition-large-blob-scan.enabled = true;
+```
+
+Set it to `false` to disable `RepartitionLargeBlobScan`. When disabled, the
rule
+returns the plan unchanged without planning scan splits or input partitions.
+
+This rule only applies to tables whose full schema contains a `BLOB`,
`ARRAY<BLOB>`,
+or `MAP<..., BLOB>` column, even if that column is pruned from the scan. Other
tables
+are skipped without planning scan splits or input partitions.
+
+When an eligible Paimon batch scan has an input partition larger than twice
`filesMaxPartitionBytes`,
+the optimizer inserts a shuffle immediately after the scan. The trigger
threshold and
+output partition count are calculated separately:
+
+```text
+threshold = 2 * filesMaxPartitionBytes
+output partition count = ceil(total input partition bytes /
filesMaxPartitionBytes)
+```
+
+If the computed output partition count is less than or equal to the existing
input partition
+count, the rule keeps the scan unchanged without adding a shuffle. A shuffle
is inserted only
+when the computed count increases parallelism.
+
+`filesMaxPartitionBytes` uses Paimon's `source.split.target-size`, falling
back to an explicitly
+configured `spark.sql.files.maxPartitionBytes`, then Paimon's default split
size.
+All input partitions contribute to the total bytes. An input partition exactly
at the
+threshold does not trigger repartitioning. File-open costs and minimum
partition counts
+do not affect this rule's threshold or output partition count.
+
+Sizes are the sum of data file sizes reported by the splits in each input
partition,
+including Blob file sizes. They are metadata estimates, not decoded row sizes
or
+measured shuffle bytes. The rule increases parallelism for downstream
operators;
+it does not split source scan tasks. It leaves streaming scans and an existing
+shuffle directly above the scan unchanged. Runtime filtering can subsequently
+reduce the scan input without changing this statically planned repartition
count.
+
+The shuffle is placed above an existing dynamic partition pruning filter,
keeping the
+filter adjacent to its scan operation so Spark can pass runtime filters to
Paimon.
+Scans below expressions using `input_file_name()`, `input_file_block_start()`,
or
+`input_file_block_length()` are skipped: these functions need the source
reader's
+task-local file context, which is not carried through a shuffle.
+
+### Filter Pushdown
+
It is highly recommended to specify partition and primary key filters
along with the query, which will speed up the data skipping of the query.
diff --git a/docs/generated/spark_connector_configuration.html
b/docs/generated/spark_connector_configuration.html
index 61684d295f..281ac494c3 100644
--- a/docs/generated/spark_connector_configuration.html
+++ b/docs/generated/spark_connector_configuration.html
@@ -56,6 +56,12 @@ under the License.
<td>Boolean</td>
<td>Whether to read row in the form of changelog (add rowkind
column in row to represent its change type).</td>
</tr>
+ <tr>
+ <td><h5>read.repartition-large-blob-scan.enabled</h5></td>
+ <td style="word-wrap: break-word;">false</td>
+ <td>Boolean</td>
+ <td>Whether to insert a shuffle after a batch scan of a table with
BLOB columns (including ARRAY<BLOB> and MAP values of BLOB) when any
input partition exceeds twice filesMaxPartitionBytes. The shuffle partition
count is the total input bytes divided by filesMaxPartitionBytes, rounded up.
Skips the shuffle if this count is no greater than the existing input partition
count. Does not change the number of scan tasks.</td>
+ </tr>
<tr>
<td><h5>read.stream.maxBytesPerTrigger</h5></td>
<td style="word-wrap: break-word;">(none)</td>
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 82d6855589..7368a7a03f 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
@@ -168,6 +168,18 @@ public class SparkConnectorOptions {
.withDescription(
"Whether to allow full scan when reading a
partitioned table.");
+ public static final ConfigOption<Boolean>
READ_REPARTITION_LARGE_BLOB_SCAN_ENABLED =
+ key("read.repartition-large-blob-scan.enabled")
+ .booleanType()
+ .defaultValue(false)
+ .withDescription(
+ "Whether to insert a shuffle after a batch scan of
a table with BLOB columns "
+ + "(including ARRAY<BLOB> and MAP values
of BLOB) when any input partition "
+ + "exceeds twice filesMaxPartitionBytes.
The shuffle partition count is "
+ + "the total input bytes divided by
filesMaxPartitionBytes, rounded up. "
+ + "Skips the shuffle if this count is no
greater than the existing input partition count. "
+ + "Does not change the number of scan
tasks.");
+
public static final ConfigOption<Boolean>
FORMAT_TABLE_REPAIR_COLLECT_STATISTICS =
key("format-table.repair.collect-statistics")
.booleanType()
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/RepartitionLargeBlobScan.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/RepartitionLargeBlobScan.scala
new file mode 100644
index 0000000000..6d127823f7
--- /dev/null
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/RepartitionLargeBlobScan.scala
@@ -0,0 +1,131 @@
+/*
+ * 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.catalyst.optimizer
+
+import org.apache.paimon.spark.{PaimonScan, SparkConnectorOptions}
+import org.apache.paimon.spark.read.BinPackingSplits
+import org.apache.paimon.spark.util.{OptionUtils, SplitUtils}
+import org.apache.paimon.types.BlobType
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.expressions.{DynamicPruning,
InputFileBlockLength, InputFileBlockStart, InputFileName}
+import org.apache.spark.sql.catalyst.planning.PhysicalOperation
+import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan,
Repartition, RepartitionOperation}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation
+import org.apache.spark.sql.internal.SQLConf
+
+import scala.collection.JavaConverters._
+
+/** Redistributes oversized scan partitions for downstream operators,
preserving the scan itself. */
+object RepartitionLargeBlobScan extends Rule[LogicalPlan] {
+
+ // injectOptimizerRule runs before V2 scan pushdown. The final user-provided
optimizer batch
+ // runs after it, so register there as well when Spark initializes its
optimizer rules.
+ def register(spark: SparkSession): Unit = {
+ val experimental = spark.experimental
+ if (!experimental.extraOptimizations.contains(this)) {
+ experimental.extraOptimizations = experimental.extraOptimizations :+ this
+ }
+ }
+
+ override def apply(plan: LogicalPlan): LogicalPlan = {
+ if (
+ !OptionUtils
+
.getOptionString(SparkConnectorOptions.READ_REPARTITION_LARGE_BLOB_SCAN_ENABLED)
+ .toBoolean
+ ) {
+ return plan
+ }
+
+ def repartitionScan(node: LogicalPlan, relation:
DataSourceV2ScanRelation): LogicalPlan = {
+ if (relation.isStreaming) {
+ return node
+ }
+ relation.scan match {
+ case scan: PaimonScan
+ if scan.table
+ .rowType()
+ .getFields
+ .asScala
+ .exists(field => BlobType.isBlobFileField(field.`type`())) =>
+ val partitions = scan.inputPartitions
+ val targetSize =
BinPackingSplits.filesMaxPartitionBytes(scan.coreOptions, SQLConf.get)
+ val threshold = BigInt(targetSize) * 2
+ val sizes = partitions.map {
+ partition => partition.splits.map(split =>
BigInt(SplitUtils.splitSize(split))).sum
+ }
+ if (targetSize > 0 && sizes.exists(_ > threshold)) {
+ val numPartitions = (sizes.sum + targetSize - 1) / targetSize
+ require(
+ numPartitions <= Int.MaxValue,
+ s"Scan repartition count $numPartitions exceeds ${Int.MaxValue};
" +
+ "increase the scan split target size"
+ )
+ if (numPartitions <= partitions.size) {
+ node
+ } else {
+ Repartition(numPartitions.toInt, shuffle = true, node)
+ }
+ } else {
+ node
+ }
+ case _ => node
+ }
+ }
+
+ def rewrite(node: LogicalPlan): LogicalPlan = {
+ // These expressions read task-local file context, which cannot cross a
shuffle.
+ // Conservatively leave their input subtree unchanged, including through
aliases.
+ if (hasInputFileExpression(node)) {
+ return node
+ }
+ node match {
+ // Preserve explicit shuffles and remain idempotent for a scan with
pruning filters.
+ case repartition: RepartitionOperation
+ if repartition.shuffle &&
scanRelation(repartition.child).isDefined =>
+ repartition
+ case filter: Filter if
filter.condition.exists(_.isInstanceOf[DynamicPruning]) =>
+ // DataSourceV2Strategy must still extract DPP from this scan
operation. If the
+ // operation cannot be extracted, leave it intact instead of
inserting below it.
+ scanRelation(filter) match {
+ case Some(relation) if !filter.exists(hasInputFileExpression) =>
+ repartitionScan(filter, relation)
+ case _ => filter
+ }
+ case relation: DataSourceV2ScanRelation => repartitionScan(relation,
relation)
+ case other => other.mapChildren(rewrite)
+ }
+ }
+
+ rewrite(plan)
+ }
+
+ private def hasInputFileExpression(plan: LogicalPlan): Boolean = {
+ plan.expressions.exists(_.exists {
+ case _: InputFileName | _: InputFileBlockStart | _: InputFileBlockLength
=> true
+ case _ => false
+ })
+ }
+
+ private def scanRelation(plan: LogicalPlan):
Option[DataSourceV2ScanRelation] = plan match {
+ case PhysicalOperation(_, _, relation: DataSourceV2ScanRelation) =>
Some(relation)
+ case _ => None
+ }
+}
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 504c5ccae7..e7eb99d31b 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
@@ -19,7 +19,7 @@
package org.apache.paimon.spark.extensions
import org.apache.paimon.spark.catalyst.analysis.{PaimonAnalysis,
PaimonDeleteTable, PaimonFunctionResolver, PaimonIncompatibleResolutionRules,
PaimonMergeInto, PaimonPostHocResolutionRules, PaimonProcedureResolver,
PaimonUpdateTable, PaimonViewResolver, ReplacePaimonFunctions}
-import
org.apache.paimon.spark.catalyst.optimizer.{MergePaimonScalarSubqueries,
OptimizeMetadataOnlyDeleteFromPaimonTable, PushDownArrayPredicates,
PushDownLateralVectorSearchFilter, RepartitionLateralVectorSearchInput}
+import
org.apache.paimon.spark.catalyst.optimizer.{MergePaimonScalarSubqueries,
OptimizeMetadataOnlyDeleteFromPaimonTable, PushDownArrayPredicates,
PushDownLateralVectorSearchFilter, RepartitionLargeBlobScan,
RepartitionLateralVectorSearchInput}
import
org.apache.paimon.spark.catalyst.plans.logical.PaimonTableValuedFunctions
import org.apache.paimon.spark.commands.BucketExpression
import org.apache.paimon.spark.execution.{OldCompatibleStrategy,
PaimonStrategy}
@@ -108,6 +108,12 @@ class PaimonSparkSessionExtensions extends
(SparkSessionExtensions => Unit) {
extensions.injectOptimizerRule(_ => RepartitionLateralVectorSearchInput)
extensions.injectOptimizerRule(_ => PushDownLateralVectorSearchFilter)
+ extensions.injectOptimizerRule {
+ spark =>
+ RepartitionLargeBlobScan.register(spark)
+ RepartitionLargeBlobScan
+ }
+
// planner extensions
extensions.injectPlannerStrategy(spark => PaimonStrategy(spark))
// old compatible
diff --git
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/RepartitionLargeBlobScanTest.scala
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/RepartitionLargeBlobScanTest.scala
new file mode 100644
index 0000000000..3376ae7413
--- /dev/null
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/RepartitionLargeBlobScanTest.scala
@@ -0,0 +1,407 @@
+/*
+ * 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
+
+import org.apache.paimon.data.BinaryRow
+import org.apache.paimon.io.DataFileMeta
+import org.apache.paimon.manifest.FileSource
+import org.apache.paimon.spark.catalyst.optimizer.RepartitionLargeBlobScan
+import org.apache.paimon.table.InnerTable
+import org.apache.paimon.table.source.{DataSplit, Split}
+import org.apache.paimon.types.{DataField, DataType, DataTypes, RowType}
+
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.catalyst.expressions.{Alias, And,
DynamicPruningExpression, Expression, GreaterThan, InputFileBlockLength,
InputFileBlockStart, InputFileName, IsNotNull, Literal}
+import org.apache.spark.sql.catalyst.plans.logical.{Filter, Project,
Repartition}
+import org.apache.spark.sql.connector.catalog.{Table => ConnectorTable,
TableCapability}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec,
DataSourceV2Relation, DataSourceV2ScanRelation}
+import org.apache.spark.sql.types.{IntegerType, StructType}
+import org.mockito.Mockito.{mock, when}
+
+import java.util.Collections
+
+import scala.collection.JavaConverters._
+
+/** Rule tests with synthetic metadata and query regressions for file context
and runtime pruning. */
+class RepartitionLargeBlobScanTest extends PaimonSparkTestBase {
+
+ private val enabledKey =
+
s"spark.paimon.${SparkConnectorOptions.READ_REPARTITION_LARGE_BLOB_SCAN_ENABLED.key()}"
+
+ test("input-file expressions in projections and filters skip scan planning")
{
+ withSplitConf {
+ Seq[Expression](InputFileName(), InputFileBlockStart(),
InputFileBlockLength()).foreach {
+ expression =>
+ val scan = relation(Seq(Seq(120L)), failOnPlanning = true)
+ val projection = Project(Seq(Alias(expression, "file_info")()), scan)
+ assert(RepartitionLargeBlobScan(projection) eq projection)
+ val filter = Filter(IsNotNull(expression), scan)
+ assert(RepartitionLargeBlobScan(filter) eq filter)
+ val aliasProjection = Project(projection.output, projection)
+ assert(RepartitionLargeBlobScan(aliasProjection) == aliasProjection)
+ }
+ }
+ }
+
+ test("shuffle stays above dynamic pruning filters and remains idempotent") {
+ withSplitConf {
+ val scan = relation(Seq(Seq(120L), Seq(30L)))
+ val filter = Filter(
+ And(DynamicPruningExpression(Literal(true)),
GreaterThan(scan.output.head, Literal(0))),
+ scan)
+ val result = RepartitionLargeBlobScan(filter).asInstanceOf[Repartition]
+ assert(result.numPartitions == 3)
+ assert(result.child eq filter)
+ assert(RepartitionLargeBlobScan(result) eq result)
+
+ val fileProjection = Project(Seq(Alias(InputFileName(), "file_name")()),
scan)
+ val fileFilter = Filter(DynamicPruningExpression(Literal(true)),
fileProjection)
+ assert(RepartitionLargeBlobScan(fileFilter) eq fileFilter)
+ }
+ }
+
+ test("input_file_name projection and filter preserve source file context in
real queries") {
+ withSparkSQLConf(
+ "spark.sql.adaptive.enabled" -> "false",
+ "spark.paimon.write.use-v2-write" -> "false") {
+ withTable("file_context") {
+ sql("""CREATE TABLE file_context (id INT) TBLPROPERTIES (
+ |'row-tracking.enabled'='true', 'data-evolution.enabled'='true',
+ |'source.split.target-size'='4kb')""".stripMargin)
+ sql("INSERT INTO file_context SELECT /*+ COALESCE(1) */ CAST(id AS
INT) FROM range(5000)")
+ sql("ALTER TABLE file_context ADD COLUMN payload BINARY COMMENT
'__BLOB_FIELD'")
+
+ withSparkSQLConf(enabledKey -> "true") {
+ assert(sql("SELECT id FROM
file_context").queryExecution.optimizedPlan.exists {
+ case _: Repartition => true
+ case _ => false
+ })
+ }
+ val queries = Seq(
+ "SELECT id, input_file_name(), input_file_block_start(),
input_file_block_length() " +
+ "FROM file_context",
+ "SELECT id FROM file_context WHERE length(input_file_name()) > 0"
+ )
+ queries.zipWithIndex.foreach {
+ case (query, index) =>
+ var expected = Seq.empty[Row]
+ withSparkSQLConf(enabledKey -> "false") {
+ expected = sql(query).collect().toSeq
+ }
+ assert(expected.size == 5000)
+ if (index == 0) {
+ assert(expected.forall(_.getString(1).nonEmpty))
+ }
+ withSparkSQLConf(enabledKey -> "true") {
+ val result = sql(query)
+
assert(!result.queryExecution.optimizedPlan.exists(_.isInstanceOf[Repartition]))
+ checkAnswer(result, expected)
+ }
+ }
+ }
+ }
+ }
+
+ test("real Blob join retains runtime partition filters below the added
shuffle") {
+ withSparkSQLConf(
+ "spark.sql.adaptive.enabled" -> "false",
+ "spark.paimon.write.use-v2-write" -> "false",
+ "spark.sql.optimizer.dynamicPartitionPruning.enabled" -> "true",
+ "spark.sql.optimizer.dynamicPartitionPruning.useStats" -> "false",
+ "spark.sql.optimizer.dynamicPartitionPruning.reuseBroadcastOnly" ->
"false",
+ "spark.sql.optimizer.dynamicPartitionPruning.fallbackFilterRatio" ->
"1.0"
+ ) {
+ withTable("pruning_fact", "pruning_dim") {
+ sql("""CREATE TABLE pruning_fact (id INT, pt INT, payload BINARY)
PARTITIONED BY (pt)
+ |TBLPROPERTIES ('row-tracking.enabled'='true',
'data-evolution.enabled'='true',
+ |'blob-field'='payload',
'source.split.target-size'='64kb')""".stripMargin)
+ sql("""INSERT INTO pruning_fact SELECT /*+ COALESCE(1) */
+ |CAST(id AS INT), CAST(id DIV 100 AS INT), CAST(repeat('x',
8192) AS BINARY)
+ |FROM range(300)""".stripMargin)
+ sql("CREATE TABLE pruning_dim (pt INT, keep INT)")
+ sql("INSERT INTO pruning_dim VALUES (0, 1), (1, 0), (2, 0)")
+ val query = """SELECT /*+ BROADCAST(d) */ f.id, length(f.payload)
+ |FROM pruning_fact f JOIN pruning_dim d ON f.pt = d.pt
+ |WHERE d.keep = 1""".stripMargin
+ Seq(false, true).foreach {
+ enabled =>
+ withSparkSQLConf(enabledKey -> enabled.toString) {
+ val result = sql(query)
+ if (enabled) {
+
assert(result.queryExecution.optimizedPlan.exists(_.isInstanceOf[Repartition]))
+ }
+ val factScan = result.queryExecution.executedPlan.collectFirst {
+ case scan: BatchScanExec if scan.output.exists(_.name ==
"payload") => scan
+ }.get
+ assert(factScan.runtimeFilters.nonEmpty,
result.queryExecution.executedPlan.toString)
+ checkAnswer(result, (0 until 100).map(id => Row(id, 8192)))
+ assert(factScan.metrics("numOutputRows").value == 100)
+ }
+ }
+ }
+ }
+ }
+
+ test("non-Blob tables are skipped without planning scan partitions") {
+ withSplitConf {
+ Seq(None, Some(DataTypes.BYTES())).foreach {
+ blobType =>
+ val scan = relation(Seq(Seq(120L)), failOnPlanning = true, blobType
= blobType)
+ assert(RepartitionLargeBlobScan(scan) eq scan)
+ }
+ }
+ }
+
+ test("Blob tables remain eligible after Blob columns are pruned") {
+ withSplitConf {
+ Seq(
+ DataTypes.BLOB(),
+ DataTypes.ARRAY(DataTypes.BLOB()),
+ DataTypes.MAP(DataTypes.STRING(), DataTypes.BLOB())).foreach {
+ blobType =>
+ val scan = relation(Seq(Seq(120L), Seq(30L)), blobType =
Some(blobType))
+ assert(scan.scan.readSchema().fieldNames.toSeq == Seq("a"))
+
assert(RepartitionLargeBlobScan(scan).asInstanceOf[Repartition].numPartitions
== 3)
+ }
+ }
+ }
+
+ test("disabled by default without planning scan partitions") {
+ assert(spark.conf.getOption(enabledKey).isEmpty)
+ val scan = relation(Seq(Seq(120L)), failOnPlanning = true)
+ assert(RepartitionLargeBlobScan(scan) eq scan)
+ }
+
+ test("explicitly disabling the rule avoids planning scan partitions") {
+ withSplitConf {
+ withSparkSQLConf(enabledKey -> "false") {
+ val scan = relation(Seq(Seq(120L)), failOnPlanning = true)
+ assert(RepartitionLargeBlobScan(scan) eq scan)
+ }
+ val enabled = relation(Seq(Seq(120L), Seq(30L)))
+
assert(RepartitionLargeBlobScan(enabled).asInstanceOf[Repartition].numPartitions
== 3)
+ }
+ }
+
+ test("partition size includes all splits and exact threshold does not
trigger") {
+ withSplitConf {
+ val fits = relation(Seq(Seq(50L, 50L), Seq(80L)))
+ assert(RepartitionLargeBlobScan(fits) eq fits)
+ val oversized = relation(Seq(Seq(60L, 60L), Seq(30L)))
+ val result =
RepartitionLargeBlobScan(oversized).asInstanceOf[Repartition]
+ assert(result.numPartitions == 3)
+ assert(result.shuffle)
+ assert(result.child eq oversized)
+ assert(RepartitionLargeBlobScan(result) eq result)
+ }
+ }
+
+ test("repartition count includes small partitions and rounds up") {
+ withSplitConf {
+ val result = RepartitionLargeBlobScan(relation(Seq(Seq(101L), Seq(25L),
Seq(25L))))
+ assert(result.asInstanceOf[Repartition].numPartitions == 4)
+ }
+ }
+
+ test("minimum partition count does not lower the repartition threshold") {
+ withSplitConf {
+ withSparkSQLConf(
+ "spark.sql.files.maxPartitionBytes" -> "1000b",
+ "spark.sql.files.minPartitionNum" -> "4") {
+ val scan = relation(Seq(Seq(80L), Seq(20L)))
+ assert(RepartitionLargeBlobScan(scan) eq scan)
+ }
+ }
+ }
+
+ test(
+ "skip shuffle when a large Blob partition and small partitions already
have more parallelism") {
+ withSplitConf {
+ val scan = relation(
+ Seq(Seq(256L * 1024)) ++ Seq.fill(20)(Seq(1L)),
+ Map("source.split.target-size" -> "64kb"))
+ // The oversized partition triggers the size check, but ceil(262164 /
65536) = 5
+ // would reduce the existing 21 input partitions. Preserve the scan
without a shuffle.
+ assert(RepartitionLargeBlobScan(scan) eq scan)
+ }
+ }
+
+ test("skip shuffle when computed partition count equals existing input
partition count") {
+ withSplitConf {
+ val scan = relation(
+ Seq(Seq(256L * 1024)) ++ Seq.fill(4)(Seq(1L)),
+ Map("source.split.target-size" -> "64kb"))
+ assert(RepartitionLargeBlobScan(scan) eq scan)
+ }
+ }
+
+ test("open file cost does not change the threshold or shuffle partition
count") {
+ withSplitConf {
+ withSparkSQLConf(
+ "spark.sql.files.openCostInBytes" -> "1000b",
+ "spark.sql.files.minPartitionNum" -> "10") {
+ val fits = relation(Seq(Seq(80L), Seq(20L)))
+ assert(RepartitionLargeBlobScan(fits) eq fits)
+ val result = RepartitionLargeBlobScan(relation(Seq(Seq(101L),
Seq(50L))))
+ assert(result.asInstanceOf[Repartition].numPartitions == 4)
+ }
+ }
+ }
+
+ test("Paimon split settings take precedence over Spark file settings") {
+ withSplitConf {
+ val scan = relation(
+ Seq(Seq(80L), Seq(20L)),
+ Map("source.split.target-size" -> "25b", "source.split.open-file-cost"
-> "0b"))
+
assert(RepartitionLargeBlobScan(scan).asInstanceOf[Repartition].numPartitions
== 4)
+ }
+ }
+
+ test("preserve an existing shuffle and leave empty scans unchanged") {
+ withSplitConf {
+ val existing = Repartition(7, shuffle = true, relation(Seq(Seq(120L))))
+ assert(RepartitionLargeBlobScan(existing) eq existing)
+ val empty = relation(Seq.empty)
+ assert(RepartitionLargeBlobScan(empty) eq empty)
+ }
+ }
+
+ test("doubling the threshold does not overflow Long") {
+ withSplitConf {
+ withSparkSQLConf("spark.sql.files.maxPartitionBytes" ->
s"${Long.MaxValue}b") {
+ val fits = relation(Seq(Seq(Long.MaxValue, Long.MaxValue)))
+ assert(RepartitionLargeBlobScan(fits) eq fits)
+ val oversized = relation(Seq(Seq(Long.MaxValue, Long.MaxValue,
Long.MaxValue)))
+
assert(RepartitionLargeBlobScan(oversized).asInstanceOf[Repartition].numPartitions
== 3)
+ }
+ }
+ }
+
+ test("reject a partition count that does not fit Spark's integer limit") {
+ withSplitConf {
+ withSparkSQLConf("spark.sql.files.maxPartitionBytes" -> "1b") {
+ val error = intercept[IllegalArgumentException] {
+ RepartitionLargeBlobScan(relation(Seq(Seq(Long.MaxValue))))
+ }
+ assert(error.getMessage.contains("exceeds"))
+ }
+ }
+ }
+
+ test("rule is registered once in an optimizer batch after scan pushdown") {
+ val batches = spark.sessionState.optimizer.batches
+ val pushdownIndex = batches.indexWhere(_.name == "Early Filter and
Projection Push-Down")
+ val lateIndex = batches.indexWhere {
+ batch =>
+ batch.rules.contains(RepartitionLargeBlobScan) &&
+ batch.name == "User Provided Optimizers"
+ }
+ assert(pushdownIndex >= 0)
+ assert(lateIndex > pushdownIndex)
+ RepartitionLargeBlobScan.register(spark)
+ RepartitionLargeBlobScan.register(spark)
+ assert(spark.experimental.extraOptimizations.count(_ ==
RepartitionLargeBlobScan) == 1)
+ withSplitConf {
+ val scan = relation(Seq(Seq(120L), Seq(30L)))
+ val result = batches(lateIndex).rules.foldLeft(
+ scan: org.apache.spark.sql.catalyst.plans.logical.LogicalPlan)((plan,
rule) => rule(plan))
+ assert(result.asInstanceOf[Repartition].numPartitions == 3)
+ }
+ }
+
+ test("non-Paimon plans are unchanged") {
+ withSplitConf {
+ val plan = spark.range(10).queryExecution.optimizedPlan
+ assert(RepartitionLargeBlobScan(plan) eq plan)
+ }
+ }
+
+ private def withSplitConf(f: => Unit): Unit = {
+ withSparkSQLConf(
+ enabledKey -> "true",
+ "spark.sql.files.maxPartitionBytes" -> "50b",
+ "spark.sql.files.minPartitionNum" -> "1",
+ "spark.sql.files.openCostInBytes" -> "0b")(f)
+ }
+
+ private def relation(
+ partitionFileSizes: Seq[Seq[Long]],
+ options: Map[String, String] = Map.empty,
+ failOnPlanning: Boolean = false,
+ blobType: Option[DataType] = Some(DataTypes.BLOB())):
DataSourceV2ScanRelation = {
+ val schema = new StructType().add("a", IntegerType)
+ val table = mock(classOf[InnerTable])
+ when(table.options()).thenReturn(options.asJava)
+ val fields = Seq(new DataField(0, "a", DataTypes.INT())) ++
+ blobType.map(dataType => new DataField(1, "payload", dataType))
+ when(table.rowType()).thenReturn(new RowType(fields.asJava))
+ when(table.partitionKeys()).thenReturn(Collections.emptyList[String]())
+ when(table.primaryKeys()).thenReturn(Collections.emptyList[String]())
+
+ val partitions = partitionFileSizes.zipWithIndex.map {
+ case (fileSizes, partitionIndex) =>
+ val splits = fileSizes.zipWithIndex.map {
+ case (size, fileIndex) =>
+ val file = DataFileMeta.forAppend(
+ s"$partitionIndex-$fileIndex.parquet",
+ size,
+ 1,
+ null,
+ 0,
+ 0,
+ 1,
+ Collections.emptyList[String](),
+ null,
+ FileSource.APPEND,
+ null,
+ null,
+ null,
+ null)
+ DataSplit
+ .builder()
+ .withSnapshot(1)
+ .withPartition(BinaryRow.EMPTY_ROW)
+ .withBucket(0)
+ .withBucketPath("test")
+ .withDataFiles(Collections.singletonList(file))
+ .rawConvertible(false)
+ .build()
+ }
+ PaimonInputPartition(splits)
+ }
+ val scan = new PaimonScan(table, schema, Seq.empty, Seq.empty, None, None,
None) {
+ override protected def getInputSplits: Array[Split] = {
+ assert(!failOnPlanning, "Skipped scans must not plan scan splits")
+ partitions.flatMap(_.splits).toArray
+ }
+ override protected def getInputPartitions(splits: Array[Split]):
Seq[PaimonInputPartition] = {
+ assert(!failOnPlanning, "Skipped scans must not plan scan partitions")
+ partitions
+ }
+ }
+ val sparkTable = new ConnectorTable {
+ override def name(): String = "test_scan"
+ override def schema(): StructType = scan.requiredSchema
+ override def capabilities(): java.util.Set[TableCapability] =
+ Collections.emptySet[TableCapability]()
+ }
+ val source = DataSourceV2Relation.create(sparkTable, None, None)
+ DataSourceV2ScanRelation(source, scan, source.output)
+ }
+}
