andygrove commented on code in PR #5780:
URL: https://github.com/apache/datafusion-comet/pull/5780#discussion_r3970067265
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -581,13 +654,224 @@ impl ClusteredBatchSplitter {
}
}
-/// Gather the rows selected by `indices` out of `batch`. A zero-copy
`RecordBatch::slice` would
-/// be cheaper, but the parquet writer's NaN-count visitor reads list/map
children via
-/// `list_array.values()`, which ignores a slice's offset window -- sliced
list-of-float columns
-/// would over-count NaNs. `take` gathers the referenced children into fresh
compacted arrays,
-/// keeping those counts correct.
-fn materialize_run(batch: &RecordBatch, indices: &UInt32Array) ->
DFResult<RecordBatch> {
- arrow::compute::take_record_batch(batch,
indices).map_err(DataFusionError::from)
+/// Cuts contiguous row ranges out of a batch: the partition runs the
clustered writer needs, and
+/// the [`ROWS_DIVISOR`]-row pieces the rolling writer needs.
+///
+/// A zero-copy `RecordBatch::slice` is the cheap way to cut a range, but it
is only exact when
+/// the parquet writer's NaN-count visitor sees every float through the slice.
The visitor reaches
+/// list elements and map entries through `list_array.values()` /
`map_array.entries()`, which
+/// ignore the parent's offset window, so a sliced `list<float>` column would
have every NaN in
+/// the batch counted once per range cut from it -- and since the JVM carries
the native writer's
+/// NaN counts into the manifest, a `nan_value_count` that reaches
`record_count` makes Iceberg's
+/// metrics evaluator prune the file from ordinary comparison predicates.
Struct children are
+/// safe, because `StructArray::slice` slices them, so the only schemas that
need the fix are the
+/// ones with a float or double under a list or map; those ranges go through
`take`, which gathers
+/// the referenced children into fresh compacted arrays.
+#[derive(Clone, Copy)]
+struct RowSlicer {
+ gather: bool,
+}
+
+impl RowSlicer {
+ /// `schema` is the field-id-decorated target schema every batch is cast
to, so this decision
+ /// is made once per task rather than per batch.
+ fn for_schema(schema: &ArrowSchema) -> Self {
+ Self {
+ gather: schema
+ .fields()
+ .iter()
+ .any(|field| float_under_list_or_map(field.data_type())),
+ }
+ }
+
+ fn slice(&self, batch: &RecordBatch, offset: usize, len: usize) ->
DFResult<RecordBatch> {
+ if offset == 0 && len == batch.num_rows() {
+ return Ok(batch.clone());
+ }
+ if self.gather {
+ gather_rows(batch, offset, len)
+ } else {
+ Ok(batch.slice(offset, len))
+ }
+ }
+
+ /// Cuts a range that outlives the batch it came from, because it is
waiting for the rows that
+ /// complete its unit. Always gathers: a zero-copy slice would pin every
buffer of its parent
+ /// batch for the wait, and a partition that receives a handful of rows
per batch would pin one
+ /// parent per batch until its unit fills.
+ fn detach(&self, batch: &RecordBatch, offset: usize, len: usize) ->
DFResult<RecordBatch> {
+ if offset == 0 && len == batch.num_rows() {
+ // The range is the whole batch, so it pins nothing beyond the
rows it holds.
+ return Ok(batch.clone());
+ }
+ gather_rows(batch, offset, len)
+ }
+}
+
+fn gather_rows(batch: &RecordBatch, offset: usize, len: usize) ->
DFResult<RecordBatch> {
+ let indices = UInt32Array::from_iter_values(offset as u32..(offset + len)
as u32);
+ arrow::compute::take_record_batch(batch,
&indices).map_err(DataFusionError::from)
+}
+
+/// How this task cuts rows and when the rolling writer's size check can
matter: everything needed
+/// to open a [`RowPacer`] for one more file.
+#[derive(Clone, Copy)]
+struct PacingPolicy {
+ slicer: RowSlicer,
+ target_file_size_bytes: usize,
+}
+
+impl PacingPolicy {
+ fn pacer(&self) -> RowPacer {
+ RowPacer::new(*self)
+ }
+}
+
+/// Hands one iceberg-rust writer exactly [`ROWS_DIVISOR`] rows at a time, so
the rolling writer
+/// underneath only ever gets to roll on the row boundaries iceberg-java's
`RollingFileWriter` rolls
+/// on -- multiples of 1000 rows since the current file opened -- however
Spark happened to batch
+/// the rows.
+///
+/// Pacing, rather than just cutting each batch up, is what makes the roll
point independent of the
+/// batch shape: a task fed 800-row batches would otherwise be offered a
boundary every 800 rows
+/// and roll into 800-row files where the JVM writer produces 1000-row ones.
Rows left over from a
+/// batch wait here for the rows that complete their unit, so at most
`ROWS_DIVISOR - 1` rows per
+/// open file are held back.
+///
+/// Pacing costs one writer call per 1000 rows instead of one per batch, which
is measurable on a
+/// wide schema, so it stays off until it can change an outcome: see `pacing`.
+struct RowPacer {
+ policy: PacingPolicy,
+ /// Whether rows are being paced yet. While off, batches go over whole,
because a batch
+ /// boundary the writer sees can only matter if its size check might fire
there -- and it cannot
+ /// while the file is still smaller than the target.
+ ///
+ /// It never goes back off. A roll starts the file's size over, but the
pacer cannot see rolls,
+ /// so from the first file that reaches the target it paces for the rest
of the task.
+ pacing: bool,
+ /// What has been handed over so far, as an upper bound on what the open
file holds: Arrow's
+ /// in-memory footprint of a batch (allocated capacity, so it over-counts
if anything) bounds
+ /// what parquet writes for the same rows, since the writer encodes and
compresses them. If that
+ /// ever failed to hold for some schema, the only consequence is a roll up
to `ROWS_DIVISOR`
+ /// rows late -- the same order as the divergence the two size estimates
already allow.
+ handed_bytes: usize,
+ /// Rows of the current 1000-row block already handed over. Non-zero only
for the first block
+ /// after pacing turns on, where it is what the whole batches before it
left in the open file --
+ /// which is all of them, since nothing could have rolled yet. Carrying it
over is what keeps
+ /// the blocks aligned to the file's own row count rather than to where
pacing started.
+ block_rows: usize,
+ /// Rows handed in but not yet handed over; `block_rows + pending_rows <
ROWS_DIVISOR`.
+ pending: Vec<RecordBatch>,
+ pending_rows: usize,
+}
+
+impl RowPacer {
+ fn new(policy: PacingPolicy) -> Self {
+ Self {
+ policy,
+ pacing: false,
+ handed_bytes: 0,
+ block_rows: 0,
+ pending: Vec::new(),
+ pending_rows: 0,
+ }
+ }
+
+ /// The complete units `batch` makes available, in row order. Whatever
does not fill a unit is
+ /// held for the next call.
+ fn push(&mut self, batch: RecordBatch) -> DFResult<Vec<RecordBatch>> {
+ debug_assert!(
+ self.block_rows + self.pending_rows < ROWS_DIVISOR,
+ "a complete unit was left pending"
+ );
+ let rows = batch.num_rows();
+ if !self.pacing {
+ self.handed_bytes += batch.get_array_memory_size();
Review Comment:
You're right, and I've reverted the gate in 7a4a17366 — it was the newest
commit, so `edc5597e3` is gone in its entirety and pacing is unconditional
again.
Your reasoning is the part I got wrong. I was treating
`get_array_memory_size()` as an upper bound on what parquet would have written,
on the grounds that it is capacity-based and parquet encodes and compresses.
But those are different quantities measured on different sides — the rolling
decision uses parquet's encoded-size estimate, which carries dictionary and
data-page state that has no Arrow counterpart — and "usually larger" is not an
invariant. I could not find one either.
The consequence is also worse than I had written down. I had told myself the
penalty for violating the bound was a roll up to 1000 rows late, which sounded
like it sat inside the divergence the two size estimates already allow. That is
wrong: as you say, if the writer rolls while `pacing` is still false,
`block_rows` is still counting the *previous* file, so every later block is
aligned from the wrong boundary. That is a silently wrong grid for the rest of
the task, not a late roll — and the grid is the entire point of the change.
On your second option, deciding from actual writer state: I looked and it is
not reachable against the pinned iceberg-rust. `RollingFileWriter::should_roll`
consults `current_written_size()`, but both that and `current_row_num()` are
private, and Comet only reaches the writer through `DataFileWriter`, so there
is no way to ask it what it has actually written. So unconditional pacing is
the available answer and this PR takes the cost: ~4% at 20 columns and ~10-15%
at 100 columns on the A/B in the description, which I have updated along with
the reverted-commit rationale. The gate can come back if iceberg-rust exposes
the writer's own size.
Thanks for chasing this into the pinned source rather than taking the commit
message's word for the bound.
##########
docs/source/user-guide/latest/iceberg-writes.md:
##########
@@ -157,7 +157,7 @@ A write is eligible only when ALL of the following hold:
| `write.parquet.bloom-filter-enabled.column.<col>`
| unset or
`false`
|
| `write.metadata.metrics.*`
| any value
(manifest metrics are re-derived on the JVM with Iceberg's own logic)
|
| `write.spark.fanout.enabled`
| any value (the
native writer implements both clustered and fanout modes)
|
-| `write.target-file-size-bytes`
| any value
(file rolling cadence differs; see accepted divergences)
|
+| `write.target-file-size-bytes`
| any value (the
roll point can differ by less than 1000 rows; see accepted divergences)
|
Review Comment:
Fixed in 7a4a17366. You're right that the grid does not bound the distance
between the two writers' roll points, and the table was promising something
this change does not establish.
The table row now reads:
> `write.target-file-size-bytes` | any value (the two writers can choose
different roll points; see accepted divergences)
and the divergences bullet documents the shared grid without a numeric
cross-writer bound:
> ... so each writer rolls only on a 1000-row boundary of its own file.
> The shared grid is all that is shared. What each writer compares against
the target differs — flushed bytes plus parquet-rs's estimate of the open row
group, versus parquet-mr's file position plus its buffered size — and the two
use different threshold comparisons. These are independent size estimates, so
nothing bounds how far apart the two writers' roll points are: they may cross
the target several grid steps apart, and the resulting files can differ in row
count by an arbitrary number of 1000-row blocks. Do not rely on file-layout
parity between the two writers; rely only on each file rolling on its own
1000-row boundary.
I also took the one-step claim out of the PR rationale in both places it
appeared — the diagnosis paragraph and the docs-changes bullet.
Worth noting these two reviews turned out to be the same mistake in two
places. The "one 1000-row step" framing is exactly what made the Arrow-size
gate in @unikdahal's thread look acceptable to me: I reasoned that violating
its bound cost a roll up to 1000 rows late, which sounded like it fell inside a
divergence I had already documented as bounded. It was not bounded, and the
gate's real failure mode was a misaligned grid rather than a late roll. That
commit is reverted in the same push.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
