diff --git a/Cargo.lock b/Cargo.lock index 0010a4e6cec..c57b186ec64 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10517,6 +10517,7 @@ dependencies = [ "arrow-schema 58.4.0", "codspeed-divan-compat", "mimalloc", + "num-traits", "rand 0.10.2", "rstest", "smallvec", diff --git a/docs/specs/row-encoding.md b/docs/specs/row-encoding.md index 8fc3288f82b..fcfd1d6eec9 100644 --- a/docs/specs/row-encoding.md +++ b/docs/specs/row-encoding.md @@ -284,12 +284,14 @@ width is the smallest decimal value type for the decimal precision: | `5..=9` | `i32` | | `10..=18` | `i64` | | `19..=38` | `i128` | +| `39..=76` | `i256` | The storage integer is encoded with the signed integer encoding described above. Decimal columns have one precision and scale, so ordering the scaled integer storage values matches ordering the decimal values in that column. -`Decimal256` is not supported by row encoding. +The signed integer transform is identical at every width, including the 32-byte `i256` +representation used by Decimal256. ## UTF-8 and Binary @@ -395,6 +397,29 @@ For a null fixed-size list, the body is canonicalized: A fixed-size list has fixed row width only when its element type has fixed row width. +## Variable-Size List + +A variable-size list is ordered lexicographically by its elements. Null and empty lists use +the variable-width sentinels. A non-empty list is encoded as: + +```text +varlen_non_empty_sentinel || escaped_elements || 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. + +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. + +## Map + +A map is encoded as its ordered list of non-null `{key, value}` entry structs. Entry order is +significant regardless of the dtype's `keys_sorted` producer assertion. Map comparison is +therefore lexicographic first by entry, then by key and value within each entry. + ## Nested Values Nested structs and fixed-size lists apply the same rules recursively. Each nullable parent @@ -407,17 +432,16 @@ The current row encoder rejects types for which it does not define byte-sort sem | Type | Reason | | --- | --- | -| Variable-size `List` | No row encoding order is defined. | | `Variant` | No row encoding order is defined. | -| `Union` | No row encoding order is defined. | -| `Extension` | No row encoding order is defined. | -| `Decimal256` | Encoding is not implemented. | +| `Union` | Values with different active variants have no defined order. | +| `Extension` | The extension API does not declare whether logical ordering matches storage ordering. | The absence of these encodings is intentional. Adding one requires defining both the logical ordering and the exact byte representation that preserves that ordering. -Temporal extensions could be added later by normalizing them to storage arrays at the -row-encoder boundary, once the supported temporal ordering contract is made explicit. +Extensions can be added once the extension API exposes an explicit contract declaring that +logical ordering is identical to storage ordering. The row encoder must not infer that contract +from the storage dtype alone. ## Size and Output Layout diff --git a/fuzz/src/row.rs b/fuzz/src/row.rs index 22508a49a6f..f8f8f9f6b47 100644 --- a/fuzz/src/row.rs +++ b/fuzz/src/row.rs @@ -28,12 +28,13 @@ use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::map::MapArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::dtype::BigCast; use vortex_array::dtype::DType; -use vortex_array::dtype::DecimalType; use vortex_array::dtype::PType; +use vortex_array::dtype::ToI256; use vortex_array::dtype::half::f16; +use vortex_array::dtype::i256; use vortex_array::match_each_decimal_value_type; use vortex_array::match_each_integer_ptype; use vortex_error::VortexExpect; @@ -201,12 +202,10 @@ fn dtype_row_encodable(dtype: &DType) -> bool { DType::Null | DType::Bool(_) | DType::Primitive(..) | DType::Utf8(_) | DType::Binary(_) => { true } - DType::Decimal(dt, _) => !matches!( - DecimalType::smallest_decimal_value_type(dt), - DecimalType::I256 - ), + DType::Decimal(..) => true, DType::Struct(fields, _) => fields.fields().all(|f| dtype_row_encodable(&f)), - DType::FixedSizeList(elem, ..) => dtype_row_encodable(elem), + DType::List(elem, _) | DType::FixedSizeList(elem, ..) => dtype_row_encodable(elem), + DType::Map(map_dtype, _) => dtype_row_encodable(&map_dtype.entries_dtype()), _ => false, } } @@ -230,8 +229,8 @@ fn collect_row_bytes(array: &ListViewArray, ctx: &mut ExecutionCtx) -> Vec), @@ -350,7 +349,12 @@ fn row_keys(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult match_each_integer_ptype!(a.ptype(), |P| { a.as_slice::

() .iter() - .map(|v| RowKey::Int((*v).into())) + .map(|v| { + RowKey::Int( + v.to_i256() + .vortex_expect("integer ptype must convert to i256"), + ) + }) .collect::>() }), }; @@ -363,9 +367,11 @@ fn row_keys(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult::from(buf[i]).vortex_expect( - "valid decimal values fit i128 (Decimal256 dtypes are filtered)", - )) + RowKey::Int( + buf[i] + .to_i256() + .vortex_expect("decimal value must convert to i256"), + ) } else { RowKey::Null } @@ -403,6 +409,19 @@ fn row_keys(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let mask = a.as_ref().validity()?.execute_mask(len, ctx)?; + let mut keys = Vec::with_capacity(len); + for i in 0..len { + if mask.value(i) { + keys.push(RowKey::Composite(row_keys(&a.list_elements_at(i)?, ctx)?)); + } else { + keys.push(RowKey::Null); + } + } + Ok(keys) + } + Canonical::Map(a) => row_keys(a.entries(), ctx), c => unreachable!( "unsupported dtypes are rejected before oracle construction: {:?}", c.dtype() diff --git a/vortex-row/Cargo.toml b/vortex-row/Cargo.toml index bcb50bc1276..85725604d44 100644 --- a/vortex-row/Cargo.toml +++ b/vortex-row/Cargo.toml @@ -18,6 +18,7 @@ version = { workspace = true } workspace = true [dependencies] +num-traits = { workspace = true } smallvec = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } diff --git a/vortex-row/benches/row_encode.rs b/vortex-row/benches/row_encode.rs index 754d71407bf..dffe5005f0d 100644 --- a/vortex-row/benches/row_encode.rs +++ b/vortex-row/benches/row_encode.rs @@ -3,8 +3,10 @@ #![expect( clippy::unwrap_used, + clippy::expect_used, clippy::clone_on_ref_ptr, clippy::cloned_ref_to_slice_refs, + clippy::chunks_exact_to_as_chunks, clippy::redundant_clone )] @@ -15,9 +17,13 @@ use std::sync::Arc; use std::sync::LazyLock; +use arrow_array::Array; use arrow_array::Int64Array; use arrow_array::StringArray; use arrow_array::StructArray as ArrowStructArray; +use arrow_array::builder::Int64Builder; +use arrow_array::builder::ListBuilder as ArrowListBuilder; +use arrow_array::builder::StringBuilder; use arrow_row::RowConverter; use arrow_row::SortField as ArrowSortField; use arrow_schema::DataType; @@ -34,6 +40,9 @@ use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::listview::ListViewArray; +use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::validity::Validity; use vortex_row::RowEncoder; use vortex_session::VortexSession; @@ -42,11 +51,16 @@ static GLOBAL: MiMalloc = MiMalloc; // Sized so the slowest scenario (struct_mixed) stays within the CodSpeed budget. const N: usize = 1_000; +const LIST_LEN: usize = 8; static SESSION: LazyLock = LazyLock::new(array_session); fn main() { LazyLock::force(&SESSION); + if std::env::args().any(|arg| arg == "--output-sizes") { + print_list_output_sizes(); + return; + } divan::main(); } @@ -70,6 +84,82 @@ fn gen_words(n: usize, mean_len: usize, seed: u64) -> Vec { .collect() } +fn make_vortex_list(elements: vortex_array::ArrayRef) -> vortex_array::ArrayRef { + let offsets = PrimitiveArray::from_iter( + (0..N).map(|i| u32::try_from(i * LIST_LEN).expect("benchmark offset must fit u32")), + ) + .into_array(); + let list_len = u32::try_from(LIST_LEN).expect("benchmark list length must fit u32"); + let sizes = PrimitiveArray::from_iter(std::iter::repeat_n(list_len, N)).into_array(); + ListViewArray::new(elements, offsets, sizes, Validity::NonNullable).into_array() +} + +fn arrow_output_bytes(rows: &arrow_row::Rows) -> u64 { + u64::try_from(rows.lengths().sum::()).expect("encoded output size must fit u64") +} + +fn vortex_output_bytes(rows: &ListViewArray) -> u64 { + u64::try_from(rows.elements().len()).expect("encoded output size must fit u64") +} + +fn list_i64_inputs() -> (arrow_array::ArrayRef, vortex_array::ArrayRef) { + let values = gen_i64(N * LIST_LEN, 11); + let mut builder = ArrowListBuilder::new(Int64Builder::with_capacity(values.len())); + for list in values.chunks_exact(LIST_LEN) { + builder.values().append_slice(list); + builder.append(true); + } + let arrow = Arc::new(builder.finish()) as arrow_array::ArrayRef; + let vortex = make_vortex_list(PrimitiveArray::from_iter(values).into_array()); + (arrow, vortex) +} + +fn list_utf8_inputs() -> (arrow_array::ArrayRef, vortex_array::ArrayRef) { + let values = gen_words(N * LIST_LEN, 16, 13); + let mut builder = ArrowListBuilder::new(StringBuilder::with_capacity( + values.len(), + values.iter().map(String::len).sum(), + )); + for list in values.chunks_exact(LIST_LEN) { + for value in list { + builder.values().append_value(value); + } + builder.append(true); + } + let arrow = Arc::new(builder.finish()) as arrow_array::ArrayRef; + let elements = VarBinViewArray::from_iter_str(values.iter().map(String::as_str)).into_array(); + let vortex = make_vortex_list(elements); + (arrow, vortex) +} + +fn list_output_sizes(arrow: &arrow_array::ArrayRef, vortex: &vortex_array::ArrayRef) -> (u64, u64) { + let converter = RowConverter::new(vec![ArrowSortField::new(arrow.data_type().clone())]) + .expect("benchmark dtype must be supported by arrow-row"); + let arrow_bytes = arrow_output_bytes( + &converter + .convert_columns(&[Arc::clone(arrow)]) + .expect("arrow-row benchmark encode must succeed"), + ); + let mut ctx = SESSION.create_execution_ctx(); + let vortex_bytes = vortex_output_bytes( + &RowEncoder::default() + .encode(&[vortex.clone()], &mut ctx) + .expect("Vortex benchmark encode must succeed"), + ); + (arrow_bytes, vortex_bytes) +} + +fn print_list_output_sizes() { + println!("case,arrow_row_bytes,vortex_bytes"); + for (name, (arrow, vortex)) in [ + ("list_i64", list_i64_inputs()), + ("list_utf8", list_utf8_inputs()), + ] { + let (arrow_bytes, vortex_bytes) = list_output_sizes(&arrow, &vortex); + println!("{name},{arrow_bytes},{vortex_bytes}"); + } +} + // ---------- primitive_i64 ---------- #[divan::bench] @@ -180,3 +270,51 @@ fn struct_mixed_vortex(bencher: divan::Bencher) { .with_inputs(|| SESSION.create_execution_ctx()) .bench_local_values(|mut ctx| encoder.encode(&[struct_arr.clone()], &mut ctx).unwrap()) } + +// ---------- list_i64 ---------- + +#[divan::bench] +fn list_i64_arrow_row(bencher: divan::Bencher) { + let (arr, _) = list_i64_inputs(); + let conv = RowConverter::new(vec![ArrowSortField::new(arr.data_type().clone())]).unwrap(); + let output_bytes = arrow_output_bytes(&conv.convert_columns(&[arr.clone()]).unwrap()); + bencher + .counter(BytesCount::new(output_bytes)) + .bench_local(|| conv.convert_columns(&[arr.clone()]).unwrap()) +} + +#[divan::bench] +fn list_i64_vortex(bencher: divan::Bencher) { + let (_, list) = list_i64_inputs(); + let encoder = RowEncoder::default(); + let mut ctx = SESSION.create_execution_ctx(); + let output_bytes = vortex_output_bytes(&encoder.encode(&[list.clone()], &mut ctx).unwrap()); + bencher + .counter(BytesCount::new(output_bytes)) + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_local_values(|mut ctx| encoder.encode(&[list.clone()], &mut ctx).unwrap()) +} + +// ---------- list_utf8 ---------- + +#[divan::bench] +fn list_utf8_arrow_row(bencher: divan::Bencher) { + let (arr, _) = list_utf8_inputs(); + let conv = RowConverter::new(vec![ArrowSortField::new(arr.data_type().clone())]).unwrap(); + let output_bytes = arrow_output_bytes(&conv.convert_columns(&[arr.clone()]).unwrap()); + bencher + .counter(BytesCount::new(output_bytes)) + .bench_local(|| conv.convert_columns(&[arr.clone()]).unwrap()) +} + +#[divan::bench] +fn list_utf8_vortex(bencher: divan::Bencher) { + let (_, list) = list_utf8_inputs(); + let encoder = RowEncoder::default(); + let mut ctx = SESSION.create_execution_ctx(); + let output_bytes = vortex_output_bytes(&encoder.encode(&[list.clone()], &mut ctx).unwrap()); + bencher + .counter(BytesCount::new(output_bytes)) + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_local_values(|mut ctx| encoder.encode(&[list.clone()], &mut ctx).unwrap()) +} diff --git a/vortex-row/src/codec.rs b/vortex-row/src/codec.rs index 58fbc3ccbbb..2ec0d297c8c 100644 --- a/vortex-row/src/codec.rs +++ b/vortex-row/src/codec.rs @@ -29,6 +29,9 @@ use vortex_array::ExecutionCtx; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::ListView; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::MapArray; use vortex_array::arrays::NullArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -37,12 +40,16 @@ use vortex_array::arrays::decimal::DecimalArrayExt; use vortex_array::arrays::decimal::converted_buffer; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::arrays::map::MapArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::DecimalType; use vortex_array::dtype::NativePType; use vortex_array::dtype::half::f16; +use vortex_array::dtype::i256; +use vortex_array::match_each_integer_ptype; use vortex_array::match_each_native_ptype; use vortex_array::validity::Validity; use vortex_error::VortexExpect; @@ -65,6 +72,13 @@ 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. +const LIST_END_ASCENDING: u8 = 0x00; +/// List terminator for descending fields; sorts after another escaped element byte. +const LIST_END_DESCENDING: u8 = 0x02; + /// Returns the size in bytes of the encoded form of a non-empty variable-length value. /// /// Includes the leading sentinel byte plus `ceil(len/32) * 33` block bytes (32 content + 1 @@ -147,7 +161,9 @@ fn varlen_non_empty_sentinel(field: RowSortField) -> u8 { /// nested struct/FSL when used as a variable-width child) it is the fixed-width null sentinel. fn child_canonical_null_byte(child_dtype: &DType, field: RowSortField) -> u8 { match child_dtype { - DType::Utf8(_) | DType::Binary(_) => varlen_null_sentinel(field), + DType::Utf8(_) | DType::Binary(_) | DType::List(..) | DType::Map(..) => { + varlen_null_sentinel(field) + } _ => field.null_sentinel(), } } @@ -187,14 +203,15 @@ pub(crate) fn row_width_for_dtype(dtype: &DType) -> VortexResult { )))), DType::Decimal(dt, _) => { let vt = decimal_key_type(dt); - if matches!(vt, DecimalType::I256) { - vortex_bail!("row encoding for Decimal256 is not yet implemented"); - } Ok(RowWidth::Fixed(encoded_size_for_fixed(byte_width_u32( vt.byte_width(), )))) } DType::Utf8(_) | DType::Binary(_) => Ok(RowWidth::Variable), + DType::List(elem, _) => { + row_width_for_dtype(elem)?; + Ok(RowWidth::Variable) + } DType::FixedSizeList(elem, n, _) => match row_width_for_dtype(elem)? { // FSL is fixed iff its element type is fixed. Add a sentinel byte for the FSL // itself, then `n` copies of the element width. @@ -224,10 +241,9 @@ pub(crate) fn row_width_for_dtype(dtype: &DType) -> VortexResult { } Ok(RowWidth::Fixed(total)) } - DType::List(..) | DType::Map(..) => { - vortex_bail!( - "row encoding does not support variable-size List or Map arrays (no well-defined ordering)" - ) + DType::Map(map_dtype, _) => { + row_width_for_dtype(&map_dtype.entries_dtype())?; + Ok(RowWidth::Variable) } DType::Variant(_) => { vortex_bail!("row encoding does not support Variant arrays (no well-defined ordering)") @@ -259,10 +275,8 @@ pub(crate) fn field_size( Canonical::VarBinView(arr) => add_size_varbinview(arr, sizes, ctx)?, Canonical::Struct(arr) => add_size_struct(arr, field, sizes, ctx)?, Canonical::FixedSizeList(arr) => add_size_fsl(arr, field, sizes, ctx)?, - Canonical::List(_) => vortex_bail!( - "row encoding does not support canonical List arrays: {:?}", - canonical.dtype() - ), + Canonical::List(arr) => add_size_list(arr, field, sizes, ctx)?, + Canonical::Map(arr) => add_size_map(arr, field, sizes, ctx)?, Canonical::Variant(_) => { vortex_bail!("row encoding does not support Variant arrays (no well-defined ordering)") } @@ -351,10 +365,8 @@ pub(crate) fn field_encode( Canonical::VarBinView(arr) => encode_varbinview(arr, field, offsets, cursors, out, ctx)?, Canonical::Struct(arr) => encode_struct(arr, field, offsets, cursors, out, ctx)?, Canonical::FixedSizeList(arr) => encode_fsl(arr, field, offsets, cursors, out, ctx)?, - Canonical::List(_) => vortex_bail!( - "row encoding does not support canonical List arrays: {:?}", - canonical.dtype() - ), + Canonical::List(arr) => encode_list(arr, field, offsets, cursors, out, ctx)?, + Canonical::Map(arr) => encode_map(arr, field, offsets, cursors, out, ctx)?, Canonical::Variant(_) => { vortex_bail!("row encoding does not support Variant arrays (no well-defined ordering)") } @@ -392,6 +404,82 @@ fn add_size_decimal(arr: &DecimalArray, sizes: &mut [u32]) { add_size_const(sizes, encoded_size_for_fixed(width)); } +fn list_ranges(arr: &ListViewArray, ctx: &mut ExecutionCtx) -> VortexResult> { + use num_traits::ToPrimitive; + + let offsets = arr.offsets().clone().execute::(ctx)?; + let sizes = arr.sizes().clone().execute::(ctx)?; + let offsets = match_each_integer_ptype!(offsets.ptype(), |O| { + offsets + .as_slice::() + .iter() + .map(|value| { + value + .to_usize() + .vortex_expect("validated list offset must fit usize") + }) + .collect::>() + }); + let sizes = match_each_integer_ptype!(sizes.ptype(), |S| { + sizes + .as_slice::() + .iter() + .map(|value| { + value + .to_usize() + .vortex_expect("validated list size must fit usize") + }) + .collect::>() + }); + Ok(offsets.into_iter().zip(sizes).collect()) +} + +fn add_size_list( + arr: &ListViewArray, + field: RowSortField, + sizes: &mut [u32], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + debug_assert_eq!(arr.len(), sizes.len()); + 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)?; + + for (i, (offset, len)) in ranges.into_iter().enumerate() { + let contribution = if !mask.value(i) || len == 0 { + 1 + } else { + let body = 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"))? + }; + sizes[i] = sizes[i] + .checked_add(contribution) + .ok_or_else(|| vortex_error::vortex_err!("per-row size overflow"))?; + } + Ok(()) +} + +fn add_size_map( + arr: &MapArray, + field: RowSortField, + sizes: &mut [u32], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + add_size_list( + &arr.entries().as_::().into_owned(), + field, + sizes, + ctx, + ) +} + /// The decimal type every chunk of a decimal column encodes its keys at. /// /// Derived from the declared decimal dtype rather than the chunk's physical `values_type`, @@ -663,9 +751,102 @@ fn encode_decimal( encode_decimal_typed::(arr, &mask, field, row_offsets, col_offset, out) } DecimalType::I256 => { - vortex_bail!("row encoding for Decimal256 is not yet implemented") + encode_decimal_typed::(arr, &mask, field, row_offsets, col_offset, out) + } + } +} + +fn encode_list( + arr: &ListViewArray, + field: RowSortField, + row_offsets: &[u32], + col_offset: &mut [u32], + 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 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()]; + field_encode( + &elements, + field, + &element_offsets, + &mut element_cursors, + &mut scratch, + ctx, + )?; + + let null = varlen_null_sentinel(field); + let empty = varlen_empty_sentinel(field); + let non_empty = varlen_non_empty_sentinel(field); + let end = if field.descending { + LIST_END_DESCENDING + } else { + LIST_END_ASCENDING + }; + for (i, (offset, len)) in ranges.into_iter().enumerate() { + let start = (row_offsets[i] + col_offset[i]) as usize; + if !mask.value(i) { + out[start] = null; + col_offset[i] += 1; + continue; + } + if len == 0 { + out[start] = empty; + col_offset[i] += 1; + continue; } + + 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; + } + } + out[dst] = end; + let written = + u32::try_from(dst + 1 - start).vortex_expect("validated list row size must fit u32"); + col_offset[i] = col_offset[i] + .checked_add(written) + .vortex_expect("list row offset overflow"); } + Ok(()) +} + +fn encode_map( + arr: &MapArray, + field: RowSortField, + row_offsets: &[u32], + col_offset: &mut [u32], + out: &mut [u8], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + encode_list( + &arr.entries().as_::().into_owned(), + field, + row_offsets, + col_offset, + out, + ctx, + ) } fn encode_decimal_typed( @@ -1219,6 +1400,7 @@ impl_row_encode_signed!(i16); impl_row_encode_signed!(i32); impl_row_encode_signed!(i64); impl_row_encode_signed!(i128); +impl_row_encode_signed!(i256); impl RowEncode for f32 { fn encode_to(self, out: &mut [u8], descending: bool) { diff --git a/vortex-row/src/encoder.rs b/vortex-row/src/encoder.rs index c1d7a8f2a54..b725cbd34c8 100644 --- a/vortex-row/src/encoder.rs +++ b/vortex-row/src/encoder.rs @@ -139,6 +139,10 @@ fn reject_extension_dtype(dtype: &DType) -> VortexResult<()> { DType::FixedSizeList(elem, ..) | DType::List(elem, _) => { reject_extension_dtype(elem)?; } + DType::Map(map_dtype, _) => { + reject_extension_dtype(&map_dtype.key_dtype())?; + reject_extension_dtype(&map_dtype.value_dtype())?; + } _ => {} } Ok(()) diff --git a/vortex-row/src/lib.rs b/vortex-row/src/lib.rs index f33fa8a4fcb..d266e2c6042 100644 --- a/vortex-row/src/lib.rs +++ b/vortex-row/src/lib.rs @@ -22,10 +22,10 @@ //! those sizes to allocate one contiguous elements buffer, then writes each column's bytes //! into the per-row slots from left to right. //! -//! Supported logical types are nulls, booleans, primitive integers and floats, decimals up to -//! 128 bits, UTF-8 and binary values, structs, and fixed-size lists. Extension, variant, -//! union, and variable-size list arrays are rejected because this crate does not define an -//! ordering for them. +//! Supported logical types are nulls, booleans, primitive integers and floats, decimals, +//! UTF-8 and binary values, structs, fixed-size and variable-size lists, and maps. Extension, +//! variant, and union arrays are rejected because this crate does not define a total ordering +//! for them. //! //! The byte-level format is documented in the row encoding spec: //! . diff --git a/vortex-row/src/tests.rs b/vortex-row/src/tests.rs index c6ba88dc498..766171d745f 100644 --- a/vortex-row/src/tests.rs +++ b/vortex-row/src/tests.rs @@ -13,13 +13,18 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::MapArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::MapDType; use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::i256; use vortex_array::extension::datetime::Date; use vortex_array::extension::datetime::TimeUnit; use vortex_array::validity::Validity; @@ -598,25 +603,87 @@ fn primitive_f16_sort_order() -> VortexResult<()> { Ok(()) } +#[rstest] +#[case::ascending(RowSortField::ascending(), vec![4, 2, 1, 0, 5, 3])] +#[case::descending_nulls_last( + RowSortField::descending().nulls_last(), + vec![3, 5, 0, 1, 2, 4] +)] +fn variable_list_sort_order( + #[case] field: RowSortField, + #[case] expected_indices: Vec, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // [[1, 2], [1], [], [2], null, [1, 3]] + let elements = PrimitiveArray::from_iter([1i32, 2, 2, 1, 3]).into_array(); + let offsets = buffer![0u32, 0, 0, 2, 0, 3].into_array(); + let sizes = buffer![2u32, 1, 0, 1, 2, 2].into_array(); + let list = ListViewArray::new( + elements, + offsets, + sizes, + Validity::from_iter([true, true, true, true, false, true]), + ); + + let rows = collect_row_bytes(&convert_columns(&[list.into_array()], &[field], &mut ctx)?); + let mut actual_indices: Vec = (0..rows.len()).collect(); + actual_indices.sort_by(|&a, &b| rows[a].cmp(&rows[b])); + assert_eq!(actual_indices, expected_indices); + Ok(()) +} + #[test] -fn reject_list_dtype_early() { - use vortex_array::ArrayRef; - use vortex_array::arrays::ListArray; - use vortex_array::validity::Validity; - use vortex_buffer::buffer; +fn variable_list_prefix_order_precedes_following_column() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // The second column deliberately has the opposite order. The list terminator must decide + // [1] < [1, 2] before comparison can observe that following column. + let lists = ListViewArray::new( + PrimitiveArray::from_iter([1i32, 2]).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(); - let offsets = PrimitiveArray::new(buffer![0u32, 1, 2], Validity::NonNullable).into_array(); - let elements = PrimitiveArray::from_iter([10i32, 20]).into_array(); - let list: ArrayRef = ListArray::try_new(elements, offsets, Validity::NonNullable) - .unwrap() - .into_array(); - let err = convert_columns(&[list], &[RowSortField::default()], &mut ctx) - .expect_err("List should not be accepted"); - assert!( - err.to_string().contains("List"), - "expected error mentioning List, got: {err}" + let map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Primitive(PType::I32, Nullability::NonNullable), + false, + )?; + let keys = PrimitiveArray::from_iter([1i32, 2, 2]).into_array(); + let values = PrimitiveArray::from_iter([10i32, 20, 0]).into_array(); + let entries = StructArray::from_fields(&[("key", keys), ("value", values)])?.into_array(); + // [{1: 10, 2: 20}, {1: 10}, {}, {2: 0}, null] + let entry_lists = ListViewArray::new( + entries, + buffer![0u32, 0, 0, 2, 0].into_array(), + buffer![2u32, 1, 0, 1, 2].into_array(), + Validity::from_iter([true, true, true, true, false]), ); + let maps = MapArray::new(map_dtype, entry_lists).into_array(); + + let rows = collect_row_bytes(&convert_columns( + &[maps], + &[RowSortField::ascending()], + &mut ctx, + )?); + let mut actual_indices: Vec = (0..rows.len()).collect(); + actual_indices.sort_by(|&a, &b| rows[a].cmp(&rows[b])); + assert_eq!(actual_indices, vec![4, 2, 1, 0, 3]); + Ok(()) } /// Chunks of one decimal column can compress to different physical value widths. The key @@ -692,3 +759,32 @@ fn decimal_value_not_fitting_key_width_errors() { "expected a does-not-fit error, got: {err}" ); } + +#[rstest] +#[case::ascending(false)] +#[case::descending(true)] +fn decimal256_sort_order(#[case] descending: bool) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let large = i256::from_parts(0, 1); + let values = vec![ + large, + -large, + i256::from_i128(-1), + i256::ZERO, + i256::from_i128(1), + ]; + let decimal = DecimalArray::from_iter(values.clone(), DecimalDType::new(76, 0)).into_array(); + let field = RowSortField::new(descending, true); + let rows = collect_row_bytes(&convert_columns(&[decimal], &[field], &mut ctx)?); + + assert!(rows.iter().all(|row| row.len() == 33)); + let mut actual_indices: Vec = (0..rows.len()).collect(); + actual_indices.sort_by(|&a, &b| rows[a].cmp(&rows[b])); + let mut expected_indices: Vec = (0..values.len()).collect(); + expected_indices.sort_by(|&a, &b| { + let order = values[a].cmp(&values[b]); + if descending { order.reverse() } else { order } + }); + assert_eq!(actual_indices, expected_indices); + Ok(()) +}