This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 665462d5b6 optimize(parquet): Nested list batching child.write calls
(#10085)
665462d5b6 is described below
commit 665462d5b6810aeda3a9945eb8515141a3920dd3
Author: mwish <[email protected]>
AuthorDate: Tue Jul 21 04:18:02 2026 +0800
optimize(parquet): Nested list batching child.write calls (#10085)
# Which issue does this PR close?
- Closes #10079 .
# Rationale for this change
Optimize nested list call recursive write counts.
# What changes are included in this PR?
Separate list write function to direct-by-offsets and by backward scan
# Are these changes tested?
Covered by existing
# Are there any user-facing changes?
No
---------
Co-authored-by: Copilot Autofix powered by AI
<[email protected]>
Co-authored-by: Andrew Lamb <[email protected]>
---
parquet/src/arrow/arrow_writer/levels.rs | 242 ++++++++++++++-----------------
1 file changed, 110 insertions(+), 132 deletions(-)
diff --git a/parquet/src/arrow/arrow_writer/levels.rs
b/parquet/src/arrow/arrow_writer/levels.rs
index 8918cfb148..7dfb875efd 100644
--- a/parquet/src/arrow/arrow_writer/levels.rs
+++ b/parquet/src/arrow/arrow_writer/levels.rs
@@ -384,131 +384,134 @@ impl LevelInfoBuilder {
return;
}
- // Fast path for "last-level list": when the child has no nested
rep_levels,
- // each child element produces exactly one rep_level entry. We can
batch
- // contiguous non-empty list slots into a single child.write() call,
then
- // fix up the rep_levels at list-slot boundaries using offsets
directly.
- //
- // Kept as a separate function so the compiler can optimize
write_list's
+ // Dispatch to separate functions so the compiler can optimize each
// hot loop independently (function body size affects codegen quality).
if is_last_level {
- Self::write_list_last_level(child, ctx, offsets, nulls, range);
- return;
+ Self::write_list_direct(child, ctx, offsets, nulls, range);
+ } else {
+ Self::write_list_scan(child, ctx, offsets, nulls, range);
}
+ }
- let offsets = &offsets[range.start..range.end + 1];
-
- let write_non_null_slice =
- |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| {
- child.write(start_idx..end_idx);
- child.visit_leaves(|leaf| {
- let rep_levels =
leaf.rep_levels.materialize_mut().unwrap();
- let mut rev = rep_levels.iter_mut().rev();
- let mut remaining = end_idx - start_idx;
+ /// Batch write for lists whose child has no nested repetition.
+ ///
+ /// "direct" means writing the child rep levels using offsets without
scanning.
+ fn write_list_direct<O: OffsetSizeTrait>(
+ child: &mut LevelInfoBuilder,
+ ctx: &LevelContext,
+ offsets: &[O],
+ nulls: Option<&NullBuffer>,
+ range: Range<usize>,
+ ) {
+ let list_start_rep = ctx.rep_level - 1;
- loop {
- let next = rev.next().unwrap();
- if *next > ctx.rep_level {
- // Nested element - ignore
- continue;
- }
+ let emit_non_empty_run = |child: &mut LevelInfoBuilder, run_offsets:
&[O]| {
+ debug_assert!(run_offsets.len() >= 2);
+ let values_start = run_offsets[0].as_usize();
+ let values_end = run_offsets[run_offsets.len() - 1].as_usize();
+ debug_assert!(values_end > values_start);
- remaining -= 1;
- if remaining == 0 {
- *next = ctx.rep_level - 1;
- break;
- }
- }
- })
- };
+ child.write(values_start..values_end);
- // In a list column, each row falls into one of three categories:
- // - "null": the list slot is absent (!is_valid), encoded at def_level
- 2
- // - "empty": the list slot is present but has zero elements
- // (offsets[i] == offsets[i+1]), encoded at def_level - 1
- // - non-empty: the list slot has child values, which are recursed into
- //
- // Consecutive runs of null or empty rows are batched and written
together.
- let write_null_run = |child: &mut LevelInfoBuilder, count: usize| {
- if count > 0 {
- child.visit_leaves(|leaf| {
- leaf.append_rep_level_run(ctx.rep_level - 1, count);
- leaf.append_def_level_run(ctx.def_level - 2, count);
- });
- }
+ // The first element of each list slot needs rep_level =
+ // list_start_rep to mark a new list boundary. Because there's a
1:1
+ // mapping between child elements and rep_level entries, the
position
+ // of each slot's first element is directly computable from
offsets.
+ child.visit_leaves(|leaf| {
+ debug_assert!(leaf.max_rep_level == ctx.rep_level);
+ let rep_levels = leaf.rep_levels.materialize_mut().unwrap();
+ let batch_len = values_end - values_start;
+ let batch_base = rep_levels.len() - batch_len;
+ for slot_offset in run_offsets.iter().take(run_offsets.len() -
1) {
+ let pos = batch_base + (slot_offset.as_usize() -
values_start);
+ rep_levels[pos] = list_start_rep;
+ }
+ });
};
- let write_empty_run = |child: &mut LevelInfoBuilder, count: usize| {
- if count > 0 {
- child.visit_leaves(|leaf| {
- leaf.append_rep_level_run(ctx.rep_level - 1, count);
- leaf.append_def_level_run(ctx.def_level - 1, count);
- });
- }
- };
+ Self::write_list_impl(child, ctx, offsets, nulls, range,
emit_non_empty_run);
+ }
- match nulls {
- Some(nulls) => {
- let null_offset = range.start;
- let mut pending_nulls: usize = 0;
- let mut pending_empties: usize = 0;
+ /// Batch write for lists whose child has nested repetition.
+ ///
+ /// After batch-writing child elements, scans backward through rep_levels
+ /// counting child-element starts to find and stamp slot boundaries.
+ ///
+ /// Scan backward because we don't know start offset before writing.
+ fn write_list_scan<O: OffsetSizeTrait>(
+ child: &mut LevelInfoBuilder,
+ ctx: &LevelContext,
+ offsets: &[O],
+ nulls: Option<&NullBuffer>,
+ range: Range<usize>,
+ ) {
+ let list_start_rep = ctx.rep_level - 1;
- // TODO: Faster bitmask iteration (#1757)
- for (idx, w) in offsets.windows(2).enumerate() {
- let is_valid = nulls.is_valid(idx + null_offset);
- let start_idx = w[0].as_usize();
- let end_idx = w[1].as_usize();
+ let emit_non_empty_run = |child: &mut LevelInfoBuilder, run_offsets:
&[O]| {
+ debug_assert!(run_offsets.len() >= 2);
+ let values_start = run_offsets[0].as_usize();
+ let values_end = run_offsets[run_offsets.len() - 1].as_usize();
+ debug_assert!(values_end > values_start);
- if !is_valid {
- write_empty_run(child, pending_empties);
- pending_empties = 0;
- pending_nulls += 1;
- } else if start_idx == end_idx {
- write_null_run(child, pending_nulls);
- pending_nulls = 0;
- pending_empties += 1;
- } else {
- write_null_run(child, pending_nulls);
- pending_nulls = 0;
- write_empty_run(child, pending_empties);
- pending_empties = 0;
- write_non_null_slice(child, start_idx, end_idx);
+ child.write(values_start..values_end);
+
+ child.visit_leaves(|leaf| {
+ let rep_levels = leaf.rep_levels.materialize_mut().unwrap();
+
+ if leaf.max_rep_level == ctx.rep_level {
+ // This algorithm is the same as write_list_direct.
+ // Use a separate function because the branch code size
would affect codegen
+ // quality of the hot loop of write_list_direct.
+ let batch_len = values_end - values_start;
+ let batch_base = rep_levels.len() - batch_len;
+ for slot_offset in
run_offsets.iter().take(run_offsets.len() - 1) {
+ let pos = batch_base + (slot_offset.as_usize() -
values_start);
+ rep_levels[pos] = list_start_rep;
}
- }
- write_null_run(child, pending_nulls);
- write_empty_run(child, pending_empties);
- }
- None => {
- let mut pending_empties: usize = 0;
- for w in offsets.windows(2) {
- let start_idx = w[0].as_usize();
- let end_idx = w[1].as_usize();
- if start_idx == end_idx {
- pending_empties += 1;
- } else {
- write_empty_run(child, pending_empties);
- pending_empties = 0;
- write_non_null_slice(child, start_idx, end_idx);
+ } else {
+ // Backward scan: count child-element starts (rep <=
ctx.rep_level)
+ // and stamp list_start_rep at slot boundaries.
+ let mut slot_bounds = run_offsets[..run_offsets.len() -
1].iter().rev();
+ let mut next_stamp_at = values_end -
slot_bounds.next().unwrap().as_usize();
+ let mut seen = 0usize;
+
+ for rep in rep_levels.iter_mut().rev() {
+ // Asserting rep >= ctx.rep_level to ensure low level
depth haven't
+ // being written.
+ debug_assert!(*rep >= ctx.rep_level);
+ // Count element starts by skipping nested reps (rep >
ctx.rep_level).
+ //
+ // `==` would also work here: the child is written
before the
+ // parent, so no entry within the batch has rep <
ctx.rep_level.
+ // Benchmarks show no difference, so keep the more
defensive `<=`.
+ if *rep <= ctx.rep_level {
+ seen += 1;
+ if seen == next_stamp_at {
+ *rep = list_start_rep;
+ match slot_bounds.next() {
+ Some(offset) => next_stamp_at = values_end
- offset.as_usize(),
+ None => break,
+ }
+ }
+ }
}
}
- write_empty_run(child, pending_empties);
- }
- }
+ });
+ };
+
+ Self::write_list_impl(child, ctx, offsets, nulls, range,
emit_non_empty_run);
}
- /// Optimized write path for lists whose child has no nested repetition
levels.
- ///
- /// When the child is a leaf (or a struct of leaves), each child element
maps to
- /// exactly one rep_level entry. This lets us batch contiguous non-empty
list
- /// slots into a single `child.write()` call, then stamp the list-start
markers
- /// at positions computed directly from offsets — avoiding per-slot
`write` +
- /// reverse-scan overhead.
- fn write_list_last_level<O: OffsetSizeTrait>(
+ /// Shared run-classification loop for write_list_direct and
write_list_scan.
+ /// Monomorphized per `emit_non_empty_run` closure type, giving the
compiler
+ /// separate optimization contexts for each backfill strategy.
+ fn write_list_impl<O: OffsetSizeTrait>(
child: &mut LevelInfoBuilder,
ctx: &LevelContext,
offsets: &[O],
nulls: Option<&NullBuffer>,
range: Range<usize>,
+ mut emit_non_empty_run: impl FnMut(&mut LevelInfoBuilder, &[O]),
) {
let null_offset = range.start;
let offsets = &offsets[range.start..range.end + 1];
@@ -528,33 +531,6 @@ impl LevelInfoBuilder {
});
};
- let emit_non_empty_run = |child: &mut LevelInfoBuilder, run_offsets:
&[O]| {
- debug_assert!(run_offsets.len() >= 2);
- let values_start = run_offsets[0].as_usize();
- let values_end = run_offsets[run_offsets.len() - 1].as_usize();
- debug_assert!(values_end > values_start);
-
- // Write all leaf values in one batch. Since the child has no
nested
- // rep, this emits (values_end - values_start) rep_levels all equal
- // to ctx.rep_level (= "continuation within list").
- child.write(values_start..values_end);
-
- // The first element of each list slot needs rep_level =
- // list_start_rep to mark a new list boundary. Because there's a
1:1
- // mapping between child elements and rep_level entries, the
position
- // of each slot's first element is directly computable from
offsets.
- child.visit_leaves(|leaf| {
- let rep_levels = leaf.rep_levels.materialize_mut().unwrap();
- let batch_len = values_end - values_start;
- let batch_base = rep_levels.len() - batch_len;
-
- for slot_offset in run_offsets.iter().take(run_offsets.len() -
1) {
- let list_start_pos = batch_base + (slot_offset.as_usize()
- values_start);
- rep_levels[list_start_pos] = list_start_rep;
- }
- });
- };
-
// Classify each slot, detect run boundaries, flush on transition.
#[derive(Clone, Copy, PartialEq)]
enum SlotKind {
@@ -591,7 +567,9 @@ impl LevelInfoBuilder {
}
match nulls {
- Some(nulls) => {
+ // A null buffer without any null can skip the per-slot validity
+ // checks and use the null-free classification loop below.
+ Some(nulls) if nulls.null_count() > 0 => {
let mut run_kind = classify!(0, nulls);
let mut run_start: usize = 0;
for i in 1..num_slots {
@@ -604,7 +582,7 @@ impl LevelInfoBuilder {
}
flush_run!(run_kind, run_start, num_slots);
}
- None => {
+ _ => {
let mut run_kind = if offsets[0] == offsets[1] {
SlotKind::Empty
} else {
