fix: ship Iceberg scan pool entries per partition so task memory does not grow with the table by parthchandra · Pull Request #7491 · apache/datafusion-comet

parthchandra · 2026-11-04T21:44:57Z

Which issue does this PR close?

Closes #7490.

Rationale for this change

Each native Iceberg scan task deserialized and serialized the whole scan's delete-file metadata, even when its own files needed a tiny fraction of it. The scan's pools of delete files, delete-file paths, delete-file sets, partition values and residuals lived in the common block. Every task got the common block in the task binary, and IcebergPlanDataInjector.inject copied it into each partition's native plan. On large merge-on-read tables, executor heap use grew with the size of the table, and tasks ran out of memory where Spark does not.

What changes are included in this PR?

Per-partition pools

CometIcebergNativeScan.serializePartitions now interns those five pools per partition, in the pass that already interned them. Each partition's entries go into its own IcebergScan.common, with indices starting at 0. Each task gets only the entries its own files use, as Spark's Iceberg reader ships each task's delete files with the task.
The common block keeps only what doesn't grow with the files: schemas, partition specs and types, name mappings and projected field ids.

Executor merge

IcebergPlanDataInjector.inject merges a partition's pools into the common block with mergeFrom, which appends repeated fields. The plan native receives has the same shape as before, so native code doesn't change.

Logging

The driver log now reports per-partition pool entry totals and per-partition message sizes.

Behavior compared with the latest release

Query results, fallback decisions and configs are unchanged.
Per-task executor memory for Iceberg scan planning data no longer grows with the table.
An entry that many partitions share, such as an unpartitioned equality delete, is now serialized into each of those partitions' messages, not once in the common block. The common block reached every task anyway, so no task's heap grows. Each task's serialized size does grow by the shared entries it actually uses, which is the same thing Spark's Iceberg reader pays.

Known limits (documented in the Iceberg user guide)

A coalesce without a shuffle combines several input partitions into one task. That task carries a copy of a shared delete file for every input partition that uses it.
The driver still builds every partition's message up front. Building them lazily is tracked separately in Native Iceberg scan builds every partition's plan data on the driver up front #7523.

How are these changes tested?

Unit tests

PlanDataInjectorSuite has a new merge test: an injected scan resolves the partition's own delete files, paths and partition values through the merged pools. The test also checks that the plan holds only the common block plus that partition.

Integration tests, compared with Spark

New SQL file test sql-tests/iceberg/merge_on_read_partition_layout.sql:
- A merge-on-read table partitioned by identity on five columns: a string with non-ASCII values, a date, a decimal, a timestamp and a timestamp without time zone.
- The values include dates and timestamps before 1970, the epoch, a leap day, both ends of the decimal range, daylight-saving gap and overlap times, and a partition where every value is null.
- It runs with preserve-data-grouping both on and off, with and without a coalesce, in a non-UTC session time zone.
New CometIcebergNativeSuite tests:
- delete files shipped with each partition that uses them, for grouped and ungrouped reads, with the common block holding no per-file entries
- V3 deletion vectors sharing one Puffin file across partitions
- an unpartitioned equality delete followed by a newer append: older files' partitions ship the delete, and the newer file's partition and its matching row are unaffected
Existing assertions that read pools from the common block now read them from each partition's message:
- the shared-Puffin and delete-pool deduplication checks, now checked per block
- the residual deduplication and pushdown checks in CometIcebergNativeSuite and CometFuzzIcebergSuite

… not grow with the table Native Iceberg scans put the whole scan's delete files, delete-file paths and sets, partition values and residuals in the common block. Every task deserialized it from the task binary, and the injector copied it into each partition's native plan, so tasks on large merge-on-read tables ran out of heap. Each partition now carries only the pool entries its tasks use, in IcebergScan.own_pools (a new IcebergFilePools message), and the common block keeps only the small pools and settings. The driver builds each partition's message lazily, while Spark serializes its task (PartitionPlanData / PlanDataSlice), so it never holds every partition's message at once. The executor merges the partition's pools into the common block before planning, and native rejects a plan that still carries own_pools. A coalesce without a shuffle repeats shared delete files once per input partition; this is documented and pinned by a test.

…ck, and hold driver pools as bytes Entries that every non-empty partition uses, such as an unpartitioned equality delete, now ship once in the common block instead of in every task. Every task needs them anyway, so executors pay nothing extra, and tasks no longer grow past spark.rpc.message.maxSize when many shared delete files exist. Partition messages index the common entries first, then their own. After planning, the driver keeps the full pools as the serialized bytes plus an offset per entry, and parses only the entries a task uses when that task is serialized. The parsed message took about 3.9 times its serialized size; bytes plus the index take about 1.1 times. Adds coalesce queries to the merge-on-read SQL test.

…se in the common block An equality delete applies only to older data files, so newer appends, or a bucket with no deletes, leave some partitions without it. Requiring every partition to use an entry sent such deletes back into every task. Entries used by at least half of the non-empty partitions now ship once in the common block. Every entry there is needed by at least half of the tasks, so the common block is at most twice what an average partition's tasks need.

mbutrovich · 2026-11-04T22:59:18Z

Thanks @parthchandra, I'll take a look in the morning!

mbutrovich

Thanks @parthchandra for tracking this down, and for the numbers in #7490. Sending each task only the pool entries its own files use looks like the right fix to me, and I'd like to get there by building on the interning that #4982 and #5853 added.

mbutrovich · 2026-11-05T16:38:41Z

+private[apache] trait PartitionPlanData {
+  def length: Int
+
+  /** Partition i's bytes. Must be safe to call from any thread, and more than once. */
+  def apply(i: Int): Array[Byte]
+
+  /**
+   * Whether `other` holds the same data, decided without building any partition. The default only
+   * knows identity.
+   */
+  def sameAs(other: PartitionPlanData): Boolean = this eq other
+
+  /** Every partition's bytes. For tests and diagnostics only: it builds them all. */
+  def toArray: Array[Array[Byte]] = Array.tabulate(length)(apply)
+}
+
+private[apache] object PartitionPlanData {
+
+  private final class Built(val bytes: Array[Array[Byte]]) extends PartitionPlanData {
+    override def length: Int = bytes.length
+    override def apply(i: Int): Array[Byte] = bytes(i)
+    override def sameAs(other: PartitionPlanData): Boolean = other match {
+      case that: Built =>
+        that.bytes.length == bytes.length &&
+        bytes.indices.forall(i => java.util.Arrays.equals(bytes(i), that.bytes(i)))
+      case _ => false
+    }
+  }
+
+  /** Data already built for every partition. */
+  def apply(bytes: Array[Array[Byte]]): PartitionPlanData = new Built(bytes)
+
+  /** Whether two scans carry the same data. */
+  def same(a: PartitionPlanData, b: PartitionPlanData): Boolean = (a eq b) || a.sameAs(b)
+}
+
+/**
+ * The planning data one task needs, keyed by scan. On the driver it holds only a reference to
+ * every scan's data and builds this partition's bytes while the task is serialized. On the
+ * executor it holds the bytes it received.
+ */
+private[comet] final class PlanDataSlice(
+    @transient private val source: Map[String, PartitionPlanData],
+    index: Int)
+    extends Serializable {
+
+  @transient private var received: Map[String, Array[Byte]] = _
+
+  def get: Map[String, Array[Byte]] =
+    if (received != null) received else source.map { case (key, data) => key -> data(index) }
+
+  private def writeObject(out: ObjectOutputStream): Unit = {
+    out.defaultWriteObject()
+    out.writeObject(get)
+  }
+
+  private def readObject(in: ObjectInputStream): Unit = {
+    in.defaultReadObject()
+    received = in.readObject().asInstanceOf[Map[String, Array[Byte]]]
+  }
+}


Should the lazy driver-side build go in its own PR? #7490 reports executor OOM. PartitionPlanData, PlanDataSlice, SerializedPools, and IcebergPartitionPlanData address driver memory, and they change CometExecRDD, CometNativeShuffleInputRDD, NativeExecContext, CometNativeScanExec, and the Delta contrib scan. They also move each partition's build inside Spark's scheduler lock, as the description notes.

On the table from #7490 (partition delete granularity, grouped read), as I read it each delete file is used by one partition. So eager per-partition messages from the previous comment would total about what the common block holds on main today. The cost of eager messages is driver duplication of entries many partitions share, such as a global equality delete. For context, the driver already holds every FileScanTask and DeleteFile for the life of the plan, because Iceberg's SparkPartitioningAwareScan caches its task groups (line 69, filled at lines 204-213). How were the 1.1x and 3.9x heap figures in the description measured? Driver heap numbers for eager messages on a shared-delete layout would tell us whether the lazy build is needed, and a separate issue and PR could carry them.

Agreed, it's out of this PR and I'll open a separate issue with driver numbers. The 3.9× and 1.1× figures came from a standalone JVM program. It built an IcebergScanCommon with 200k and 1M synthetic position-delete entries (path, delete file, set, partition value), serialized it, and parsed it back. It then measured the used-heap difference with a serial GC, after several System.gc() calls, keeping the bytes alive. The 1.1× figure is the serialized bytes plus a per-entry offset index. These are microbenchmark numbers, not driver numbers from a real job, so the follow-up will measure eager messages on a shared-delete layout.

#7523

… serde pass Review feedback: build each partition's pools in the pass that already interns them, and ship them in the partition's existing IcebergScan.common, so the injector merges the two with mergeFrom. This drops the global-then-split step (IcebergScanPools, Layout, SerializedPools), the IcebergFilePools message and own_pools field, and the half-of-partitions threshold: the common block now holds none of the per-file pools. The lazy driver-side build (PartitionPlanData, PlanDataSlice) moves out of this PR, so CometExecRDD, CometNativeShuffleInputRDD, NativeExecContext, CometNativeScanExec and the Delta contrib scan are unchanged from main.

parthchandra

Thanks Matt. Reworked the PR to follow your approach,

mbutrovich · 2026-11-05T18:49:33Z

Thanks Matt. Reworked the PR to follow your approach,

Thanks @parthchandra! I will take another pass soon!

comphead

Thanks @parthchandra, the write-up of what ships where made this easy to review, and the new revision is much smaller.

One cross-PR note that has no diff anchor here. #7443, #7156, #7233 and #6847 add CometIcebergNativeSuite tests that parse IcebergScanCommon.parseFrom(scan.commonData) and read its residual pool, for example assert(common.getResidualPoolCount == (if (pushed) 1 else 0)) in #7156 and #7233. All four merge with this branch without a textual conflict, but with this change the common block holds no residuals, so I expect those assertions to need the per-partition data, whichever PR lands second. I have not run them.

comphead · 2026-11-05T20:19:25Z

+that its own files need, the way Spark's Iceberg reader does. A coalesce without a shuffle
+(`coalesce(n)`, or a write that coalesces its input) combines several input partitions into one
+task, and that task then carries a copy of a shared delete file for every input partition that
+uses it. Scans where many partitions share delete files, such as tables with unpartitioned
+equality deletes, can make coalesced tasks large. Compacting delete files, or repartitioning
+instead of coalescing, avoids this.


Thanks for documenting this limit. As far as I can tell Spark's reader avoids it. A coalesced task writes all of its input partitions into one ObjectOutputStream (CoalescedRDDPartition.writeObject writes parents through defaultWriteObject), and Iceberg's DeleteFileIndex.forDataFile hands out references to the same DeleteFile objects, so a shared delete file is written once per task and later occurrences are back-references. The latest release ships it once too, in the common block. With the threshold gone, this now also applies to unpartitioned equality deletes, which every partition uses, so a coalesced task over such a table carries one copy per input partition where the release carries one.

Could each partition hold its entries as references to one shared object per entry, rather than as copied bytes, so that Java serialization shares them within a task the same way?

Agreed, this was a regression against the release for coalesced reads. Each delete file is now serialized once per scan, and every partition that uses it holds the same array and path string (IcebergPartitionPlanData). Java serialization writes it once per task stream, so a coalesced task carries one copy, and the driver keeps one copy too. The executor sets file_path_idx when it builds the partition's message, because the path pool order differs per partition. PlanDataInjectorSuite has a serialization test for the coalesced case, and the doc paragraph is updated.

…artition

…k carries each once Review feedback: a coalesce without a shuffle writes several input partitions into one task, and each partition held its own copy of a shared delete file, so the task carried one copy per input partition where the latest release carried one. Per-partition plan data is now a PartitionPlanData object instead of bytes. An Iceberg partition holds its own message plus references to delete files serialized once for the whole scan, so Java serialization writes a shared delete file once per task stream and the driver keeps one copy. The executor sets each delete file's path index when it builds the partition's message, since the path pool order differs per partition. Native Parquet and contrib scans wrap their bytes, so the CometScanWithPlanData trait is unchanged. Move the Iceberg pool test helpers into IcebergScanPlanData, with distinctResiduals for tests that read the residual pool.

parthchandra

Thanks for flagging this. The helpers are now in org.apache.comet.IcebergScanPlanData, with a distinctResiduals(scan) helper, so whichever PR lands second can change getResidualPoolCount == n to distinctResiduals(scan).size == n.

parthchandra · 2026-11-05T23:59:05Z

+that its own files need, the way Spark's Iceberg reader does. A coalesce without a shuffle
+(`coalesce(n)`, or a write that coalesces its input) combines several input partitions into one
+task, and that task then carries a copy of a shared delete file for every input partition that
+uses it. Scans where many partitions share delete files, such as tables with unpartitioned
+equality deletes, can make coalesced tasks large. Compacting delete files, or repartitioning
+instead of coalescing, avoids this.


Agreed, this was a regression against the release for coalesced reads. Each delete file is now serialized once per scan, and every partition that uses it holds the same array and path string (IcebergPartitionPlanData). Java serialization writes it once per task stream, so a coalesced task carries one copy, and the driver keeps one copy too. The executor sets file_path_idx when it builds the partition's message, because the path pool order differs per partition. PlanDataInjectorSuite has a serialization test for the coalesced case, and the doc paragraph is updated.

parthchandra requested review from andygrove, comphead and mbutrovich November 4, 2026 21:44

parthchandra added run-iceberg-tests run-all-spark-profiles Run the Comet test suites against every Spark profile on this pull request, ahead of the merge queue labels Nov 4, 2026

github-actions Bot added bug Something isn't working area:shuffle Shuffle (JVM and native) area:scan Parquet scan / data reading area:Iceberg area:Delta Delta Lake contrib (native scan, deletion vectors) labels Nov 4, 2026

parthchandra changed the title ~~fix: ship Iceberg scan pool entries per partition so task memory does…~~ fix: ship Iceberg scan pool entries per partition so task memory does not grow with the table Nov 4, 2026

parthchandra added 2 commits November 4, 2026 14:06

style: drop an interpolator with no placeholders

b1fc2b8

mbutrovich requested changes Nov 5, 2026

View reviewed changes

parthchandra commented Nov 5, 2026

View reviewed changes

Comment thread native/proto/src/proto/operator.proto Resolved

Comment thread spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala Resolved

Comment thread spark/src/main/scala/org/apache/comet/serde/operator/IcebergScanPools.scala Resolved

mbutrovich self-requested a review November 5, 2026 18:49

comphead reviewed Nov 5, 2026

View reviewed changes

parthchandra added 2 commits November 5, 2026 13:46

docs: note that the driver keeps a copy of a shared delete file per p…

44cf085

…artition

parthchandra commented Nov 6, 2026

View reviewed changes

style: drop an interpolator with no placeholders

d7aafc3

Conversation

parthchandra commented Nov 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mbutrovich commented Nov 4, 2026

Uh oh!

mbutrovich left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mbutrovich Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

parthchandra Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

parthchandra left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mbutrovich commented Nov 5, 2026

Uh oh!

comphead left a comment

Choose a reason for hiding this comment

Uh oh!

comphead Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

parthchandra Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

parthchandra left a comment

Choose a reason for hiding this comment

Uh oh!

parthchandra Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

parthchandra commented Nov 4, 2026 •

edited

Loading