unikdahal commented on code in PR #5780:
URL: https://github.com/apache/datafusion-comet/pull/5780#discussion_r3962403226
##########
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:
I dug a bit further into the pinned iceberg-rust/parquet-rs code here, and I
don't think this bound is safe.
`get_array_memory_size()` measures Arrow-side allocation, whereas the
rolling decision uses Parquet's encoded-size estimate, which also includes
things like dictionary/data-page state. I couldn't find an invariant
guaranteeing the Arrow size is always >= that estimate.
If Parquet reaches the target while `pacing` is still false, the rolling
writer can roll before the next batch is written. `block_rows` would then still
be tracking the previous file, so subsequent 1000-row blocks can be aligned
from the wrong file boundary.
I think this needs to avoid using the Arrow-memory estimate as a correctness
gate, either pace unconditionally, or make the decision from actual writer
state.
--
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]
