feat: implement native Iceberg V2 writer via iceberg-rust - #5361
Conversation
|
@andygrove here is 3/3! I'd greatly appreciate a CI kick. cc @parthchandra , @comphead , @mbutrovich . @unikdahal - I have addressed your comments regarding CTAS and RTAS in this one, thank you for helping to validate this feature!! |
8416066 to
b092202
Compare
|
Sigh..sorry @andygrove I needed to cargo fmt, feel free to kick off the CI whenever |
andygrove
left a comment
There was a problem hiding this comment.
Thanks for pushing this all the way through, and for splitting it into three reviewable parts. The JVM-side metrics rebuild is a really nice idea: re-deriving each DataFile's metrics through the version-matched ParquetUtil.footerMetrics + MetricsConfig makes manifest metadata iceberg-java's decision by construction rather than something we have to imitate, and it let you drop the counts/none gates entirely. The comments throughout are genuinely helpful, especially the ones recording why a piece of reflection is shaped the way it is.
A few things I verified while reading, recorded here so nobody re-does the work:
DataFiles.Builder.copydoes carrysortOrderIdandsplitOffsetsacross (DataFiles.java:203on 1.8.1), so the stamp-then-rebuild ordering works.ParquetUtil.footerMetricsfalls back to footer stats for fields absent fromfieldMetrics, and sourcesnanValueCountsonly fromfieldMetrics, which matches whatbuildFloatFieldMetricssynthesises.ctx.propertiesistableProperties ++ writeProperties, sorequireParseableCompressionLeveldoes see per-write option overrides.- The
arrowcast indecorate_batch_with_field_idsshort-circuits when types already match, so it is not a per-batch copy.
CI
Lint is red on both Linux and macOS. I reproduced it locally with rustfmt --check: it is cargo fmt on native/core/src/execution/operators/iceberg_write.rs in three spots.
- around line 406, the
parts.sort_by_key(...)closure fits on one line - around line 551, the
ManifestWriterBuilder::new(...)call wants collapsing - around line 773, a double blank line before the
integrationmodule comment
make format should clear all three. iceberg_common.rs and iceberg-writes.md are both already clean.
The rest of my comments are inline.
| requireRowGroupCheckMaxRecordCountAtDefault, | ||
| requireParquetPageVersionDefault, | ||
| requireShredVariantsDisabled, | ||
| requireParseableCompressionLevel, |
There was a problem hiding this comment.
I think a table with a uuid column would be judged eligible here and then fail the task rather than falling back.
Iceberg's TypeToSparkType maps uuid to StringType, so Comet hands the native writer a Utf8 column, while schema_to_arrow_schema makes the target FixedSizeBinary(16). Since decorate_batch_with_field_ids casts with safe: false, a 36-character UUID string has nowhere to go. The gate is entirely table-property based with no schema-type check, so detection would return Compatible first and the failure would surface at execution time.
Have you tried a table with a uuid column? If it does fail, a trigger rule that declines on column types the native writer cannot reproduce (plus a detection test pinning it) would keep this a fallback rather than a task error. fixed(N) should be fine via Binary -> FixedSizeBinary as long as every value is exactly N bytes, but it would be good to have that pinned too.
There was a problem hiding this comment.
You're right — a uuid column passed detection and then failed the task inside decorate_batch_with_field_ids (Arrow has no Utf8 -> FixedSizeBinary(16) cast). Added a type-based trigger rule: detection now walks the write schema (nested fields included) and declines any uuid column with a fall-back reason. Pinned by a new detection test that evolves a uuid column onto a table through the Iceberg API (Spark DDL can't declare one) and asserts the write falls back and still succeeds through iceberg-java.
fixed(N) is pinned too: a new round-trip test writes exact-N-byte values through the native path and compares row values and readable_metrics bounds against a JVM-written twin — Binary -> FixedSizeBinary(4) casts fine as long as lengths match, and a wrong-length value fails the task the same way the JVM writer rejects it. I audited the remaining Iceberg types: uuid is the only Spark-writable type with a cast gap (time / timestamp_ns can't be written through Spark's Iceberg integration at all, and the V3-only types are behind the format-version gate).
| } | ||
| } | ||
|
|
||
| test("native acceleration: complex types (struct, array, map) round-trip with field IDs") { |
There was a problem hiding this comment.
The native-acceleration tests all write (id INT, region STRING, amount DOUBLE), and this complex-types test adds STRING and INT leaves, so the type surface actually exercised through the native writer is fairly narrow.
Would you mind adding a round-trip parity test over the wider primitive set: DATE, TIMESTAMP, TIMESTAMP_NTZ, DECIMAL at a couple of precisions, BINARY, BOOLEAN, BIGINT, FLOAT? The one I would most like pinned is timestamps under a non-UTC spark.sql.session.timeZone, with both timestamp and timestamptz columns. That is exactly where a silent value shift would hide, and none of the current tests would notice it. Your manifestMetricsParity helper looks like it would adapt nicely to compare row values as well as metrics.
There was a problem hiding this comment.
Added native acceleration: wide primitive types keep JVM-parity values and manifest metrics: BIGINT, BOOLEAN, FLOAT, DECIMAL(9,2), DECIMAL(38,10), DATE, TIMESTAMP, TIMESTAMP_NTZ, BINARY written under spark.sql.session.timeZone=America/New_York, comparing both the full row sets and the aggregated readable_metrics bounds/counts against a JVM-written twin. On top of the twin parity there's an absolute check: the timestamptz value must read back as the same zoned instant it was written as and the ntz value must pass through untouched, so a shift that happened to hit both paths equally would still fail. The fixed(N) test from the other thread rides the same helper shape. Green on all four Spark/Iceberg profiles.
| let (calculator, partition_type) = clustered_order | ||
| .expect("clustered order helper must be Some for clustered writes"); | ||
| let order = partition_first_occurrence_order(&batch, calculator, partition_type)?; | ||
| parts.sort_by_key(|(key, _)| order.get(key.data()).copied().unwrap_or(usize::MAX)); |
There was a problem hiding this comment.
This computes the partition values a second time. partition_first_occurrence_order runs PartitionValueCalculator::calculate over the whole batch, but RecordBatchPartitionSplitter::split (two lines up) has already evaluated the same partition transforms internally. Clustered is the default for partitioned tables, so this doubles the partition-transform cost on the common path.
Is there a way to recover the ordering without the second pass? The input is partition-sorted by construction, so anything that gives each part's first row index would do. Alternatively, a single calculate whose result feeds both the split and the ordering, if iceberg-rust exposes that shape.
There was a problem hiding this comment.
Fixed — the clustered path no longer uses RecordBatchPartitionSplitter at all. A new ClusteredBatchSplitter runs PartitionValueCalculator::calculate once per batch, walks the literals for contiguous-run boundaries (preserving batch order by construction, so the sort and the first-occurrence map are gone too), and emits one part per run. Input that isn't actually clustered produces multiple runs with the same key and surfaces the same ClusteredWriter error as before.
One deliberate non-optimisation, recorded in a comment on materialize_run: the runs are materialised with filter_record_batch (what the splitter produced before) rather than zero-copy RecordBatch::slice. iceberg-rust's NaN-count visitor reads list/map children via list_array.values(), which ignores a slice's offset window, so sliced list-of-float columns would over-count NaNs in the manifest. The transform double-compute is gone; the copy the splitter was already paying stays.
| * normal `SparkWrite` path populates from the table's `outputSortOrderId`; we mirror that here | ||
| * by writing the same value via reflection before handing files to the committer. | ||
| */ | ||
| def stampSortOrderId(dataFiles: java.util.List[_], sortOrderId: Int): Unit = { |
There was a problem hiding this comment.
This is a step beyond the rest of the reflection in this file, which reads fields and calls package-private constructors but does not mutate iceberg-java state. An uncaught NoSuchFieldException here would also be a task failure rather than a fallback, unlike newDataManifestFile which soft-fails on exactly that.
DataFiles.Builder.withSortOrder(SortOrder) is public, and rebuildDataFilesWithJavaMetrics already builds through that builder immediately afterwards. Could the sort order be applied there instead, by resolving the SortOrder on the driver (it is Serializable) and shipping it into the task closure alongside metricsConfig? That would drop the private-field write entirely. If you would rather keep the current shape, could it at least soft-fail on NoSuchFieldException for consistency?
There was a problem hiding this comment.
Done — stampSortOrderId (and its setAccessible field write) is gone. The driver resolves the write's SortOrder from Table.sortOrders() (falling back to SortOrder.unsorted() for id 0, which isn't always in the map), ships it in the task closure (SortOrder is Serializable), and rebuildDataFilesWithJavaMetrics applies it through the public DataFiles.Builder.withSortOrder it was already building through. Verified withSortOrder(SortOrder) is identical on 1.5.2 / 1.8.1 / 1.10.0 / 1.11.0. The existing sort_order_id JVM-parity test still passes on all four profiles.
| task_attempt_id, | ||
| output_schema, | ||
| plan_properties, | ||
| metrics: ExecutionPlanMetricsSet::new(), |
There was a problem hiding this comment.
I do not see anything writing to this ExecutionPlanMetricsSet, so metrics() always returns an empty set, and CometIcebergWriteExec passes CometMetricNode(metrics, Nil) with no child nodes either. That leaves no native-side visibility at all for the writer.
Some counters would be valuable here: time spent inside the iceberg-rust writer stack, rows and bytes handed to parquet-rs, files rolled. If you would rather defer that, dropping the field for now would be clearer than keeping a metrics set that never reports anything.
There was a problem hiding this comment.
Wired up. The native operator now reports write_time — time inside the iceberg-rust writer stack (write + close), excluding waiting on upstream — which flows through the root CometMetricNode by name into a new nano-timing SQL metric ("time in native Iceberg writer") on CometIcebergWriteExec. Rows / bytes / files-rolled deliberately stay JVM-derived: they're already surfaced as numFiles / numOutputRows / numOutputBytes from the decoded manifest (the committed values, which is what the stock Spark UI row shows), so native counters with the same meaning would just double-report.
| def fromCometPlan(cometPlan: SparkPlan): CometMetricNode = { | ||
| val children = cometPlan.children.map(fromCometPlan) | ||
| CometMetricNode(cometPlan.metrics, children) | ||
| def fromCometPlan(cometPlan: SparkPlan): CometMetricNode = cometPlan match { |
There was a problem hiding this comment.
The AQE reasoning in the doc comment makes sense, but this changes behaviour for every caller, not just the Iceberg write: CometNativeScanExec, CometShuffleExchangeExec, CometNativeWriteExec and operators.scala all come through here, and a whole subtree now silently reports Map.empty where it previously would have contributed metrics.
I went looking and could not find a non-CometPlan node living inside a native block today, so I think it is safe in practice. Could you either narrow the guard to the node type that actually triggers this (or guard the metrics access itself), or add a test pinning that no existing operator's metrics regress? A silent metric drop is the kind of thing that goes unnoticed for a long time.
There was a problem hiding this comment.
Narrowed to the metrics access itself. fromCometPlan now recurses through every node exactly as upstream does; the only remaining change is that a node whose session is null (the actual NPE trigger — metrics is a lazy val that calls SQLMetrics.createMetric(sparkContext, ...)) contributes an empty metric map while its subtree is still walked. Any node with a live session — Comet or not — reports exactly as before, so no existing operator's metrics can regress, and the null-session node's children are no longer pruned either (they were under the old guard).
There was a problem hiding this comment.
Correction to the above after the first full CI run: pre-checking session != null was wrong. fromCometPlan also runs inside task closures on executors, where the @transient session is always null after deserialisation — but metrics is safe there because the lazy val was materialised on the driver and shipped with the plan. The pre-check therefore blanked every native operator's metrics on executors, which CometTaskMetricsSuite's "native parquet write reports task-level output metrics" caught on all five [exec] shards.
Now the guard wraps the metrics access itself and only swallows the NPE from forcing an unmaterialised lazy val with a null session (the original AQE stage-finalisation case). Every reachable metrics map — driver or executor, Comet node or not — reports exactly as upstream did; CometTaskMetricsSuite passes locally again.
Separately, the macos-14 [scans] job died in a hotspot-level crash ("error occurred during error reporting, SIGTRAP") right after an unrelated timestamp suite — that one just needs a re-run.
| // dispatch key / access mode are inert here. | ||
| let memory_io = load_file_io( | ||
| &std::collections::HashMap::new(), | ||
| "memory:///", |
There was a problem hiding this comment.
Small thing to confirm: does this memory:// FileIO own its backing store, so it is dropped when the function returns, or is opendal's memory service process-global? If the store is shared, every task's manifest bytes would stay resident for the lifetime of the executor, since nothing deletes the entry after the read() below.
There was a problem hiding this comment.
It owns its store. opendal's memory service creates a fresh MemoryCore { data: Arc<Mutex<BTreeMap>> } inside Builder::build() — there's no process-global registry — and load_file_io builds a new FileIO per call, so the manifest bytes are freed when the function returns and the FileIO/operator drops. Added a sentence to the comment at the call site so the next reader doesn't have to re-derive it.
b092202 to
2035caf
Compare
|
@andygrove comments all addressed again, happy to do another round of CI! Thanks for iterating so quickly here! |
2035caf to
810000b
Compare
|
@andygrove all seven review comments are addressed — each thread above has a reply describing the change. The first full CI run also caught a real bug in my Local verification on
Could you approve a CI run when you get a chance? One caveat from the last run: the macos-14 [scans] job died in a hotspot-level runner crash ( |
andygrove
left a comment
There was a problem hiding this comment.
Thanks for turning all of that around so quickly. Every one of the seven points is addressed, and a few of the answers taught me something: I had not realised opendal's memory service has no global registry, and the correction on session != null not being a usable pre-check (because fromCometPlan also runs in task closures where session is always null) is exactly the kind of thing that only a full CI run surfaces. Thanks for writing that reasoning into the doc comment rather than just fixing it.
I checked the claim in your new requireNoUuidColumns comment that uuid is the only Spark-writable Iceberg type with an Arrow mismatch, since that is the sort of thing that quietly stops being true. It holds: TypeToSparkType still throws UnsupportedOperationException for TIME as of Iceberg 1.11, so a time column cannot be planned by Spark at all, and everything else new in 1.11 (variant, geometry, geography, the nanosecond timestamps, unknown) is V3 and already excluded by the format-version gate. So the gate is complete for V1 and V2, and the comment is accurate.
I also confirmed the write_time plumbing works the way you describe. subset_time("write_time", partition) matches the existing convention in scan.rs and parquet_writer.rs, and CometMetricNode.set_all zips the native children against the JVM children, so the extra native Projection and Scan nodes under a Nil-children root are silently ignored rather than mismatched. Unknown metric names are a debug no-op, so nothing else is disturbed.
On formatting: cargo fmt --all -- --check now passes on this head. I re-ran it against the previous head to be sure the original report was real, and it was. Worth noting that CI has not run at all on 810000ba yet, so none of this is machine-verified beyond the format check.
Three things left, all on the new code.
| /// would over-count NaNs. `filter_record_batch` materialises compacted children (exactly what | ||
| /// `RecordBatchPartitionSplitter` produces), keeping those counts correct. | ||
| fn materialize_run(batch: &RecordBatch, start: usize, len: usize) -> DFResult<RecordBatch> { | ||
| let mut mask = vec![false; batch.num_rows()]; |
There was a problem hiding this comment.
The run-boundary split is a much better shape than what it replaced, and the comment on why a plain slice is unsafe here is a good catch. But I think this reintroduces the cost problem from the other direction.
materialize_run allocates a Vec<bool> the length of the whole batch, bit-packs it into a BooleanArray, and then filter_record_batch has to scan that full-length predicate to build its filter strategy. That is O(batch rows) per run rather than O(run length), so a batch with R runs costs O(R x batch_rows) across every column. For a low-cardinality partition that is nothing, but a clustered write on something like bucket(1000, id) or an hourly partition over a wide range can easily produce hundreds of runs in one 8192-row batch, and then this is a good deal more work than the two transform evaluations I originally asked you to remove.
Would building one sequential UInt32Array of 0..num_rows per batch and then take-ing indices.slice(start, len) per run work here? take on a list or map column gathers the referenced children into fresh compacted arrays, same as filter does, so the NaN-visitor concern in your comment should still be satisfied, and each run becomes O(run length) with the total back to O(batch rows).
Separately, PartitionKey::new(self.partition_spec.as_ref().clone(), ...) clones the whole PartitionSpec once per run. Small next to the above, but it scales the same way, and it would be nice not to have both.
There was a problem hiding this comment.
Done, exactly as you sketched: one sequential UInt32Array of 0..num_rows per batch, and each run gathers through take_record_batch(batch, indices.slice(start, len)) — O(run length) per run, O(batch rows) per batch total, and take compacts list/map children the same way filter did, so the NaN-visitor constraint holds. Single-run batches (the common case) short-circuit to a zero-copy batch.clone() — arrays keep their original zero offsets there, so the slice hazard doesn't apply.
On the PartitionSpec clone: I couldn't remove it — PartitionKey::new takes its spec by value, and copy_with_data clones internally too, so one spec clone per run is the floor with the current iceberg-rust API. Left a comment saying so; it's now the only per-run cost that isn't O(run length), and specs are a handful of fields.
| CometMetricNode(cometPlan.metrics, children) | ||
| val nodeMetrics = | ||
| try cometPlan.metrics | ||
| catch { case _: NullPointerException => Map.empty[String, SQLMetric] } |
There was a problem hiding this comment.
This is much better than pruning the subtree, and the note about why the session pre-check does not work is worth having. My remaining worry is that the catch is wider than the cause you diagnosed. It will also swallow an NPE thrown from inside any operator's own metrics definition, so a genuine bug in a Comet operator's metric construction turns into that operator silently reporting nothing instead of failing loudly.
Since the failure you are guarding is specifically a foreign JVM node built off the planning thread, would gating on that narrow the blast radius enough? Something like only swallowing when the node is not a CometPlan. Comet operators reach here either on the driver with a live session or on executors with metrics already materialised, so neither case should need the guard, and a real NPE inside a Comet operator's metrics would still surface.
There was a problem hiding this comment.
Narrowed as suggested: CometPlan nodes read metrics unguarded, so a genuine NPE inside a Comet operator's metric construction fails loudly; only foreign nodes get the catch. Doc comment updated with the reasoning (Comet operators reach here on the driver with a live session or on executors with metrics already materialised, so neither case needs it). CometTaskMetricsSuite still passes on all four profiles.
| } | ||
| } | ||
|
|
||
| test("native acceleration: empty append commits exactly once with zero data files") { |
There was a problem hiding this comment.
This pins the empty append on an unpartitioned table, so it exercises UnpartitionedWriter::close() with nothing written. The partitioned equivalent goes through ClusteredWriter::close() or FanoutWriter::close() having never seen a batch, and that path is not covered anywhere. An INSERT INTO ... SELECT whose filter matches nothing is an entirely ordinary thing to do against a partitioned table, so if either of those writers errors or panics with no current writer it would be a hard failure on a common case.
Could you add a partitioned variant here? It should be a couple of lines given the helper already exists, and it would also cover the ClusteredBatchSplitter never being invoked at all.
There was a problem hiding this comment.
Added native acceleration: empty append to a partitioned table commits with zero data files, covering both partitioned writers: an identity-partitioned table on the default clustered path and a fanout twin (write.spark.fanout.enabled=true), each fed an INSERT INTO ... SELECT matching nothing. Both commit exactly once with zero data files through CometIcebergWriteExec, with the splitter never invoked. I also read both writers' close() at the pinned rev — ClusteredWriter takes the (absent) current writer and FanoutWriter iterates an empty map, so both return empty output rather than erroring — and the test now pins that. Green on all four profiles.
810000b to
85cecd1
Compare
|
@andygrove good for a CI kick |
|
@andygrove failure seems transient |
Which issue does this PR close?
(Third of three) Closes #4322 and closes #5308. Part one was #4658 (split-operator plan, merged); part two was #5298 (detection + fall-back reasons, merged).
Rationale for this change
We want to speed up writing iceberg data for ETL jobs and table maintenance!
With detection in place (#5298), this PR wires in the native writer itself: when
spark.comet.iceberg.write.enabled=trueand a write passes the eligibility gate, theIcebergWriteoperator's per-task Parquet write is delegated to iceberg-rust through Comet's native execution pipeline. Metadata writing (this diverges from my original approach where I had rust responsible for writing it all) and commit semantics are unchanged and go through the JVM, ensuring identical semantics to Java.What changes are included in this PR?
High level design:
In some areas, iceberg-java and iceberg-rust just simply can't be byte identical. Those are documented in
iceberg-writes.md.How are these changes tested?
CometIcebergWriteActionSuitegains a "native acceleration" section (~20 new tests): appends (unpartitioned/partitioned/fanout with unsorted input), overwrite static/dynamic, CoW DELETE/UPDATE natively and a pin that CoW MERGE intentionally falls back (MergeRowsExecis not Comet-native), CTAS/RTAS on Spark 3.5+ (pinning the empty-metadata-location path), complex types with nested field IDs, multi-file rolling viawrite.target-file-size-bytes, empty appends (one commit, zero files), andsort_order_idparity with the JVM writer (0 through Iceberg 1.10; the resolved order id on 1.11+).truncate(16), metrics modenonewith a per-columnfulloverride, and modecountswithWRITE ORDERED BY(pinning the sorted-column promotion insideMetricsConfig.forTable).CometIcebergRewriteActionSuite) now assert compaction data files were physically written by iceberg-rust via the parquet footercreated_bysignature.FieldMetricsconstructors, and 1.11 changed whatsort_order_idthe Spark writer stamps).