[spark] Infer postpone bucket count from data size

docs/docs/primary-key-table/data-distribution.md

@@ -70,12 +70,27 @@ 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,
+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. 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
 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`.

docs/generated/core_configuration.html

@@ -1294,7 +1294,13 @@
             <td><h5>postpone.target-row-num-per-bucket</h5></td>
             <td style="word-wrap: break-word;">(none)</td>
             <td>Long</td>
-            <td>Target row number per bucket for partitions compacted from postpone bucket files for the first time.</td>
+            <td>Target row number per bucket when batch writing fixed buckets or compacting postpone bucket files for a partition without real bucket data.</td>
+        </tr>
+        <tr>
+            <td><h5>postpone.target-size-per-bucket</h5></td>
+            <td style="word-wrap: break-word;">1 gb</td>
+            <td>MemorySize</td>
+            <td>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.</td>
         </tr>
         <tr>
             <td><h5>primary-key</h5></td>

paimon-api/src/main/java/org/apache/paimon/CoreOptions.java

@@ -2733,7 +2733,15 @@ 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<MemorySize> 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<Long> GLOBAL_INDEX_ROW_COUNT_PER_SHARD =
             key("global-index.row-count-per-shard")
@@ -4406,6 +4414,10 @@ public Optional<Long> 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);
     }

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala

@@ -52,13 +52,16 @@ import java.io.IOException
 import java.util.Collections.singletonMap
 
 import scala.collection.JavaConverters._
+import scala.collection.mutable
 
 case class PaimonSparkWriter(
     table: FileStoreTable,
     writeRowTracking: Boolean = false,
     batchId: Option[Long] = None)
   extends WriteHelper {
 
+  import PaimonSparkWriter._
+
   private lazy val tableSchema = table.schema
 
   private lazy val bucketMode = table.bucketMode
@@ -106,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._
@@ -118,20 +127,19 @@ case class PaimonSparkWriter(
           .withColumn(BUCKET_COL, lit(-1))
       case _ => data.withColumn(BUCKET_COL, lit(-1))
     }
-    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] =
+    val postponeBucketAssignment =
       if (postponeBatchWriteFixedBucket) {
-        val knownNumBuckets = PostponeUtils.getKnownNumBuckets(table)
-        val defaultPostponeNumBuckets = Math.min(
-          withInitBucketCol.rdd.getNumPartitions,
-          table.coreOptions().postponeBatchWriteFixedBucketMaxParallelism)
-        Some((p: BinaryRow) => knownNumBuckets.getOrDefault(p, defaultPostponeNumBuckets))
+        Some(preparePostponeBucketAssignment(withInitBucketCol, overwriteExistingData))
       } else {
         None
       }
 
+    val rowKindColIdx = SparkRowUtils.getFieldIndex(withInitBucketCol.schema, ROW_KIND_COL)
+    val bucketColIdx = SparkRowUtils.getFieldIndex(withInitBucketCol.schema, BUCKET_COL)
+    val encoderGroupWithBucketCol = EncoderSerDeGroup(withInitBucketCol.schema)
+    val postponePartitionBucketComputer =
+      postponeBucketAssignment.map(_.partitionBucketComputer)
+
     def newWrite() = PaimonDataWrite(
       writeBuilder,
       writeType,
@@ -349,7 +357,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 (postponeBucketAssignment.exists(_.dataPersisted)) {
+        withInitBucketCol.unpersist()
+      }
+    }
   }
 
   /**
@@ -536,6 +550,136 @@ case class PaimonSparkWriter(
       .toSeq
   }
 
+  private def preparePostponeBucketAssignment(
+      df: DataFrame,
+      overwriteExistingData: Boolean): PostponeBucketAssignment = {
+    val knownNumBuckets = PostponeUtils.getKnownNumBuckets(table)
+    val maxNumBuckets = coreOptions.postponeBatchWriteFixedBucketMaxParallelism()
+    val unpartitionedTableHasKnownNumBuckets =
+      tableSchema.partitionKeys().isEmpty &&
+        knownNumBuckets.containsKey(BinaryRow.EMPTY_ROW)
+    val inferBucketNumFromData =
+      maxNumBuckets != 1 && !unpartitionedTableHasKnownNumBuckets
+    if (inferBucketNumFromData) {
+      df.persist()
+    }
+
+    try {
+      val defaultNumBuckets = Math.min(df.rdd.getNumPartitions, maxNumBuckets)
+      val inferredNumBuckets: Map[BinaryRow, Int] =
+        if (inferBucketNumFromData) {
+          val targetRowNum = coreOptions.postponeTargetRowNumPerBucket()
+          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 {
+            case (partition, stats) =>
+              val postponeRowCount = postponeRowCounts.getOrDefault(partition, 0L)
+              val numBuckets =
+                if (targetRowNum.isPresent) {
+                  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, inferBucketNumFromData)
+    } catch {
+      case e: Throwable =>
+        if (inferBucketNumFromData) {
+          df.unpersist()
+        }
+        throw e
+    }
+  }
+
+  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(
       serializer: CommitMessageSerializer,
       bytes: Array[Byte]): CommitMessage = {
@@ -551,10 +695,42 @@ case class PaimonSparkWriter(
     override def numPartitions: Int = partitions
     override def getPartition(key: Any): Int = Math.abs(key.asInstanceOf[Int] % numPartitions)
   }
+
 }
 
 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)
+
+  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)
   }

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