From 87996f3371c3c7b504626218af38bc5e990f85c6 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Sat, 22 Aug 2026 20:50:52 -0400 Subject: [PATCH] perf(row): compact variable-list keys Signed-off-by: Nicholas Gates --- docs/specs/row-encoding.md | 16 ++-- vortex-row/src/codec.rs | 163 +++++++++++++++++++++++++++++-------- vortex-row/src/encode.rs | 5 +- vortex-row/src/size.rs | 17 ++-- vortex-row/src/tests.rs | 23 ++++++ 5 files changed, 173 insertions(+), 51 deletions(-) diff --git a/docs/specs/row-encoding.md b/docs/specs/row-encoding.md index fcfd1d6eec9..e555735ac26 100644 --- a/docs/specs/row-encoding.md +++ b/docs/specs/row-encoding.md @@ -403,13 +403,19 @@ A variable-size list is ordered lexicographically by its elements. Null and empt the variable-width sentinels. A non-empty list is encoded as: ```text -varlen_non_empty_sentinel || escaped_elements || list_terminator +varlen_non_empty_sentinel + || element_marker || encoded_element_0 + || element_marker || encoded_element_1 + || ... + || list_terminator ``` -Each byte of each recursively encoded element is escaped as `0x01 || byte`. The terminator is -`0x00` for ascending fields and `0x02` for descending fields. This makes a shorter list that -is an element-wise prefix sort before the longer list in ascending order and after it in -descending order, without allowing a following column to affect that comparison. +The element marker is `0x01`. Element row keys are self-delimiting and prefix-safe, so their +bytes are copied without further escaping. The terminator is `0x00` for ascending fields and +`0x02` for descending fields. It therefore sorts before or after the next element marker, +making a shorter list that is an element-wise prefix sort before the longer list in ascending +order and after it in descending order, without allowing a following column to affect that +comparison. Element encodings use the same `RowSortField` as the list. Consequently, nested null placement remains independent of sort direction, consistent with structs and fixed-size lists. diff --git a/vortex-row/src/codec.rs b/vortex-row/src/codec.rs index 2ec0d297c8c..6e722ed7f4d 100644 --- a/vortex-row/src/codec.rs +++ b/vortex-row/src/codec.rs @@ -72,11 +72,11 @@ pub(crate) const VARLEN_NULL_SIZE: u32 = 1; /// Size in bytes of an encoded empty varlen value (just the sentinel byte). pub(crate) const VARLEN_EMPTY_SIZE: u32 = 1; -/// Prefix before each byte of a recursively encoded variable-list element. -const LIST_BYTE_ESCAPE: u8 = 0x01; -/// List terminator for ascending fields; sorts before another escaped element byte. +/// Prefix before each recursively encoded variable-list element. +const LIST_ELEMENT_MARKER: u8 = 0x01; +/// List terminator for ascending fields; sorts before another element marker. const LIST_END_ASCENDING: u8 = 0x00; -/// List terminator for descending fields; sorts after another escaped element byte. +/// List terminator for descending fields; sorts after another element marker. const LIST_END_DESCENDING: u8 = 0x02; /// Returns the size in bytes of the encoded form of a non-empty variable-length value. @@ -181,6 +181,31 @@ pub(crate) enum RowWidth { Variable, } +/// A canonical column plus any nested work retained from the sizing pass for encoding. +pub(crate) enum PreparedField { + Canonical(Canonical), + List(PreparedList), +} + +/// Prepared child row keys and list ranges shared by list sizing and encoding. +pub(crate) struct PreparedList { + mask: vortex_mask::Mask, + ranges: Vec<(usize, usize)>, + elements: Canonical, + element_sizes: Vec, + element_offsets: Vec, + total_element_bytes: usize, +} + +impl PreparedField { + pub(crate) fn as_canonical(&self) -> Option<&Canonical> { + match self { + Self::Canonical(canonical) => Some(canonical), + Self::List(_) => None, + } + } +} + /// Classify a column's per-row encoded width by inspecting only its [`DType`]. /// /// Returns `Fixed(w)` when every row encodes to exactly `w` bytes (sentinel + value), @@ -290,6 +315,27 @@ pub(crate) fn field_size( Ok(()) } +/// Size a top-level field and retain list child preparation for the encode pass. +pub(crate) fn prepare_field( + canonical: Canonical, + field: RowSortField, + sizes: &mut [u32], + ctx: &mut ExecutionCtx, +) -> VortexResult { + let prepared = match canonical { + Canonical::List(arr) => prepare_list(&arr, field, ctx)?, + Canonical::Map(arr) => { + prepare_list(&arr.entries().as_::().into_owned(), field, ctx)? + } + canonical => { + field_size(&canonical, field, sizes, ctx)?; + return Ok(PreparedField::Canonical(canonical)); + } + }; + add_size_prepared_list(&prepared, sizes)?; + Ok(PreparedField::List(prepared)) +} + /// Encode a fixed-width column at arithmetic offsets, without reading or writing any per-row /// cursor. /// @@ -380,6 +426,23 @@ pub(crate) fn field_encode( Ok(()) } +/// Encode a field prepared by [`prepare_field`]. +pub(crate) fn field_encode_prepared( + prepared: &PreparedField, + field: RowSortField, + offsets: &[u32], + cursors: &mut [u32], + out: &mut [u8], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + match prepared { + PreparedField::Canonical(canonical) => { + field_encode(canonical, field, offsets, cursors, out, ctx) + } + PreparedField::List(list) => encode_prepared_list(list, field, offsets, cursors, out, ctx), + } +} + fn add_size_const(sizes: &mut [u32], add: u32) { for s in sizes.iter_mut() { *s += add; @@ -441,23 +504,54 @@ fn add_size_list( ctx: &mut ExecutionCtx, ) -> VortexResult<()> { debug_assert_eq!(arr.len(), sizes.len()); + let prepared = prepare_list(arr, field, ctx)?; + add_size_prepared_list(&prepared, sizes) +} + +fn prepare_list( + arr: &ListViewArray, + field: RowSortField, + ctx: &mut ExecutionCtx, +) -> VortexResult { let mask = arr.as_ref().validity()?.execute_mask(arr.len(), ctx)?; let ranges = list_ranges(arr, ctx)?; let elements = arr.elements().clone().execute::(ctx)?; let mut element_sizes = vec![0u32; elements.len()]; field_size(&elements, field, &mut element_sizes, ctx)?; + let mut element_offsets = Vec::with_capacity(elements.len()); + let mut total = 0u32; + for &size in &element_sizes { + element_offsets.push(total); + total = total + .checked_add(size) + .ok_or_else(|| vortex_error::vortex_err!("list element bytes overflow u32"))?; + } + Ok(PreparedList { + mask, + ranges, + elements, + element_sizes, + element_offsets, + total_element_bytes: usize::try_from(total) + .vortex_expect("list element bytes must fit usize"), + }) +} - for (i, (offset, len)) in ranges.into_iter().enumerate() { - let contribution = if !mask.value(i) || len == 0 { +fn add_size_prepared_list(prepared: &PreparedList, sizes: &mut [u32]) -> VortexResult<()> { + for (i, &(offset, len)) in prepared.ranges.iter().enumerate() { + let contribution = if !prepared.mask.value(i) || len == 0 { 1 } else { - let body = element_sizes[offset..offset + len] + let body = prepared.element_sizes[offset..offset + len] .iter() .try_fold(0u32, |sum, &size| sum.checked_add(size)) .ok_or_else(|| vortex_error::vortex_err!("list element sizes overflow u32"))?; - body.checked_mul(2) - .and_then(|size| size.checked_add(2)) - .ok_or_else(|| vortex_error::vortex_err!("list row size overflows u32"))? + body.checked_add( + u32::try_from(len) + .map_err(|_| vortex_error::vortex_err!("list element count overflows u32"))?, + ) + .and_then(|size| size.checked_add(2)) + .ok_or_else(|| vortex_error::vortex_err!("list row size overflows u32"))? }; sizes[i] = sizes[i] .checked_add(contribution) @@ -764,26 +858,24 @@ fn encode_list( out: &mut [u8], ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let mask = arr.as_ref().validity()?.execute_mask(arr.len(), ctx)?; - let ranges = list_ranges(arr, ctx)?; - let elements = arr.elements().clone().execute::(ctx)?; - let mut element_sizes = vec![0u32; elements.len()]; - field_size(&elements, field, &mut element_sizes, ctx)?; + let prepared = prepare_list(arr, field, ctx)?; + encode_prepared_list(&prepared, field, row_offsets, col_offset, out, ctx) +} - let mut element_offsets = Vec::with_capacity(elements.len()); - let mut total = 0u32; - for &size in &element_sizes { - element_offsets.push(total); - total = total - .checked_add(size) - .ok_or_else(|| vortex_error::vortex_err!("list element bytes overflow u32"))?; - } - let mut scratch = vec![0u8; total as usize]; - let mut element_cursors = vec![0u32; elements.len()]; +fn encode_prepared_list( + prepared: &PreparedList, + field: RowSortField, + row_offsets: &[u32], + col_offset: &mut [u32], + out: &mut [u8], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let mut scratch = vec![0u8; prepared.total_element_bytes]; + let mut element_cursors = vec![0u32; prepared.elements.len()]; field_encode( - &elements, + &prepared.elements, field, - &element_offsets, + &prepared.element_offsets, &mut element_cursors, &mut scratch, ctx, @@ -797,9 +889,9 @@ fn encode_list( } else { LIST_END_ASCENDING }; - for (i, (offset, len)) in ranges.into_iter().enumerate() { + for (i, &(offset, len)) in prepared.ranges.iter().enumerate() { let start = (row_offsets[i] + col_offset[i]) as usize; - if !mask.value(i) { + if !prepared.mask.value(i) { out[start] = null; col_offset[i] += 1; continue; @@ -813,13 +905,12 @@ fn encode_list( out[start] = non_empty; let mut dst = start + 1; for element_index in offset..offset + len { - let src = element_offsets[element_index] as usize; - let size = element_sizes[element_index] as usize; - for &byte in &scratch[src..src + size] { - out[dst] = LIST_BYTE_ESCAPE; - out[dst + 1] = byte; - dst += 2; - } + let src = prepared.element_offsets[element_index] as usize; + let size = prepared.element_sizes[element_index] as usize; + out[dst] = LIST_ELEMENT_MARKER; + dst += 1; + out[dst..dst + size].copy_from_slice(&scratch[src..src + size]); + dst += size; } out[dst] = end; let written = diff --git a/vortex-row/src/encode.rs b/vortex-row/src/encode.rs index 5383e3461a0..c88f164c10a 100644 --- a/vortex-row/src/encode.rs +++ b/vortex-row/src/encode.rs @@ -231,6 +231,9 @@ fn execute_row_encode( before_varlen: true, .. } => { + let canonical = canonical + .as_canonical() + .vortex_expect("fixed-width field must retain its canonical array"); codec::field_encode_fixed_arithmetic( canonical, options.fields[i], @@ -243,7 +246,7 @@ fn execute_row_encode( )?; } ColKind::Fixed { .. } | ColKind::Variable { .. } => { - codec::field_encode( + codec::field_encode_prepared( canonical, options.fields[i], listview_offsets_slice, diff --git a/vortex-row/src/size.rs b/vortex-row/src/size.rs index 526f7c7bab5..d475d3e6412 100644 --- a/vortex-row/src/size.rs +++ b/vortex-row/src/size.rs @@ -6,7 +6,6 @@ use std::sync::Arc; use vortex_array::ArrayRef; -use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; @@ -31,6 +30,7 @@ use vortex_error::vortex_bail; use vortex_session::VortexSession; use crate::codec; +use crate::codec::PreparedField; use crate::codec::RowWidth; use crate::options::RowEncodingOptions; use crate::options::deserialize_row_encoding_options; @@ -57,15 +57,14 @@ pub(crate) enum ColKind { /// Result of the size pass: enough information for both [`RowSize::execute`] and the /// downstream [`RowEncode`](super::encode::RowEncode) pipeline. /// -/// `columns` holds the canonicalized form of each input so the encode pass can write bytes -/// without re-decoding — a single canonicalization per column is shared between size and -/// encode. +/// `columns` holds canonicalized inputs plus any nested list preparation retained for the +/// encode pass. pub(crate) struct SizePassResult { pub fixed_per_row: u32, pub var_lengths: Option>, pub col_kinds: Vec, pub first_varlen_idx: Option, - pub columns: Vec, + pub columns: Vec, } /// Walk N input columns once, classifying each as fixed-width or variable-length and @@ -96,7 +95,7 @@ pub(crate) fn compute_sizes( } let nrows = args.row_count(); - let mut columns: Vec = Vec::with_capacity(n_inputs); + let mut columns: Vec = Vec::with_capacity(n_inputs); let mut col_kinds: Vec = Vec::with_capacity(n_inputs); let mut fixed_per_row: u32 = 0; let mut var_lengths: Option> = None; @@ -115,7 +114,7 @@ pub(crate) fn compute_sizes( } let width = codec::row_width_for_dtype(col.dtype())?; // Canonicalize once and reuse for both sizing (variable columns) and encoding. - let canonical = col.execute::(ctx)?; + let canonical = col.execute::(ctx)?; match width { RowWidth::Fixed(w) => { col_kinds.push(ColKind::Fixed { @@ -126,19 +125,19 @@ pub(crate) fn compute_sizes( || vortex_error::vortex_err!("per-row fixed width overflows u32 at column {i}"); fixed_per_row = fixed_per_row.checked_add(w).ok_or_else(overflow)?; running_fixed_prefix = running_fixed_prefix.checked_add(w).ok_or_else(overflow)?; + columns.push(PreparedField::Canonical(canonical)); } RowWidth::Variable => { if first_varlen_idx.is_none() { first_varlen_idx = Some(i); } let v = var_lengths.get_or_insert_with(|| vec![0u32; nrows]); - codec::field_size(&canonical, options.fields[i], v, ctx)?; + columns.push(codec::prepare_field(canonical, options.fields[i], v, ctx)?); col_kinds.push(ColKind::Variable { fixed_prefix: running_fixed_prefix, }); } } - columns.push(canonical); } Ok(SizePassResult { diff --git a/vortex-row/src/tests.rs b/vortex-row/src/tests.rs index 766171d745f..0d856df0ad9 100644 --- a/vortex-row/src/tests.rs +++ b/vortex-row/src/tests.rs @@ -655,6 +655,29 @@ fn variable_list_prefix_order_precedes_following_column() -> VortexResult<()> { Ok(()) } +#[test] +fn variable_list_null_element_prefix_precedes_following_column() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // The child null sentinel and ascending list terminator are both zero. The element marker + // must still make [null] sort before [null, 1] without observing the following column. + let lists = ListViewArray::new( + PrimitiveArray::from_option_iter([None, Some(1i32)]).into_array(), + buffer![0u32, 0].into_array(), + buffer![1u32, 2].into_array(), + Validity::NonNullable, + ) + .into_array(); + let suffix = PrimitiveArray::from_iter([i64::MAX, i64::MIN]).into_array(); + let rows = collect_row_bytes(&convert_columns( + &[lists, suffix], + &[RowSortField::ascending(), RowSortField::ascending()], + &mut ctx, + )?); + + assert!(rows[0] < rows[1]); + Ok(()) +} + #[test] fn map_sort_order() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx();