perf(row): compact variable-list keys

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.

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<u32>,
+    element_offsets: Vec<u32>,
+    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<PreparedField> {
+    let prepared = match canonical {
+        Canonical::List(arr) => prepare_list(&arr, field, ctx)?,
+        Canonical::Map(arr) => {
+            prepare_list(&arr.entries().as_::<ListView>().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<PreparedList> {
     let mask = arr.as_ref().validity()?.execute_mask(arr.len(), ctx)?;
     let ranges = list_ranges(arr, ctx)?;
     let elements = arr.elements().clone().execute::<Canonical>(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::<Canonical>(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 =

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,

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<Vec<u32>>,
     pub col_kinds: Vec<ColKind>,
     pub first_varlen_idx: Option<usize>,
-    pub columns: Vec<Canonical>,
+    pub columns: Vec<PreparedField>,
 }
 
 /// 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<Canonical> = Vec::with_capacity(n_inputs);
+    let mut columns: Vec<PreparedField> = Vec::with_capacity(n_inputs);
     let mut col_kinds: Vec<ColKind> = Vec::with_capacity(n_inputs);
     let mut fixed_per_row: u32 = 0;
     let mut var_lengths: Option<Vec<u32>> = 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::<Canonical>(ctx)?;
+        let canonical = col.execute::<vortex_array::Canonical>(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 {

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();