From fe8389c920feed6746209cecd6cb693139f1ee7e Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Mon, 27 Jul 2026 21:30:32 +0800 Subject: [PATCH 1/6] [spark] Infer postpone bucket count from row count --- .../primary-key-table/data-distribution.md | 20 +++-- docs/generated/core_configuration.html | 2 +- .../java/org/apache/paimon/CoreOptions.java | 2 +- .../spark/commands/PaimonSparkWriter.scala | 78 ++++++++++++++++--- .../spark/sql/PostponeBucketTableTest.scala | 70 +++++++++++++++++ 5 files changed, 156 insertions(+), 16 deletions(-) diff --git a/docs/docs/primary-key-table/data-distribution.md b/docs/docs/primary-key-table/data-distribution.md index 49846b2ea3b8..bbc41cdcde7d 100644 --- a/docs/docs/primary-key-table/data-distribution.md +++ b/docs/docs/primary-key-table/data-distribution.md @@ -70,12 +70,22 @@ Postpone bucket mode is configured by `'bucket' = '-2'`. This mode aims to solve the difficulty to determine a fixed number of buckets and support different buckets for different partitions. -When writing records into the table, -all records will first be stored in the `bucket-postpone` directory of each partition +By default, batch writes set `postpone.batch-write-fixed-bucket` to `true` +and write records directly to real buckets. +For a Spark batch write to a partition without real bucket files, +you can configure `postpone.target-row-num-per-bucket` to calculate the bucket number +from the incoming row count plus the row count of existing postpone files. +The calculated bucket number is +`ceil(row_count / postpone.target-row-num-per-bucket)`, is at least `1`, +and is limited by `postpone.batch-write-fixed-bucket.max-parallelism`. +Without this target, Spark uses the number of input writers, also limited by the maximum parallelism. +Inferring from row counts adds a counting stage; Spark caches the batch so that counting and writing +use the same input rows. + +When `postpone.batch-write-fixed-bucket` is `false`, +records are first stored in the `bucket-postpone` directory of each partition and are not available to readers. - -To move the records into the correct bucket and make them readable, -you need to run a compaction job. +To move these records into the correct bucket and make them readable, run a compaction job. See `compact` [procedure](../flink/procedures). The bucket number for the partitions compacted for the first time is configured by the option `postpone.default-bucket-num`, whose default value is `1`. diff --git a/docs/generated/core_configuration.html b/docs/generated/core_configuration.html index 4d1d8f87f7b1..f0eb67698ad0 100644 --- a/docs/generated/core_configuration.html +++ b/docs/generated/core_configuration.html @@ -1294,7 +1294,7 @@
postpone.target-row-num-per-bucket
(none) Long - Target row number per bucket for partitions compacted from postpone bucket files for the first time. + Target row number per bucket when batch writing fixed buckets or compacting postpone bucket files for a partition without real bucket data.
primary-key
diff --git a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java index 3040a981a4fc..6934109718ef 100644 --- a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java +++ b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java @@ -2733,7 +2733,7 @@ public String toString() { .longType() .noDefaultValue() .withDescription( - "Target row number per bucket for partitions compacted from postpone bucket files for the first time."); + "Target row number per bucket when batch writing fixed buckets or compacting postpone bucket files for a partition without real bucket data."); public static final ConfigOption GLOBAL_INDEX_ROW_COUNT_PER_SHARD = key("global-index.row-count-per-shard") diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala index 1518a77934f6..f3d0379e36f6 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala @@ -110,6 +110,8 @@ case class PaimonSparkWriter( val uriReaderFactory = uriReaderFactoryForBlobDescriptor import sparkSession.implicits._ + val inferPostponeBucketNum = + postponeBatchWriteFixedBucket && coreOptions.postponeTargetRowNumPerBucket().isPresent val withInitBucketCol = bucketMode match { case BUCKET_UNAWARE => data case KEY_DYNAMIC if !data.schema.fieldNames.contains(ROW_KIND_COL) => @@ -118,18 +120,49 @@ case class PaimonSparkWriter( .withColumn(BUCKET_COL, lit(-1)) case _ => data.withColumn(BUCKET_COL, lit(-1)) } + if (inferPostponeBucketNum) { + withInitBucketCol.persist() + } + val rowKindColIdx = SparkRowUtils.getFieldIndex(withInitBucketCol.schema, ROW_KIND_COL) val bucketColIdx = SparkRowUtils.getFieldIndex(withInitBucketCol.schema, BUCKET_COL) val encoderGroupWithBucketCol = EncoderSerDeGroup(withInitBucketCol.schema) val postponePartitionBucketComputer: Option[BinaryRow => Integer] = - if (postponeBatchWriteFixedBucket) { - val knownNumBuckets = PostponeUtils.getKnownNumBuckets(table) - val defaultPostponeNumBuckets = Math.min( - withInitBucketCol.rdd.getNumPartitions, - table.coreOptions().postponeBatchWriteFixedBucketMaxParallelism) - Some((p: BinaryRow) => knownNumBuckets.getOrDefault(p, defaultPostponeNumBuckets)) - } else { - None + try { + if (postponeBatchWriteFixedBucket) { + val knownNumBuckets = PostponeUtils.getKnownNumBuckets(table) + val maxNumBuckets = table.coreOptions().postponeBatchWriteFixedBucketMaxParallelism + val defaultPostponeNumBuckets = + Math.min(withInitBucketCol.rdd.getNumPartitions, maxNumBuckets) + val inferredNumBuckets: Map[BinaryRow, Int] = + if (inferPostponeBucketNum) { + val targetRowNum = coreOptions.postponeTargetRowNumPerBucket().get() + val postponeRowCounts = PostponeUtils.getPostponeRowCounts(table) + countRowsByPartition(withInitBucketCol).map { + case (partition, rowCount) => + partition -> Math.min( + PostponeUtils.computeBucketNumByRowCount( + Math.addExact(rowCount, postponeRowCounts.getOrDefault(partition, 0L)), + targetRowNum), + maxNumBuckets) + } + } else { + Map.empty[BinaryRow, Int] + } + Some( + (p: BinaryRow) => + knownNumBuckets.getOrDefault( + p, + Integer.valueOf(inferredNumBuckets.getOrElse(p, defaultPostponeNumBuckets)))) + } else { + None + } + } catch { + case e: Throwable => + if (inferPostponeBucketNum) { + withInitBucketCol.unpersist() + } + throw e } def newWrite() = PaimonDataWrite( @@ -349,7 +382,13 @@ case class PaimonSparkWriter( throw new UnsupportedOperationException(s"Spark doesn't support $bucketMode mode.") } - WriteTaskResult.merge(written.collect()) + try { + WriteTaskResult.merge(written.collect()) + } finally { + if (inferPostponeBucketNum) { + withInitBucketCol.unpersist() + } + } } /** @@ -536,6 +575,27 @@ case class PaimonSparkWriter( .toSeq } + private def countRowsByPartition(df: DataFrame): Map[BinaryRow, Long] = { + val partitionColumns = partitionCols(df) + if (partitionColumns.isEmpty) { + Map(BinaryRow.EMPTY_ROW -> df.count()) + } else { + val partitionType = tableSchema.logicalPartitionType() + val serializer = InternalSerializers.create(partitionType) + df.groupBy(partitionColumns: _*) + .count() + .collect() + .map { + row => + val partitionRow = Row.fromSeq(row.toSeq.dropRight(1)) + val partition = + serializer.toBinaryRow(new SparkRow(partitionType, partitionRow)).copy() + partition -> row.getLong(row.length - 1) + } + .toMap + } + } + private def deserializeCommitMessage( serializer: CommitMessageSerializer, bytes: Array[Byte]): CommitMessage = { diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala index 805b747bb6c4..943ae01c7f9c 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala @@ -91,6 +91,76 @@ class PostponeBucketTableTest extends PaimonSparkTestBase { } } + test("Postpone bucket table: infer bucket number from incoming row count") { + withTable("t") { + sql(""" + |CREATE TABLE t ( + | k INT, + | v STRING, + | pt INT + |) PARTITIONED BY (pt) + |TBLPROPERTIES ( + | 'primary-key' = 'k, pt', + | 'bucket' = '-2', + | 'postpone.batch-write-fixed-bucket' = 'true', + | 'postpone.target-row-num-per-bucket' = '200' + |) + |""".stripMargin) + + sql(""" + |INSERT INTO t SELECT /*+ REPARTITION(20) */ + |id AS k, + |CAST(id AS STRING) AS v, + |CASE WHEN id < 100 THEN 0 ELSE 1 END AS pt + |FROM range (0, 550) + |""".stripMargin) + + checkAnswer( + sql("SELECT distinct(bucket) FROM `t$buckets` WHERE partition = '{0}' ORDER BY bucket"), + Seq(Row(0)) + ) + checkAnswer( + sql("SELECT distinct(bucket) FROM `t$buckets` WHERE partition = '{1}' ORDER BY bucket"), + Seq(Row(0), Row(1), Row(2)) + ) + + // Existing partitions keep their bucket number even when the new data volume changes. + sql(""" + |INSERT INTO t SELECT /*+ REPARTITION(20) */ + |id AS k, + |CAST(id AS STRING) AS v, + |0 AS pt + |FROM range (1000, 2000) + |""".stripMargin) + checkAnswer( + sql("SELECT distinct(bucket) FROM `t$buckets` WHERE partition = '{0}' ORDER BY bucket"), + Seq(Row(0)) + ) + + // Postpone rows are included when a partition gets real buckets for the first time. + withSparkSQLConf("spark.paimon.postpone.batch-write-fixed-bucket" -> "false") { + sql(""" + |INSERT INTO t SELECT + |id AS k, + |CAST(id AS STRING) AS v, + |2 AS pt + |FROM range (2000, 2150) + |""".stripMargin) + } + sql(""" + |INSERT INTO t SELECT /*+ REPARTITION(20) */ + |id AS k, + |CAST(id AS STRING) AS v, + |2 AS pt + |FROM range (3000, 3100) + |""".stripMargin) + checkAnswer( + sql("SELECT distinct(bucket) FROM `t$buckets` WHERE partition = '{2}' ORDER BY bucket"), + Seq(Row(-2), Row(0), Row(1)) + ) + } + } + test("Postpone bucket table: write fix bucket then write postpone bucket") { withTable("t") { sql(""" From 4b829d690b881ec13b47fd141aa727f339b75dfa Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Mon, 27 Jul 2026 21:55:42 +0800 Subject: [PATCH 2/6] [spark] Infer postpone bucket count from data size --- .../primary-key-table/data-distribution.md | 20 ++- docs/generated/core_configuration.html | 6 + .../java/org/apache/paimon/CoreOptions.java | 12 ++ .../spark/commands/PaimonSparkWriter.scala | 137 +++++++++++++----- .../spark/sql/PostponeBucketTableTest.scala | 73 +++++++++- 5 files changed, 200 insertions(+), 48 deletions(-) diff --git a/docs/docs/primary-key-table/data-distribution.md b/docs/docs/primary-key-table/data-distribution.md index bbc41cdcde7d..2bb26974bbb1 100644 --- a/docs/docs/primary-key-table/data-distribution.md +++ b/docs/docs/primary-key-table/data-distribution.md @@ -73,14 +73,18 @@ and support different buckets for different partitions. By default, batch writes set `postpone.batch-write-fixed-bucket` to `true` and write records directly to real buckets. For a Spark batch write to a partition without real bucket files, -you can configure `postpone.target-row-num-per-bucket` to calculate the bucket number -from the incoming row count plus the row count of existing postpone files. -The calculated bucket number is -`ceil(row_count / postpone.target-row-num-per-bucket)`, is at least `1`, -and is limited by `postpone.batch-write-fixed-bucket.max-parallelism`. -Without this target, Spark uses the number of input writers, also limited by the maximum parallelism. -Inferring from row counts adds a counting stage; Spark caches the batch so that counting and writing -use the same input rows. +Spark calculates the bucket number from the uncompressed Paimon `BinaryRow` serialized size. +The target size is configured by `postpone.target-size-per-bucket` and defaults to `1 GB`. +Existing postpone files do not record their uncompressed serialized size, so Spark estimates their +size using their row count and the average serialized size of incoming rows in the same partition. +You can instead configure `postpone.target-row-num-per-bucket` to calculate the bucket number +from row counts; this option takes precedence over the target size. +The calculated bucket number is at least `1` and is limited by +`postpone.batch-write-fixed-bucket.max-parallelism`. +The serialized size is measured before file encoding and compression, so it is not the exact +ORC or Parquet size on storage. +Inferring the data amount adds a statistics stage; Spark caches the batch so that inference and +writing use the same input rows. When `postpone.batch-write-fixed-bucket` is `false`, records are first stored in the `bucket-postpone` directory of each partition diff --git a/docs/generated/core_configuration.html b/docs/generated/core_configuration.html index f0eb67698ad0..7bc238d4ff23 100644 --- a/docs/generated/core_configuration.html +++ b/docs/generated/core_configuration.html @@ -1296,6 +1296,12 @@ Long Target row number per bucket when batch writing fixed buckets or compacting postpone bucket files for a partition without real bucket data. + +
postpone.target-size-per-bucket
+ 1 gb + MemorySize + Target uncompressed serialized data size per bucket when Spark batch writes fixed buckets for a partition without real bucket data. This option is ignored when 'postpone.target-row-num-per-bucket' is configured. +
primary-key
(none) diff --git a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java index 6934109718ef..ec95ab6ec537 100644 --- a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java +++ b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java @@ -2735,6 +2735,14 @@ public String toString() { .withDescription( "Target row number per bucket when batch writing fixed buckets or compacting postpone bucket files for a partition without real bucket data."); + public static final ConfigOption POSTPONE_TARGET_SIZE_PER_BUCKET = + key("postpone.target-size-per-bucket") + .memoryType() + .defaultValue(MemorySize.parse("1 gb")) + .withDescription( + "Target uncompressed serialized data size per bucket when Spark batch writes fixed buckets for a partition without real bucket data. " + + "This option is ignored when 'postpone.target-row-num-per-bucket' is configured."); + public static final ConfigOption GLOBAL_INDEX_ROW_COUNT_PER_SHARD = key("global-index.row-count-per-shard") .longType() @@ -4406,6 +4414,10 @@ public Optional postponeTargetRowNumPerBucket() { return options.getOptional(POSTPONE_TARGET_ROW_NUM_PER_BUCKET); } + public long postponeTargetSizePerBucket() { + return options.get(POSTPONE_TARGET_SIZE_PER_BUCKET).getBytes(); + } + public long globalIndexRowCountPerShard() { return options.get(GLOBAL_INDEX_ROW_COUNT_PER_SHARD); } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala index f3d0379e36f6..db8589062d75 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala @@ -52,6 +52,7 @@ import java.io.IOException import java.util.Collections.singletonMap import scala.collection.JavaConverters._ +import scala.collection.mutable case class PaimonSparkWriter( table: FileStoreTable, @@ -59,6 +60,8 @@ case class PaimonSparkWriter( batchId: Option[Long] = None) extends WriteHelper { + import PaimonSparkWriter._ + private lazy val tableSchema = table.schema private lazy val bucketMode = table.bucketMode @@ -110,8 +113,7 @@ case class PaimonSparkWriter( val uriReaderFactory = uriReaderFactoryForBlobDescriptor import sparkSession.implicits._ - val inferPostponeBucketNum = - postponeBatchWriteFixedBucket && coreOptions.postponeTargetRowNumPerBucket().isPresent + val inferPostponeBucketNum = postponeBatchWriteFixedBucket val withInitBucketCol = bucketMode match { case BUCKET_UNAWARE => data case KEY_DYNAMIC if !data.schema.fieldNames.contains(ROW_KIND_COL) => @@ -134,21 +136,29 @@ case class PaimonSparkWriter( val maxNumBuckets = table.coreOptions().postponeBatchWriteFixedBucketMaxParallelism val defaultPostponeNumBuckets = Math.min(withInitBucketCol.rdd.getNumPartitions, maxNumBuckets) - val inferredNumBuckets: Map[BinaryRow, Int] = - if (inferPostponeBucketNum) { - val targetRowNum = coreOptions.postponeTargetRowNumPerBucket().get() - val postponeRowCounts = PostponeUtils.getPostponeRowCounts(table) - countRowsByPartition(withInitBucketCol).map { - case (partition, rowCount) => - partition -> Math.min( + val targetRowNum = coreOptions.postponeTargetRowNumPerBucket() + val postponeRowCounts = PostponeUtils.getPostponeRowCounts(table) + val dataStats = + collectDataStatsByPartition(withInitBucketCol, collectSize = !targetRowNum.isPresent) + val inferredNumBuckets: Map[BinaryRow, Int] = dataStats.map { + case (partition, stats) => + val postponeRowCount = postponeRowCounts.getOrDefault(partition, 0L) + val numBuckets = + if (targetRowNum.isPresent) { + Math.min( PostponeUtils.computeBucketNumByRowCount( - Math.addExact(rowCount, postponeRowCounts.getOrDefault(partition, 0L)), - targetRowNum), + Math.addExact(stats.rowCount, postponeRowCount), + targetRowNum.get()), maxNumBuckets) - } - } else { - Map.empty[BinaryRow, Int] - } + } else { + computeBucketNumBySize( + stats, + postponeRowCount, + coreOptions.postponeTargetSizePerBucket(), + maxNumBuckets) + } + partition -> numBuckets + } Some( (p: BinaryRow) => knownNumBuckets.getOrDefault( @@ -575,25 +585,72 @@ case class PaimonSparkWriter( .toSeq } - private def countRowsByPartition(df: DataFrame): Map[BinaryRow, Long] = { - val partitionColumns = partitionCols(df) - if (partitionColumns.isEmpty) { - Map(BinaryRow.EMPTY_ROW -> df.count()) - } else { - val partitionType = tableSchema.logicalPartitionType() - val serializer = InternalSerializers.create(partitionType) - df.groupBy(partitionColumns: _*) - .count() - .collect() - .map { - row => - val partitionRow = Row.fromSeq(row.toSeq.dropRight(1)) - val partition = - serializer.toBinaryRow(new SparkRow(partitionType, partitionRow)).copy() - partition -> row.getLong(row.length - 1) - } - .toMap + private def collectDataStatsByPartition( + df: DataFrame, + collectSize: Boolean): Map[BinaryRow, PartitionDataStats] = { + val schema = tableSchema + val rowType = writeType + val toPaimonRow = SparkRowUtils.toPaimonRow( + rowType, + SparkRowUtils.getFieldIndex(df.schema, ROW_KIND_COL), + uriReaderFactoryForBlobDescriptor) + df.rdd + .mapPartitions { + rows => + val partitionKeyExtractor = new RowPartitionKeyExtractor(schema) + val rowSerializer = InternalSerializers.create(rowType) + val stats = mutable.HashMap.empty[SerializedPartition, PartitionDataStats] + rows.foreach { + row => + val paimonRow = toPaimonRow(row) + val partition = SerializedPartition( + SerializationUtils.serializeBinaryRow(partitionKeyExtractor.partition(paimonRow))) + val rowSize = + if (collectSize) rowSerializer.toBinaryRow(paimonRow).getSizeInBytes.toLong else 0L + val previous = stats.getOrElse(partition, PartitionDataStats(0L, 0L)) + stats.put( + partition, + PartitionDataStats( + Math.addExact(previous.rowCount, 1L), + Math.addExact(previous.serializedSize, rowSize))) + } + stats.iterator + } + .reduceByKey { + (left, right) => + PartitionDataStats( + Math.addExact(left.rowCount, right.rowCount), + Math.addExact(left.serializedSize, right.serializedSize)) + } + .collect() + .map { + case (partition, stats) => + SerializationUtils.deserializeBinaryRow(partition.bytes) -> stats + } + .toMap + } + + private def computeBucketNumBySize( + dataStats: PartitionDataStats, + postponeRowCount: Long, + targetSizePerBucket: Long, + maxNumBuckets: Int): Int = { + if (targetSizePerBucket <= 0) { + throw new IllegalArgumentException( + "Option 'postpone.target-size-per-bucket' must be greater than 0.") } + + // Previous postpone files do not record their uncompressed serialized size. Estimate it with + // the average serialized size of incoming rows from the same partition. + val estimatedTotalSizeNumerator = + BigInt(dataStats.serializedSize) * + (BigInt(dataStats.rowCount) + BigInt(postponeRowCount)) + val rowCount = BigInt(dataStats.rowCount) + val estimatedTotalSize = (estimatedTotalSizeNumerator + rowCount - 1) / rowCount + val bucketNum = + if (estimatedTotalSize == 0) BigInt(1) + else (estimatedTotalSize - 1) / BigInt(targetSizePerBucket) + 1 + bucketNum.min(BigInt(maxNumBuckets)).toInt } private def deserializeCommitMessage( @@ -611,10 +668,24 @@ case class PaimonSparkWriter( override def numPartitions: Int = partitions override def getPartition(key: Any): Int = Math.abs(key.asInstanceOf[Int] % numPartitions) } + } object PaimonSparkWriter { + private case class PartitionDataStats(rowCount: Long, serializedSize: Long) + + private case class SerializedPartition(bytes: Array[Byte]) { + override def equals(other: Any): Boolean = { + other match { + case that: SerializedPartition => java.util.Arrays.equals(bytes, that.bytes) + case _ => false + } + } + + override def hashCode(): Int = java.util.Arrays.hashCode(bytes) + } + def apply(table: FileStoreTable): PaimonSparkWriter = { new PaimonSparkWriter(table) } diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala index 943ae01c7f9c..5742ebc8a69e 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala @@ -60,9 +60,11 @@ class PostponeBucketTableTest extends PaimonSparkTestBase { checkAnswer(sql("SELECT sum(k) FROM t"), Seq(Row(499500))) checkAnswer( sql("SELECT distinct(bucket) FROM `t$buckets` ORDER BY bucket"), - Seq(Row(0), Row(1), Row(2), Row(3)) + Seq(Row(0)) ) + sql("ALTER TABLE t SET TBLPROPERTIES ('postpone.target-size-per-bucket' = '8 kb')") + // Write to existing partition, the bucket number should not change sql(""" |INSERT INTO t SELECT /*+ REPARTITION(6) */ @@ -73,7 +75,7 @@ class PostponeBucketTableTest extends PaimonSparkTestBase { |""".stripMargin) checkAnswer( sql("SELECT distinct(bucket) FROM `t$buckets` WHERE partition = '{3}' ORDER BY bucket"), - Seq(Row(0), Row(1), Row(2), Row(3)) + Seq(Row(0)) ) // Write to new partition, the bucket number should change @@ -86,7 +88,7 @@ class PostponeBucketTableTest extends PaimonSparkTestBase { |""".stripMargin) checkAnswer( sql("SELECT distinct(bucket) FROM `t$buckets` WHERE partition = '{5}' ORDER BY bucket"), - Seq(Row(0), Row(1), Row(2), Row(3), Row(4), Row(5)) + Seq(Row(0), Row(1), Row(2)) ) } } @@ -161,6 +163,63 @@ class PostponeBucketTableTest extends PaimonSparkTestBase { } } + test("Postpone bucket table: infer bucket number from serialized data size") { + withTable("t") { + sql(""" + |CREATE TABLE t ( + | k INT, + | v STRING, + | pt INT + |) PARTITIONED BY (pt) + |TBLPROPERTIES ( + | 'primary-key' = 'k, pt', + | 'bucket' = '-2', + | 'postpone.batch-write-fixed-bucket' = 'true', + | 'postpone.target-size-per-bucket' = '32 kb' + |) + |""".stripMargin) + + sql(""" + |INSERT INTO t SELECT /*+ REPARTITION(20) */ + |id AS k, + |CASE WHEN id < 100 THEN repeat('x', 100) ELSE repeat('x', 1000) END AS v, + |CASE WHEN id < 100 THEN 0 ELSE 1 END AS pt + |FROM range (0, 200) + |""".stripMargin) + + checkAnswer( + sql("SELECT distinct(bucket) FROM `t$buckets` WHERE partition = '{0}' ORDER BY bucket"), + Seq(Row(0)) + ) + checkAnswer( + sql("SELECT distinct(bucket) FROM `t$buckets` WHERE partition = '{1}' ORDER BY bucket"), + Seq(Row(0), Row(1), Row(2), Row(3)) + ) + + // Estimate existing postpone data with the average serialized size of incoming rows. + withSparkSQLConf("spark.paimon.postpone.batch-write-fixed-bucket" -> "false") { + sql(""" + |INSERT INTO t SELECT + |id AS k, + |repeat('x', 1000) AS v, + |2 AS pt + |FROM range (1000, 1100) + |""".stripMargin) + } + sql(""" + |INSERT INTO t SELECT /*+ REPARTITION(20) */ + |id AS k, + |repeat('x', 1000) AS v, + |2 AS pt + |FROM range (2000, 2100) + |""".stripMargin) + checkAnswer( + sql("SELECT distinct(bucket) FROM `t$buckets` WHERE partition = '{2}' ORDER BY bucket"), + Seq(Row(-2), Row(0), Row(1), Row(2), Row(3), Row(4), Row(5), Row(6)) + ) + } + } + test("Postpone bucket table: write fix bucket then write postpone bucket") { withTable("t") { sql(""" @@ -186,7 +245,7 @@ class PostponeBucketTableTest extends PaimonSparkTestBase { checkAnswer(sql("SELECT sum(k) FROM t"), Seq(Row(499500))) checkAnswer( sql("SELECT distinct(bucket) FROM `t$buckets` ORDER BY bucket"), - Seq(Row(0), Row(1), Row(2), Row(3)) + Seq(Row(0)) ) // write postpone bucket @@ -201,7 +260,7 @@ class PostponeBucketTableTest extends PaimonSparkTestBase { checkAnswer(sql("SELECT sum(k) FROM t"), Seq(Row(499500))) checkAnswer( sql("SELECT distinct(bucket) FROM `t$buckets` ORDER BY bucket"), - Seq(Row(-2), Row(0), Row(1), Row(2), Row(3)) + Seq(Row(-2), Row(0)) ) } } @@ -675,7 +734,7 @@ class PostponeBucketTableTest extends PaimonSparkTestBase { checkAnswer(sql("SELECT sum(k) FROM t"), Seq(Row(499500))) checkAnswer( sql("SELECT distinct(bucket) FROM `t$buckets` ORDER BY bucket"), - Seq(Row(-2), Row(0), Row(1), Row(2), Row(3), Row(4), Row(5)) + Seq(Row(-2), Row(0)) ) } @@ -691,7 +750,7 @@ class PostponeBucketTableTest extends PaimonSparkTestBase { checkAnswer(sql("SELECT sum(k) FROM t"), Seq(Row(124750))) checkAnswer( sql("SELECT distinct(bucket) FROM `t$buckets` ORDER BY bucket"), - Seq(Row(0), Row(1), Row(2), Row(3), Row(4), Row(5)) + Seq(Row(0)) ) } } From 188df1746f338a359916b8dc2062884cfd86b488 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Mon, 27 Jul 2026 22:13:38 +0800 Subject: [PATCH 3/6] [spark] Skip unnecessary postpone bucket statistics --- .../primary-key-table/data-distribution.md | 3 +- .../spark/commands/PaimonSparkWriter.scala | 83 ++++++++++++------- 2 files changed, 56 insertions(+), 30 deletions(-) diff --git a/docs/docs/primary-key-table/data-distribution.md b/docs/docs/primary-key-table/data-distribution.md index 2bb26974bbb1..a7cf7aa16dc0 100644 --- a/docs/docs/primary-key-table/data-distribution.md +++ b/docs/docs/primary-key-table/data-distribution.md @@ -84,7 +84,8 @@ The calculated bucket number is at least `1` and is limited by The serialized size is measured before file encoding and compression, so it is not the exact ORC or Parquet size on storage. Inferring the data amount adds a statistics stage; Spark caches the batch so that inference and -writing use the same input rows. +writing use the same input rows. Spark skips this stage for an unpartitioned table that already +has real bucket files, or when `postpone.batch-write-fixed-bucket.max-parallelism` is `1`. When `postpone.batch-write-fixed-bucket` is `false`, records are first stored in the `bucket-postpone` directory of each partition diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala index db8589062d75..f0cf8cca7970 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala @@ -113,7 +113,6 @@ case class PaimonSparkWriter( val uriReaderFactory = uriReaderFactoryForBlobDescriptor import sparkSession.implicits._ - val inferPostponeBucketNum = postponeBatchWriteFixedBucket val withInitBucketCol = bucketMode match { case BUCKET_UNAWARE => data case KEY_DYNAMIC if !data.schema.fieldNames.contains(ROW_KIND_COL) => @@ -122,7 +121,26 @@ case class PaimonSparkWriter( .withColumn(BUCKET_COL, lit(-1)) case _ => data.withColumn(BUCKET_COL, lit(-1)) } - if (inferPostponeBucketNum) { + val knownPostponeNumBuckets = + if (postponeBatchWriteFixedBucket) { + Some(PostponeUtils.getKnownNumBuckets(table)) + } else { + None + } + val maxPostponeNumBuckets = + if (postponeBatchWriteFixedBucket) { + coreOptions.postponeBatchWriteFixedBucketMaxParallelism() + } else { + 0 + } + val unpartitionedTableHasKnownNumBuckets = + tableSchema.partitionKeys().isEmpty && + knownPostponeNumBuckets.exists(_.containsKey(BinaryRow.EMPTY_ROW)) + val collectPostponeBucketStats = + postponeBatchWriteFixedBucket && + maxPostponeNumBuckets != 1 && + !unpartitionedTableHasKnownNumBuckets + if (collectPostponeBucketStats) { withInitBucketCol.persist() } @@ -132,33 +150,40 @@ case class PaimonSparkWriter( val postponePartitionBucketComputer: Option[BinaryRow => Integer] = try { if (postponeBatchWriteFixedBucket) { - val knownNumBuckets = PostponeUtils.getKnownNumBuckets(table) - val maxNumBuckets = table.coreOptions().postponeBatchWriteFixedBucketMaxParallelism + val knownNumBuckets = knownPostponeNumBuckets.get + val maxNumBuckets = maxPostponeNumBuckets val defaultPostponeNumBuckets = Math.min(withInitBucketCol.rdd.getNumPartitions, maxNumBuckets) - val targetRowNum = coreOptions.postponeTargetRowNumPerBucket() - val postponeRowCounts = PostponeUtils.getPostponeRowCounts(table) - val dataStats = - collectDataStatsByPartition(withInitBucketCol, collectSize = !targetRowNum.isPresent) - val inferredNumBuckets: Map[BinaryRow, Int] = dataStats.map { - case (partition, stats) => - val postponeRowCount = postponeRowCounts.getOrDefault(partition, 0L) - val numBuckets = - if (targetRowNum.isPresent) { - Math.min( - PostponeUtils.computeBucketNumByRowCount( - Math.addExact(stats.rowCount, postponeRowCount), - targetRowNum.get()), - maxNumBuckets) - } else { - computeBucketNumBySize( - stats, - postponeRowCount, - coreOptions.postponeTargetSizePerBucket(), - maxNumBuckets) - } - partition -> numBuckets - } + val inferredNumBuckets: Map[BinaryRow, Int] = + if (collectPostponeBucketStats) { + val targetRowNum = coreOptions.postponeTargetRowNumPerBucket() + val postponeRowCounts = PostponeUtils.getPostponeRowCounts(table) + val dataStats = + collectDataStatsByPartition( + withInitBucketCol, + collectSize = !targetRowNum.isPresent) + dataStats.map { + case (partition, stats) => + val postponeRowCount = postponeRowCounts.getOrDefault(partition, 0L) + val numBuckets = + if (targetRowNum.isPresent) { + Math.min( + PostponeUtils.computeBucketNumByRowCount( + Math.addExact(stats.rowCount, postponeRowCount), + targetRowNum.get()), + maxNumBuckets) + } else { + computeBucketNumBySize( + stats, + postponeRowCount, + coreOptions.postponeTargetSizePerBucket(), + maxNumBuckets) + } + partition -> numBuckets + } + } else { + Map.empty + } Some( (p: BinaryRow) => knownNumBuckets.getOrDefault( @@ -169,7 +194,7 @@ case class PaimonSparkWriter( } } catch { case e: Throwable => - if (inferPostponeBucketNum) { + if (collectPostponeBucketStats) { withInitBucketCol.unpersist() } throw e @@ -395,7 +420,7 @@ case class PaimonSparkWriter( try { WriteTaskResult.merge(written.collect()) } finally { - if (inferPostponeBucketNum) { + if (collectPostponeBucketStats) { withInitBucketCol.unpersist() } } From 8edeae147a8af84f36073d09437f6de0c55889de Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Mon, 27 Jul 2026 22:20:24 +0800 Subject: [PATCH 4/6] [spark] Extract postpone bucket assignment preparation --- .../spark/commands/PaimonSparkWriter.scala | 136 +++++++++--------- 1 file changed, 65 insertions(+), 71 deletions(-) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala index f0cf8cca7970..4186263d342a 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala @@ -121,84 +121,18 @@ case class PaimonSparkWriter( .withColumn(BUCKET_COL, lit(-1)) case _ => data.withColumn(BUCKET_COL, lit(-1)) } - val knownPostponeNumBuckets = + val postponeBucketAssignment = if (postponeBatchWriteFixedBucket) { - Some(PostponeUtils.getKnownNumBuckets(table)) + Some(preparePostponeBucketAssignment(withInitBucketCol)) } else { None } - val maxPostponeNumBuckets = - if (postponeBatchWriteFixedBucket) { - coreOptions.postponeBatchWriteFixedBucketMaxParallelism() - } else { - 0 - } - val unpartitionedTableHasKnownNumBuckets = - tableSchema.partitionKeys().isEmpty && - knownPostponeNumBuckets.exists(_.containsKey(BinaryRow.EMPTY_ROW)) - val collectPostponeBucketStats = - postponeBatchWriteFixedBucket && - maxPostponeNumBuckets != 1 && - !unpartitionedTableHasKnownNumBuckets - if (collectPostponeBucketStats) { - withInitBucketCol.persist() - } val rowKindColIdx = SparkRowUtils.getFieldIndex(withInitBucketCol.schema, ROW_KIND_COL) val bucketColIdx = SparkRowUtils.getFieldIndex(withInitBucketCol.schema, BUCKET_COL) val encoderGroupWithBucketCol = EncoderSerDeGroup(withInitBucketCol.schema) - val postponePartitionBucketComputer: Option[BinaryRow => Integer] = - try { - if (postponeBatchWriteFixedBucket) { - val knownNumBuckets = knownPostponeNumBuckets.get - val maxNumBuckets = maxPostponeNumBuckets - val defaultPostponeNumBuckets = - Math.min(withInitBucketCol.rdd.getNumPartitions, maxNumBuckets) - val inferredNumBuckets: Map[BinaryRow, Int] = - if (collectPostponeBucketStats) { - val targetRowNum = coreOptions.postponeTargetRowNumPerBucket() - val postponeRowCounts = PostponeUtils.getPostponeRowCounts(table) - val dataStats = - collectDataStatsByPartition( - withInitBucketCol, - collectSize = !targetRowNum.isPresent) - dataStats.map { - case (partition, stats) => - val postponeRowCount = postponeRowCounts.getOrDefault(partition, 0L) - val numBuckets = - if (targetRowNum.isPresent) { - Math.min( - PostponeUtils.computeBucketNumByRowCount( - Math.addExact(stats.rowCount, postponeRowCount), - targetRowNum.get()), - maxNumBuckets) - } else { - computeBucketNumBySize( - stats, - postponeRowCount, - coreOptions.postponeTargetSizePerBucket(), - maxNumBuckets) - } - partition -> numBuckets - } - } else { - Map.empty - } - Some( - (p: BinaryRow) => - knownNumBuckets.getOrDefault( - p, - Integer.valueOf(inferredNumBuckets.getOrElse(p, defaultPostponeNumBuckets)))) - } else { - None - } - } catch { - case e: Throwable => - if (collectPostponeBucketStats) { - withInitBucketCol.unpersist() - } - throw e - } + val postponePartitionBucketComputer = + postponeBucketAssignment.map(_.partitionBucketComputer) def newWrite() = PaimonDataWrite( writeBuilder, @@ -420,7 +354,7 @@ case class PaimonSparkWriter( try { WriteTaskResult.merge(written.collect()) } finally { - if (collectPostponeBucketStats) { + if (postponeBucketAssignment.exists(_.dataPersisted)) { withInitBucketCol.unpersist() } } @@ -610,6 +544,62 @@ case class PaimonSparkWriter( .toSeq } + private def preparePostponeBucketAssignment(df: DataFrame): PostponeBucketAssignment = { + val knownNumBuckets = PostponeUtils.getKnownNumBuckets(table) + val maxNumBuckets = coreOptions.postponeBatchWriteFixedBucketMaxParallelism() + val unpartitionedTableHasKnownNumBuckets = + tableSchema.partitionKeys().isEmpty && + knownNumBuckets.containsKey(BinaryRow.EMPTY_ROW) + val collectBucketStats = + maxNumBuckets != 1 && !unpartitionedTableHasKnownNumBuckets + if (collectBucketStats) { + df.persist() + } + + try { + val defaultNumBuckets = Math.min(df.rdd.getNumPartitions, maxNumBuckets) + val inferredNumBuckets: Map[BinaryRow, Int] = + if (collectBucketStats) { + val targetRowNum = coreOptions.postponeTargetRowNumPerBucket() + val postponeRowCounts = PostponeUtils.getPostponeRowCounts(table) + val dataStats = + collectDataStatsByPartition(df, collectSize = !targetRowNum.isPresent) + dataStats.map { + case (partition, stats) => + val postponeRowCount = postponeRowCounts.getOrDefault(partition, 0L) + val numBuckets = + if (targetRowNum.isPresent) { + Math.min( + PostponeUtils.computeBucketNumByRowCount( + Math.addExact(stats.rowCount, postponeRowCount), + targetRowNum.get()), + maxNumBuckets) + } else { + computeBucketNumBySize( + stats, + postponeRowCount, + coreOptions.postponeTargetSizePerBucket(), + maxNumBuckets) + } + partition -> numBuckets + } + } else { + Map.empty + } + val partitionBucketComputer = (partition: BinaryRow) => + knownNumBuckets.getOrDefault( + partition, + Integer.valueOf(inferredNumBuckets.getOrElse(partition, defaultNumBuckets))) + PostponeBucketAssignment(partitionBucketComputer, collectBucketStats) + } catch { + case e: Throwable => + if (collectBucketStats) { + df.unpersist() + } + throw e + } + } + private def collectDataStatsByPartition( df: DataFrame, collectSize: Boolean): Map[BinaryRow, PartitionDataStats] = { @@ -698,6 +688,10 @@ case class PaimonSparkWriter( object PaimonSparkWriter { + private case class PostponeBucketAssignment( + partitionBucketComputer: BinaryRow => Integer, + dataPersisted: Boolean) + private case class PartitionDataStats(rowCount: Long, serializedSize: Long) private case class SerializedPartition(bytes: Array[Byte]) { From ececb21abde2611a174f5354170b56b58d111709 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Mon, 27 Jul 2026 22:26:31 +0800 Subject: [PATCH 5/6] [spark] Clarify postpone bucket inference flag --- .../paimon/spark/commands/PaimonSparkWriter.scala | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala index 4186263d342a..99e17d06239f 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala @@ -550,16 +550,16 @@ case class PaimonSparkWriter( val unpartitionedTableHasKnownNumBuckets = tableSchema.partitionKeys().isEmpty && knownNumBuckets.containsKey(BinaryRow.EMPTY_ROW) - val collectBucketStats = + val inferBucketNumFromData = maxNumBuckets != 1 && !unpartitionedTableHasKnownNumBuckets - if (collectBucketStats) { + if (inferBucketNumFromData) { df.persist() } try { val defaultNumBuckets = Math.min(df.rdd.getNumPartitions, maxNumBuckets) val inferredNumBuckets: Map[BinaryRow, Int] = - if (collectBucketStats) { + if (inferBucketNumFromData) { val targetRowNum = coreOptions.postponeTargetRowNumPerBucket() val postponeRowCounts = PostponeUtils.getPostponeRowCounts(table) val dataStats = @@ -590,10 +590,10 @@ case class PaimonSparkWriter( knownNumBuckets.getOrDefault( partition, Integer.valueOf(inferredNumBuckets.getOrElse(partition, defaultNumBuckets))) - PostponeBucketAssignment(partitionBucketComputer, collectBucketStats) + PostponeBucketAssignment(partitionBucketComputer, inferBucketNumFromData) } catch { case e: Throwable => - if (collectBucketStats) { + if (inferBucketNumFromData) { df.unpersist() } throw e From c46623424bd4c61442c5a7c90e89f4b7a3234e82 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Tue, 28 Jul 2026 00:03:14 +0800 Subject: [PATCH 6/6] [spark] Fix postpone bucket inference edge cases --- .../spark/commands/PaimonSparkWriter.scala | 40 ++++++++-- .../spark/commands/WriteIntoPaimonTable.scala | 5 +- .../spark/sql/PostponeBucketTableTest.scala | 76 +++++++++++++++++++ 3 files changed, 113 insertions(+), 8 deletions(-) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala index 99e17d06239f..aaba0f1ca0bf 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala @@ -109,6 +109,12 @@ case class PaimonSparkWriter( } def write(data: DataFrame): Seq[CommitMessage] = { + write(data, overwriteExistingData = false) + } + + private[commands] def write( + data: DataFrame, + overwriteExistingData: Boolean): Seq[CommitMessage] = { val sparkSession = data.sparkSession val uriReaderFactory = uriReaderFactoryForBlobDescriptor import sparkSession.implicits._ @@ -123,7 +129,7 @@ case class PaimonSparkWriter( } val postponeBucketAssignment = if (postponeBatchWriteFixedBucket) { - Some(preparePostponeBucketAssignment(withInitBucketCol)) + Some(preparePostponeBucketAssignment(withInitBucketCol, overwriteExistingData)) } else { None } @@ -544,7 +550,9 @@ case class PaimonSparkWriter( .toSeq } - private def preparePostponeBucketAssignment(df: DataFrame): PostponeBucketAssignment = { + private def preparePostponeBucketAssignment( + df: DataFrame, + overwriteExistingData: Boolean): PostponeBucketAssignment = { val knownNumBuckets = PostponeUtils.getKnownNumBuckets(table) val maxNumBuckets = coreOptions.postponeBatchWriteFixedBucketMaxParallelism() val unpartitionedTableHasKnownNumBuckets = @@ -561,7 +569,12 @@ case class PaimonSparkWriter( val inferredNumBuckets: Map[BinaryRow, Int] = if (inferBucketNumFromData) { val targetRowNum = coreOptions.postponeTargetRowNumPerBucket() - val postponeRowCounts = PostponeUtils.getPostponeRowCounts(table) + val postponeRowCounts = + if (overwriteExistingData) { + java.util.Collections.emptyMap[BinaryRow, java.lang.Long]() + } else { + PostponeUtils.getPostponeRowCounts(table) + } val dataStats = collectDataStatsByPartition(df, collectSize = !targetRowNum.isPresent) dataStats.map { @@ -569,10 +582,9 @@ case class PaimonSparkWriter( val postponeRowCount = postponeRowCounts.getOrDefault(partition, 0L) val numBuckets = if (targetRowNum.isPresent) { - Math.min( - PostponeUtils.computeBucketNumByRowCount( - Math.addExact(stats.rowCount, postponeRowCount), - targetRowNum.get()), + computeBucketNumByRowCount( + Math.addExact(stats.rowCount, postponeRowCount), + targetRowNum.get(), maxNumBuckets) } else { computeBucketNumBySize( @@ -688,6 +700,20 @@ case class PaimonSparkWriter( object PaimonSparkWriter { + private[spark] def computeBucketNumByRowCount( + rowCount: Long, + targetRowNumPerBucket: Long, + maxNumBuckets: Int): Int = { + if (targetRowNumPerBucket <= 0) { + throw new IllegalArgumentException( + "Option 'postpone.target-row-num-per-bucket' must be greater than 0.") + } + + val bucketNum = + if (rowCount <= 0) 1L else (rowCount - 1) / targetRowNumPerBucket + 1 + Math.min(bucketNum, maxNumBuckets.toLong).toInt + } + private case class PostponeBucketAssignment( partitionBucketComputer: BinaryRow => Integer, dataPersisted: Boolean) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/WriteIntoPaimonTable.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/WriteIntoPaimonTable.scala index f400409fdae1..346aa0fe3ae8 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/WriteIntoPaimonTable.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/WriteIntoPaimonTable.scala @@ -71,7 +71,10 @@ case class WriteIntoPaimonTable( Snapshot.Operation.WRITE } } - val commitMessages = writer.write(replacedData) + val commitMessages = + writer.write( + replacedData, + overwriteExistingData = overwritePartition != null || dynamicPartitionOverwriteMode) writer.commit(commitMessages, operation) Seq.empty diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala index 5742ebc8a69e..1e9910bf7e6e 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PostponeBucketTableTest.scala @@ -22,6 +22,7 @@ import org.apache.paimon.catalog.{Catalog, CatalogLoader, DelegateCatalog, Ident import org.apache.paimon.deletionvectors.DeletionVectorsIndexFile.DELETION_VECTORS_INDEX import org.apache.paimon.fs.Path import org.apache.paimon.spark.{PaimonScan, PaimonSparkTestBase} +import org.apache.paimon.spark.commands.PaimonSparkWriter import org.apache.paimon.spark.procedure.SparkPostponeCompactProcedure import org.apache.paimon.table.{CatalogEnvironment, FileStoreTableFactory} @@ -33,6 +34,14 @@ import scala.collection.JavaConverters._ class PostponeBucketTableTest extends PaimonSparkTestBase { + test("Postpone bucket table: cap inferred row-count buckets before Int conversion") { + assert( + PaimonSparkWriter.computeBucketNumByRowCount( + Integer.MAX_VALUE.toLong + 1L, + targetRowNumPerBucket = 1L, + maxNumBuckets = 2048) == 2048) + } + test("Postpone bucket table: write with different bucket number") { withTable("t") { sql(""" @@ -163,6 +172,73 @@ class PostponeBucketTableTest extends PaimonSparkTestBase { } } + test("Postpone bucket table: overwrite excludes existing postpone rows from inference") { + Seq( + ( + "static", + """ + |INSERT OVERWRITE t SELECT + |id AS k, + |CAST(id AS STRING) AS v, + |0 AS pt + |FROM range (1000, 1100) + |""".stripMargin), + ( + "static", + """ + |INSERT OVERWRITE t PARTITION (pt = 0) SELECT + |id AS k, + |CAST(id AS STRING) AS v + |FROM range (1000, 1100) + |""".stripMargin), + ( + "dynamic", + """ + |INSERT OVERWRITE t SELECT + |id AS k, + |CAST(id AS STRING) AS v, + |0 AS pt + |FROM range (1000, 1100) + |""".stripMargin) + ).foreach { + case (partitionOverwriteMode, overwriteSql) => + withTable("t") { + sql(""" + |CREATE TABLE t ( + | k INT, + | v STRING, + | pt INT + |) PARTITIONED BY (pt) + |TBLPROPERTIES ( + | 'primary-key' = 'k, pt', + | 'bucket' = '-2', + | 'postpone.batch-write-fixed-bucket' = 'false', + | 'postpone.target-row-num-per-bucket' = '100' + |) + |""".stripMargin) + + sql(""" + |INSERT INTO t SELECT + |id AS k, + |CAST(id AS STRING) AS v, + |0 AS pt + |FROM range (0, 1000) + |""".stripMargin) + + withSparkSQLConf( + "spark.paimon.postpone.batch-write-fixed-bucket" -> "true", + "spark.sql.sources.partitionOverwriteMode" -> partitionOverwriteMode) { + sql(overwriteSql) + } + + checkAnswer( + sql("SELECT distinct(bucket) FROM `t$buckets` WHERE partition = '{0}'"), + Seq(Row(0)) + ) + } + } + } + test("Postpone bucket table: infer bucket number from serialized data size") { withTable("t") { sql("""