fix: ship Iceberg scan pool entries per partition so task memory does not grow with the table

contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala

@@ -210,7 +210,7 @@ case class CometDeltaNativeScanExec(
       sparkContext,
       Seq.empty,
       Map(sourceKey -> commonData),
-      Map(sourceKey -> perPartitionData),
+      Map(sourceKey -> PartitionPlanData.wrap(perPartitionData)),
       serializedPlan,
       PlanDataInjector.planFingerprint(serializedPlan),
       perPartitionData.length,

docs/source/user-guide/latest/iceberg.md

@@ -194,6 +194,15 @@ For a custom S3-compatible endpoint, configure the catalog with the endpoint, pa
 
 These `s3.*` storage properties are not specific to the Hive catalog shown here. When `s3.access-key-id` / `s3.secret-access-key` are omitted, credentials come from the standard AWS chain (environment variables, instance profiles, and so on). The region is not auto-detected: when neither the catalog (`client.region` or `s3.region`) nor the executor environment (`AWS_REGION` or `AWS_DEFAULT_REGION`) supplies one, Comet uses `us-east-1`, so set it for AWS buckets in any other region. If your REST catalog vends temporary credentials, the native reader does not consume them automatically, and wiring that requires the credential provider bridge. See Iceberg's [S3 FileIO](https://iceberg.apache.org/docs/latest/aws/#s3-fileio) docs for the full property list, and [S3 Credential Providers](s3-credential-providers.md) for vended or per-request credentials.
 
+### Plan data per task
+
+The native reader sends each task only the delete files, partition values and residual filters
+that its own files need, the way Spark's Iceberg reader does. A delete file that many partitions
+use, such as an unpartitioned equality delete, is serialized once and shared by those
+partitions. The driver keeps one copy of it, and a coalesce without a shuffle (`coalesce(n)`, or
+a write that coalesces its input), which combines several input partitions into one task, sends
+that task one copy.
+
 ### Current limitations
 
 The following scenarios will fall back to the JVM Iceberg reader:

native/proto/src/proto/operator.proto

@@ -388,7 +388,12 @@ message IcebergScanCommon {
   // Schema to read
   repeated SparkStructField required_schema = 3;
 
-  // Deduplication pools (must contain all entries for cross-partition deduplication)
+  // Deduplication pools. In the common block every task gets, the per-file pools
+  // (partition_data_pool, delete_files_pool, residual_pool, delete_file_pool,
+  // delete_file_path_pool) are empty: each partition's IcebergScan.common holds the entries its
+  // tasks use, and the JVM merges the two before the plan reaches native. The driver ships a
+  // partition's delete files and their paths apart from its message, shared with the other
+  // partitions that use them, and the executor JVM adds them to the partition's pools.
   repeated string schema_pool = 4;
   // Index-aligned with partition_spec_pool: entry i is the StructType JSON for the spec at
   // partition_spec_pool[i]. Native uses this pairing to recover each entry's spec_id when
@@ -414,8 +419,8 @@ message IcebergScanCommon {
 
   // Flat pool of unique delete files, referenced by index from DeleteFileList. Iceberg's default
   // partition delete granularity has one delete file apply to many data files, so interning each
-  // delete file here (rather than embedding copies per delete_files_pool set) keeps the broadcast
-  // common message from growing with the number of references.
+  // delete file here (rather than embedding copies per delete_files_pool set) keeps a
+  // partition's pools from growing with the number of references.
   repeated IcebergDeleteFile delete_file_pool = 14;
 
   // Flat pool of unique delete-file paths, referenced by index from IcebergDeleteFile. A V3
@@ -432,7 +437,8 @@ message IcebergScanCommon {
 }
 
 message IcebergScan {
-  // Common data shared across partitions (pools, metadata, catalog props)
+  // Common data shared across partitions (pools, metadata, catalog props). In the per-partition
+  // message the driver ships, it holds only that partition's per-file pool entries.
   IcebergScanCommon common = 1;
 
   // Single partition's file scan tasks

spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala

@@ -37,7 +37,7 @@ import org.apache.comet.{CometExecIterator, CometRuntimeException, CometShuffleB
 private[spark] class CometExecPartition(
     override val index: Int,
     val inputPartitions: Array[Partition],
-    val planDataByKey: Map[String, Array[Byte]],
+    val planDataByKey: Map[String, PartitionPlanData],
     val filePaths: Seq[String] = Seq.empty)
     extends Partition
 
@@ -58,7 +58,7 @@ private[spark] class CometExecRDD(
     sc: SparkContext,
     var inputRDDs: Seq[RDD[_]],
     commonByKey: Map[String, Array[Byte]],
-    @transient perPartitionByKey: Map[String, Array[Array[Byte]]],
+    @transient perPartitionByKey: Map[String, Array[PartitionPlanData]],
     serializedPlan: Array[Byte],
     planFingerprint: Long,
     defaultNumPartitions: Int,
@@ -134,8 +134,8 @@ private[spark] class CometExecRDD(
     // fingerprints the bytes once so the probe here does not rehash them.
     val actualPlan = if (commonByKey.nonEmpty) {
       val basePlan = PlanDataInjector.parseBasePlan(serializedPlan, planFingerprint)
-      val injected =
-        PlanDataInjector.injectPlanData(basePlan, commonByKey, partition.planDataByKey)
+      val partitionBytes = partition.planDataByKey.map { case (key, data) => key -> data.bytes }
+      val injected = PlanDataInjector.injectPlanData(basePlan, commonByKey, partitionBytes)
       PlanDataInjector.serializeOperator(injected)
     } else {
       serializedPlan
@@ -236,7 +236,7 @@ object CometExecRDD {
       sc: SparkContext,
       inputRDDs: Seq[RDD[_]],
       commonByKey: Map[String, Array[Byte]],
-      perPartitionByKey: Map[String, Array[Array[Byte]]],
+      perPartitionByKey: Map[String, Array[PartitionPlanData]],
       serializedPlan: Array[Byte],
       planFingerprint: Long,
       numPartitions: Int,

spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala

@@ -92,7 +92,7 @@ case class CometIcebergNativeScanExec(
    * by every construction site), so values resolved through `waitForSubqueries` are visible on
    * both sides.
    */
-  @transient private lazy val serializedPartitionData: (Array[Byte], Array[Array[Byte]]) = {
+  @transient private lazy val serializedPartitionData: (Array[Byte], Array[PartitionPlanData]) = {
     // Canonicalized instances set originalPlan = null and are not meant to be executed.
     // If we ever reach this lazy val on a canonicalized form, fail loud rather than NPE
     // deep inside originalPlan.inputRDD.
@@ -122,7 +122,7 @@ case class CometIcebergNativeScanExec(
 
   def commonData: Array[Byte] = serializedPartitionData._1
 
-  def perPartitionData: Array[Array[Byte]] = serializedPartitionData._2
+  def perPartitionData: Array[PartitionPlanData] = serializedPartitionData._2
 
   // numPartitions for execution - derived from actual DPP-filtered partitions
   // Only accessed during execution, not planning

spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala

@@ -264,7 +264,7 @@ case class CometNativeScanExec(
       sparkContext,
       Seq.empty,
       Map(sourceKey -> commonData),
-      Map(sourceKey -> perPartitionData),
+      Map(sourceKey -> PartitionPlanData.wrap(perPartitionData)),
       serializedPlan,
       PlanDataInjector.planFingerprint(serializedPlan),
       perPartitionData.length,

spark/src/main/scala/org/apache/spark/sql/comet/PartitionPlanData.scala

@@ -0,0 +1,104 @@
+/*
+ * 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.spark.sql.comet
+
+import scala.collection.mutable
+
+import org.apache.comet.serde.OperatorOuterClass
+
+/**
+ * One partition's planning data for one scan. It rides in the task with its partition object.
+ *
+ * A coalesce packs several partitions into one task, and Java serialization writes an object they
+ * share once per task stream. So a scan can hold an entry several partitions use as one shared
+ * object, and a coalesced task carries it once, as Spark's Iceberg reader does with its shared
+ * DeleteFile objects.
+ */
+trait PartitionPlanData extends Serializable {
+
+  /** The serialized scan message this scan's [[PlanDataInjector]] reads. */
+  def bytes: Array[Byte]
+}
+
+object PartitionPlanData {
+
+  /** Plan data that is already one serialized message, with nothing shared. */
+  def apply(bytes: Array[Byte]): PartitionPlanData = new Bytes(bytes)
+
+  def wrap(perPartition: Array[Array[Byte]]): Array[PartitionPlanData] =
+    perPartition.map(apply)
+
+  /** Whether two scans' plan data match, for [[PlanDataInjector.findAllPlanData]]. */
+  def same(a: Array[PartitionPlanData], b: Array[PartitionPlanData]): Boolean =
+    a.length == b.length && a.indices.forall { i =>
+      (a(i) eq b(i)) || java.util.Arrays.equals(a(i).bytes, b(i).bytes)
+    }
+
+  private final class Bytes(override val bytes: Array[Byte]) extends PartitionPlanData
+}
+
+/**
+ * One native Iceberg scan partition's plan data.
+ *
+ * `scan` is the partition's IcebergScan: its file tasks, plus its partition values, residuals and
+ * delete-file sets in `common`. Its delete files are kept apart, because a delete file is often
+ * used by many partitions: the driver serializes each delete file once, and every partition that
+ * uses it holds the same array and the same path string. A coalesced task then carries each
+ * delete file once, and the driver keeps one copy of it.
+ *
+ * `deleteFiles(i)` is the partition's delete file `i`, an IcebergDeleteFile without
+ * `file_path_idx`, and `deleteFilePaths(i)` is its path. The delete-file sets in `scan` index
+ * into `deleteFiles`. Each partition orders its delete files differently, so the index into the
+ * path pool is set here, when the partition's message is built.
+ */
+final class IcebergPartitionPlanData(
+    val scan: Array[Byte],
+    val deleteFiles: Array[Array[Byte]],
+    val deleteFilePaths: Array[String])
+    extends PartitionPlanData {
+
+  require(
+    deleteFiles.length == deleteFilePaths.length,
+    s"${deleteFiles.length} delete files but ${deleteFilePaths.length} paths")
+
+  override def bytes: Array[Byte] = {
+    if (deleteFiles.isEmpty) {
+      scan
+    } else {
+      val builder = OperatorOuterClass.IcebergScan.newBuilder().mergeFrom(scan)
+      val common = builder.getCommonBuilder
+      val pathIndex = mutable.HashMap[String, Int]()
+      deleteFiles.indices.foreach { i =>
+        val path = deleteFilePaths(i)
+        val pathIdx = pathIndex.getOrElseUpdate(
+          path, {
+            common.addDeleteFilePathPool(path)
+            pathIndex.size
+          })
+        common.addDeleteFilePool(
+          OperatorOuterClass.IcebergDeleteFile
+            .newBuilder()
+            .mergeFrom(deleteFiles(i))
+            .setFilePathIdx(pathIdx))
+      }
+      builder.build().toByteArray
+    }
+  }
+}

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala

@@ -21,7 +21,7 @@ package org.apache.spark.sql.comet.execution.shuffle
 
 import org.apache.spark._
 import org.apache.spark.rdd.{DeterministicLevel, RDD}
-import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode}
+import org.apache.spark.sql.comet.{CometExecRDD, CometMetricNode, PartitionPlanData}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
 import org.apache.comet.CometShuffleBlockIterator
@@ -44,7 +44,7 @@ private[shuffle] class CometNativeShuffleInputRDD(
     numPartitionsParam: Int,
     shuffleScanIndices: Set[Int],
     spillMetricNode: CometMetricNode,
-    @transient perPartitionByKey: Map[String, Array[Array[Byte]]] = Map.empty,
+    @transient perPartitionByKey: Map[String, Array[PartitionPlanData]] = Map.empty,
     positionalRoundRobin: Boolean = false)
     extends RDD[Product2[Int, ColumnarBatch]](
       sc,
@@ -144,7 +144,7 @@ private[shuffle] class CometNativeShuffleInputRDD(
 private[shuffle] class CometNativeShuffleInputPartition(
     override val index: Int,
     val inputPartitions: Array[Partition],
-    val planDataByKey: Map[String, Array[Byte]])
+    val planDataByKey: Map[String, PartitionPlanData])
     extends Partition
 
 /**
@@ -159,7 +159,7 @@ private[shuffle] class CometNativeShuffleInputIterator(
     val partitionIndex: Int,
     val inputObjects: Array[Object],
     val shuffleBlockIterators: Map[Int, CometShuffleBlockIterator],
-    val planDataByKey: Map[String, Array[Byte]])
+    val planDataByKey: Map[String, PartitionPlanData])
     extends Iterator[Product2[Int, ColumnarBatch]] {
 
   override def hasNext: Boolean = false

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala

@@ -156,7 +156,7 @@ class CometNativeShuffleWriter[K, V](
         shuffleId,
         unifiedPlan,
         ctx.commonByKey,
-        shuffleInputIter.planDataByKey)
+        shuffleInputIter.planDataByKey.map { case (key, data) => key -> data.bytes })
     } else {
       unifiedPlan
     }

spark/src/main/scala/org/apache/spark/sql/comet/operators.scala

@@ -454,7 +454,7 @@ private[comet] object PlanDataInjector extends Logging {
    *   (commonByKey, perPartitionByKey) - common data is shared, per-partition varies
    */
   private[comet] def findAllPlanData(
-      plan: SparkPlan): (Map[String, Array[Byte]], Map[String, Array[Array[Byte]]]) = {
+      plan: SparkPlan): (Map[String, Array[Byte]], Map[String, Array[PartitionPlanData]]) = {
     plan match {
       case iceberg: CometIcebergNativeScanExec =>
         // Trigger Spark's standard prepare -> waitForSubqueries lifecycle so DPP
@@ -496,7 +496,9 @@ private[comet] object PlanDataInjector extends Logging {
         s.ensureSubqueriesResolved()
         s.sendDriverMetrics()
         if (s.commonData.nonEmpty && s.perPartitionData.nonEmpty) {
-          (Map(s.sourceKey -> s.commonData), Map(s.sourceKey -> s.perPartitionData))
+          (
+            Map(s.sourceKey -> s.commonData),
+            Map(s.sourceKey -> PartitionPlanData.wrap(s.perPartitionData)))
         } else {
           (Map.empty, Map.empty)
         }
@@ -518,13 +520,10 @@ private[comet] object PlanDataInjector extends Logging {
         val results = plan.children.map(findAllPlanData)
         (
           mergeByKey(results.map(_._1))(java.util.Arrays.equals(_, _)),
-          mergeByKey(results.map(_._2))(samePartitionData))
+          mergeByKey(results.map(_._2))(PartitionPlanData.same))
     }
   }
 
-  private def samePartitionData(a: Array[Array[Byte]], b: Array[Array[Byte]]): Boolean =
-    a.length == b.length && a.indices.forall(i => java.util.Arrays.equals(a(i), b(i)))
-
   /**
    * Merges the planning data found under sibling subtrees. The injector looks data up by key, so
    * two scans sharing a key would both read whichever entry survived. The same data reached twice
@@ -573,11 +572,18 @@ private[comet] object IcebergPlanDataInjector extends PlanDataInjector {
   override def inject(op: Operator, common: Prepared, partitionBytes: Array[Byte]): Operator = {
     val scan = op.getIcebergScan
 
-    val tasksOnly = OperatorOuterClass.IcebergScan.parseFrom(partitionBytes)
+    val partition = OperatorOuterClass.IcebergScan.parseFrom(partitionBytes)
 
     val scanBuilder = scan.toBuilder
-    scanBuilder.setCommon(common)
-    scanBuilder.addAllFileScanTasks(tasksOnly.getFileScanTasksList)
+    // The partition carries the per-file pool entries its tasks use, which the common block does
+    // not hold. mergeFrom appends repeated fields, so the tasks' indices, which start at 0, stay
+    // valid.
+    if (partition.hasCommon) {
+      scanBuilder.setCommon(common.toBuilder.mergeFrom(partition.getCommon))
+    } else {
+      scanBuilder.setCommon(common)
+    }
+    scanBuilder.addAllFileScanTasks(partition.getFileScanTasksList)
 
     op.toBuilder.setIcebergScan(scanBuilder).build()
   }
@@ -887,7 +893,7 @@ private[comet] case class NativeExecContext(
     // Partition object - see CometNativeShuffleInputRDD / CometExecRDD); the executor reads its own
     // slice, never this map. Keeping it off the wire stops it from bloating the broadcast task
     // binary when this context rides on the non-transient CometShuffleDependency.nativeShuffleSpec.
-    @transient perPartitionByKey: Map[String, Array[Array[Byte]]],
+    @transient perPartitionByKey: Map[String, Array[PartitionPlanData]],
     shuffleScanIndices: Set[Int],
     hasScanInput: Boolean,
     // Like plan data, file paths are sliced onto each task's Partition on the driver.