Skip to content

optimize(parquet): Nested list batching child.write calls #10085

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mapleFU wants to merge 5 commits into
base: main
Choose a base branch
from
+103 −130
Open

optimize(parquet): Nested list batching child.write calls #10085

Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
24f236b
For performance, split to two fns
mapleFU Jun 6, 2026
d576404
More comments and checking
mapleFU Jun 6, 2026
7f52901
Enhancement
mapleFU Jun 6, 2026
22f7991
Apply suggestions from code review
mapleFU Jun 6, 2026
9e980e6
Update parquet/src/arrow/arrow_writer/levels.rs
mapleFU Jun 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
230 changes: 100 additions & 130 deletions parquet/src/arrow/arrow_writer/levels.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,131 +352,128 @@ 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 not has nested repetition.
///
/// direct means direct write the child rep-levels by offsets without scan.
Comment thread
mapleFU marked this conversation as resolved.
Outdated
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>(

@mapleFU mapleFU Jun 6, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I previously want to remove the branch for write non-nested childs, however benchmark shows that adding a more branch will hurt the performance. So I split it to two functions.

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 same as write_list_direct.
// Use separate function because the branch code size would affect codegen
// quality of the hot loop of write_list_direct.
Comment thread
mapleFU marked this conversation as resolved.
Outdated
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() {
// This can uses `==`, since list write is recursive and the child is written
// before the parent.
Comment thread
mapleFU marked this conversation as resolved.
Outdated
Comment thread
mapleFU marked this conversation as resolved.
Outdated
if *rep <= ctx.rep_level {

@mapleFU mapleFU Jun 6, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I still use <= because benchmark shows no performance enhancement between <= and ==, so just uses <=

seen += 1;
if seen == next_stamp_at {

@mapleFU mapleFU Jun 6, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we can SIMD this in the future, this is not cirtical in this branch. Or we can "batching" if list offsets is large. E.g. checking not one-by-one, and just batch by batch

*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];
Expand All @@ -496,33 +493,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 {
Expand Down
Loading