RFC: Reduce Driver Heap Usage with Early Iceberg Task Encoding

[infvg]

Problem

Before an Iceberg query can read data, it must plan the scan, determining which files to read, how to divide them into work for executors and which filters still need evaluation.

Currently, the scan prep path creates java task objects and task groups.
References to these objects remain in the scan, batch and input partitions. Gluten then makes more objects of the same information.

For queries with many files, those objects and their duplicates increase the driver's memory usage.

Goals

1. Stop retaining the complete java task and task group collections after preperation
2. Reduce duplicate representations of file metadata
3. Lower peak driver process memory
4. Reduce retained driver heap after scan preperation

Proposal

Add an early streaming prep path for iceberg reads.

This new path will take the java planning output incrementally and store any information it needs for execution in compact batches. The scan will keep these batches and planning summaries, while the references to java tasks will be removed so they can be garbage collected.

Statistics, partition constructor and executor transport uses these batches.

The plan is to support batch snapshot reads of parquet files without delete files. We can offer this as an optional incremental scan planning config.

Benefit

The main benefit is the lower retained driver heap for scans many files, which means that larger scans/more queries can be run with the same memory budget.

But in general, fewer java objects will be allocated, there should be less metadata copying & smaller task payloads.

Implementation Plan

Objects and replacements

Existing	Representation	Replacement
FileScanTask	A file range to read, including file information, range offset and length, partition info, remaining filters, delete files	execution info stored in the existing protobuf split bytes
ScanTaskGroup	Collection of file range tasks grouped into a unit of work, info about estimated size & row count	A record with entries corresponding to each unit of work + summary of values
Task collections held by the scan + SparkBatch	references to planned tasks and groups used by statistics and partition construction	one planning result containing all work + summaries
SparkInputPartition	A task group, and the read context for it	partition representation with the encoded batches and read context
IcebergLocalFilesNode	gluten's current representation for file locations, ranges, partition values, and other split info for the native scans	split bytes carried through GlutenPartition without IcebergLocalFilesNode

SparkBatch holds multiple ScanTaskGroups. Each group holds multiple FileScanTasks

SparkInputPartition references one group, and the scan caches tasks and groups so the same objects remain reachable through the scan batch and input partitions

Gluten creates IcebergLocalFilesNode, which contains split info, carried inside GlutenPartition

The new path will replace the task/group references, and the intermediate split nodes with the work + summaries.

Driver side

Add an early scan point

Add a scan path before java task caches fill. Consume tasks and groups incrementally, keeping the compact batch + summaries. Statistics and partition requests will reuse this result without rebuilding the complete task collection.

Add a compact planning result

Add a new object eg CompactPlanningResult for each scan request.
It will have:

1. shared read metadata: snapshot identity, schemas, partition specs, typed partition values, required filters
2. encoded file ranges: file paths, formats, sizes, range offsets, and range lengths
3. task assignments: which file ranges belong to each task
4. planning summaries: statistics needed by spark eg estimated rows and bytes

Statistics & partition requests will reuse this result.

Requests that can obtain statistics directly from snapshot metadata will continue to avoid file planning.

Consume java tasks & create batches

Once tasks become available, we validate them and collect required statistics, encode their execution information using the existing protobuf format. Also we'll remove references to task objects once their info is moved.

We will use the existing protobuf format.

after the task info is transferred, remove references to the temporary task object (and the same for the task groups).

We then add a new input partition for the native batches, which contain the task batches and scheduling information.

Each partition will have the information needed for independent retries with this implementation.

Keep compact batches available for retries and recomputation. Planning memory still grows with the amount of encoded work.

Executor side

Task batch decoding

Pass the encoded batches directly to the existing Velox JNI interface.

Executors won't rebuild Java split nodes or Java protobuf objects before passing the bytes to native code.

Draft: #13269

Evidence

I ran some local tests with a early compact prototype vs the existing path:

For 100k parquet files, containing 6.4 million rows:

Measurement	Existing	New path	Reduction
Live driver heap after preperation	121.5 MiB	79.7 MiB	34.4%
Heap added during scan preperation	59.8 MiB	20.3 MiB	66.0%
Sampled peak driver process memory	748.8 MiB	702.2 MiB	6.2%
Warm query latency	6.516 s	5.870 s	9.9%

Memory figures are medians from 3 fresh JVMs per path. Query times are medians from ten measured query per path.

After GC, the new path had no live BaseFileScanTask, SplitScanTask, BaseScanTaskGroup or IcebergLocalFilesNode instances, while the query was still available to execute.

I also ran a larger test with one million files containing 64 million rows

Driver heap limit	Existing	New path
512 MiB	Failed	Failed
576 MiB	Failed	Passed
896 MiB	Failed	Passed
960 MiB	Passed	Passed

Each result ran twice.

Risk

1. The compact result must keep all info needed for statistics, scheduling, filters, execution & retries.
2. Unsupported scans must use a tested fallback. Runtime filtering and storage-partitioned joins must be excluded initially.

Future work

1. Expand to support delete files
2. Runtime filtering
3. Storage-partitioned joins
4. Manifest caching
5. Shared metadata dictionaries and indexed references

[zhouyuan]

@infvg thanks for the proposal, looks promising, could we also extend this optimization to other connectors as well?
Cc: @wangyum @jinchengchenghh

[infvg]

yes, it can definitely be expanded to include them

[infvg]

Raised a draft: #13269