Repository navigation
fix: have the native Iceberg write read a converted child's Arrow stream - #7536
zhangfengcdt wants to merge 2 commits into
Conversation
A native Iceberg write over a Spark-to-Arrow conversion wrapped the conversion's executeColumnar batches in a reader that closes each batch once native has it. The conversion reuses its vectors for the next batch, and closing a struct vector drops its children, so the second batch of a partition failed to import with "no more field nodes". The write now reads such a child's Arrow stream, as native shuffle does since apache#6689.
comphead
left a comment
There was a problem hiding this comment.
The change follows #6689, and the stream is still built by CometArrowStream.stream, so I found no new leak against the old path. Since the bug is live only on the release branches, could the PR carry the backport- labels for branch-1.2 (and branch-1.1 if the fix applies there), as backporting.md asks?
| // Native reads the source's Arrow stream, as it does for a native operator's input. The | ||
| // source reuses its vectors for the next batch, and wrapping its `executeColumnar` batches | ||
| // below closes each batch once native has it; closing a struct vector drops its children, | ||
| // so the next batch would lose them (#7239). The scheduler probe counts rows as batches. | ||
| case source: CometNativeArrowSource if schedulerProbe.isEmpty => |
There was a problem hiding this comment.
As the description says, on main this failure no longer happens since #7325, so the comment describes the release branches only. Could it say why the stream route is still preferred on main (it skips the per-batch VectorSchemaRoot re-wrap and unload/load in readerBatchIter and ColumnarBatchArrowReader) and point to #7239 for the older branches?
| // source reuses its vectors for the next batch, and wrapping its `executeColumnar` batches | ||
| // below closes each batch once native has it; closing a struct vector drops its children, | ||
| // so the next batch would lose them (#7239). The scheduler probe counts rows as batches. | ||
| case source: CometNativeArrowSource if schedulerProbe.isEmpty => |
There was a problem hiding this comment.
With schedulerProbe.isEmpty in the match, every test that installs the probe over a VALUES or RDD write (CometLocalTableScanExec and CometSparkToColumnarExec are both CometNativeArrowSource) now takes the wrapped route. So the speculation, executor-loss and scheduler failure tests no longer exercise the route production uses for those children. Which suites cover failure and cleanup on the stream route, for example the handoff failpoint tests, which do not install the probe?
| // https://github.com/apache/datafusion-comet/issues/7239. A conversion of Spark rows writes | ||
| // every batch into the same vectors, and closing a struct vector drops its children. The write | ||
| // reads the conversion's Arrow stream, as native shuffle does, rather than batches it would | ||
| // close once native had them, which left the second batch of a partition without them. | ||
| test("native write reads a struct column that a conversion writes over several batches") { |
There was a problem hiding this comment.
On main the old route already hands out independent vectors since #7325, so I expect this test to pass there with and without the change (I have not run it). It guards the release branches, so could the backport PRs confirm it fails there without the fix? The #6689 test also ran a local table scan (CometLocalTableScanExec), which reaches the same branch through the trait. Looping over both, with more than one partition, would cover the new case fully.
andygrove
left a comment
There was a problem hiding this comment.
The PR backlog is huge and the maintainers don't have the bandwidth to review them all, so here is an automated review. Feel free to ignore it.
Reading the Arrow stream of a CometNativeArrowSource child instead of closing its executeColumnar batches is the right fix and matches what #6689 did for shuffle. I did not find a new leak or ownership problem in the stream route, since CometArrowStream.stream is what the old wrapped path already ended in.
The PR no longer merges with main. #6926 moved the native input handling out of doExecuteColumnar's mapPartitionsInternal body and into CometIcebergWriteRDD, which calls CometExecRDD.resolveInputObjects and takes the first ArrowArrayStream of each partition. The Array[Object](iter.next()) change and the streamRDD.mapPartitionsInternal loop in this diff have nothing to apply to now. After the rebase, the selection looks like it belongs in the non-shuffle else branch where input is built today. There, a CometNativeArrowSource child would return source.doExecuteAsArrowStream() and every other child would keep the IcebergSchedulerProbeRDD plus wrapColumnarBatchRDD path. The ShuffleScan branch has to stay untouched, because resolveInputObjects needs the CometShuffledBatchRDD itself.
Could you rebase onto current main and move the change into that branch? The new test then needs to pass on the merged code. It would also help to confirm that the stream route still hands CometIcebergWriteRDD exactly one stream per partition, since resolveInputObjects throws on an empty partition.
On behavior against branch-1.2, the only change I see is the intended one. A struct column from an RDD-backed or local-table conversion is no longer read through closed vectors after the first batch. I found no change in results, fallback decisions, configs or memory use for other child operators. The earlier reviewer's questions about the scheduler probe and the test coverage are still open and I have nothing to add to them.
…ow-stream The stream route moves into the branch that builds the Scan leaf's input, and the scheduler probe follows it: beforeInput runs ahead of the stream export and afterInput comes from a reader wrapper behind the stream.
Which issue does this PR close?
Closes #7239.
Rationale for this change
A native Iceberg write over a Spark-to-Arrow conversion wrapped the conversion's
executeColumnarbatches in a reader that closes each batch once native has it. The conversion reuses its vectors for the next batch, and closing a struct vector drops its children, so the second batch of a partition fails withno more field nodes. Native shuffle had the same bug, fixed in #6689 by reading the child's Arrow stream.On
mainthe symptom is hidden since #7325, which copies each batch into fresh vectors. The release branches have neither change, so the bug is live there. Onmainthe stream route still saves that copy and the unload and load back inColumnarBatchArrowReader.What changes are included in this PR?
CometIcebergWriteExecreads aCometNativeArrowSourcechild throughdoExecuteAsArrowStream(), as fix: have native shuffle read a converted child's Arrow stream #6689 does for shuffle, in the branch that builds theScanleaf's input. Any other columnar child is wrapped into one Arrow stream per partition, as before. TheShuffleScanroute is unchanged.beforeInputruns inCometIcebergWriteRDD.computeahead of the stream export, as it does ahead of a shuffle fetch, andafterInputis reported by a reader wrapper behind the exported stream.CometNativeArrowSource.doExecuteAsArrowStreamtakes an optional reader wrapper for that.How are these changes tested?
CometIcebergWriteActionSuite: a struct column written from an RDD conversion and from a local table scan, each in several partitions of several batches, withspark.comet.batchSize=7. With the batch reader as it was before feat: preserve columnar batches across observations #7325 the RDD case fails withno more field nodes for field a, and passes with this change. The backport PRs will show it failing on the release branches without the fix.CometIcebergWriteActionSuite: the scheduler probe over a conversion the write reads directly sees each partition start, the rows it supplied and its handoff.CometIcebergWriteActionSuite,CometIcebergWriteDetectionSuiteand the struct tests ofCometNativeShuffleSuitepass on Spark 4.1.