From 40783a62bcfdc5b5d92782144c101cd577801d11 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 17 May 2026 22:11:06 +0000 Subject: [PATCH] vortex-row: RowEncode ScalarFn Add the RowEncode variadic scalar function: encode N input columns into a single ListView in a five-phase pipeline. Phase 1: size pass via `compute_sizes`. Phase 2: allocate a zero-initialized output buffer sized to fit every row's encoded bytes; bail if the total exceeds u32::MAX. Phase 3: build per-row `listview_offsets`: i * fixed_per_row for the pure-fixed case, or i * fixed_per_row + exclusive cumsum of varlen lengths otherwise. Uses the simple `Vec::push` + `checked_add` loop. Phase 4: walk columns left-to-right and call `dispatch_encode` for every column (cursor path for all). Each call writes its per-row bytes at `offsets[i] + cursors[i]` and advances the cursor. Phase 5: build the ListView via the validating `try_new` constructor. `dispatch_encode` is the canonicalize-then-`codec::field_encode` fallback; in-crate kernel arms and the inventory registry land in PR 3. The `RowEncodeKernel` trait is defined but unused. PR 2 will iterate on this pipeline (skip zero-init, skip ListView validation, auto- vectorize the offsets loop, etc.). Signed-off-by: Claude --- vortex-row/public-api.lock | 76 ++++++++++++ vortex-row/src/encode.rs | 238 +++++++++++++++++++++++++++++++++++++ vortex-row/src/lib.rs | 3 + vortex-row/src/size.rs | 8 ++ 4 files changed, 325 insertions(+) create mode 100644 vortex-row/src/encode.rs diff --git a/vortex-row/public-api.lock b/vortex-row/public-api.lock index 85985bf7521..f999303948d 100644 --- a/vortex-row/public-api.lock +++ b/vortex-row/public-api.lock @@ -102,6 +102,46 @@ pub fn vortex_row::codec::field_size(&vortex_array::canonical::Canonical, vortex pub fn vortex_row::codec::row_width_for_dtype(&vortex_array::dtype::DType) -> vortex_error::VortexResult +pub mod vortex_row::encode + +pub struct vortex_row::encode::RowEncode + +impl core::clone::Clone for vortex_row::encode::RowEncode + +pub fn vortex_row::encode::RowEncode::clone(&self) -> vortex_row::encode::RowEncode + +impl core::fmt::Debug for vortex_row::encode::RowEncode + +pub fn vortex_row::encode::RowEncode::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result + +impl vortex_array::scalar_fn::vtable::ScalarFnVTable for vortex_row::encode::RowEncode + +pub type vortex_row::encode::RowEncode::Options = vortex_row::options::RowEncodeOptions + +pub fn vortex_row::encode::RowEncode::arity(&self, &Self::Options) -> vortex_array::scalar_fn::vtable::Arity + +pub fn vortex_row::encode::RowEncode::child_name(&self, &Self::Options, usize) -> vortex_array::scalar_fn::vtable::ChildName + +pub fn vortex_row::encode::RowEncode::deserialize(&self, &[u8], &vortex_session::VortexSession) -> vortex_error::VortexResult + +pub fn vortex_row::encode::RowEncode::execute(&self, &Self::Options, &dyn vortex_array::scalar_fn::vtable::ExecutionArgs, &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult + +pub fn vortex_row::encode::RowEncode::id(&self) -> vortex_array::scalar_fn::ScalarFnId + +pub fn vortex_row::encode::RowEncode::is_fallible(&self, &Self::Options) -> bool + +pub fn vortex_row::encode::RowEncode::is_null_sensitive(&self, &Self::Options) -> bool + +pub fn vortex_row::encode::RowEncode::return_dtype(&self, &Self::Options, &[vortex_array::dtype::DType]) -> vortex_error::VortexResult + +pub fn vortex_row::encode::RowEncode::serialize(&self, &Self::Options) -> vortex_error::VortexResult>> + +pub trait vortex_row::encode::RowEncodeKernel: vortex_array::array::vtable::VTable + +pub fn vortex_row::encode::RowEncodeKernel::row_encode_into(vortex_array::array::view::ArrayView<'_, Self>, vortex_row::options::SortField, &[u32], &mut [u32], &mut [u8], &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult> + +pub fn vortex_row::encode::dispatch_encode(&vortex_array::array::erased::ArrayRef, vortex_row::options::SortField, &[u32], &mut [u32], &mut [u8], &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<()> + pub mod vortex_row::options pub struct vortex_row::options::RowEncodeOptions @@ -222,6 +262,38 @@ pub fn vortex_row::size::RowSizeKernel::row_size_contribution(vortex_array::arra pub fn vortex_row::size::dispatch_size(&vortex_array::array::erased::ArrayRef, vortex_row::options::SortField, &mut [u32], &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<()> +pub struct vortex_row::RowEncode + +impl core::clone::Clone for vortex_row::encode::RowEncode + +pub fn vortex_row::encode::RowEncode::clone(&self) -> vortex_row::encode::RowEncode + +impl core::fmt::Debug for vortex_row::encode::RowEncode + +pub fn vortex_row::encode::RowEncode::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result + +impl vortex_array::scalar_fn::vtable::ScalarFnVTable for vortex_row::encode::RowEncode + +pub type vortex_row::encode::RowEncode::Options = vortex_row::options::RowEncodeOptions + +pub fn vortex_row::encode::RowEncode::arity(&self, &Self::Options) -> vortex_array::scalar_fn::vtable::Arity + +pub fn vortex_row::encode::RowEncode::child_name(&self, &Self::Options, usize) -> vortex_array::scalar_fn::vtable::ChildName + +pub fn vortex_row::encode::RowEncode::deserialize(&self, &[u8], &vortex_session::VortexSession) -> vortex_error::VortexResult + +pub fn vortex_row::encode::RowEncode::execute(&self, &Self::Options, &dyn vortex_array::scalar_fn::vtable::ExecutionArgs, &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult + +pub fn vortex_row::encode::RowEncode::id(&self) -> vortex_array::scalar_fn::ScalarFnId + +pub fn vortex_row::encode::RowEncode::is_fallible(&self, &Self::Options) -> bool + +pub fn vortex_row::encode::RowEncode::is_null_sensitive(&self, &Self::Options) -> bool + +pub fn vortex_row::encode::RowEncode::return_dtype(&self, &Self::Options, &[vortex_array::dtype::DType]) -> vortex_error::VortexResult + +pub fn vortex_row::encode::RowEncode::serialize(&self, &Self::Options) -> vortex_error::VortexResult>> + pub struct vortex_row::RowEncodeOptions pub vortex_row::RowEncodeOptions::fields: smallvec::SmallVec<[vortex_row::options::SortField; 4]> @@ -330,6 +402,10 @@ impl core::marker::Copy for vortex_row::options::SortField impl core::marker::StructuralPartialEq for vortex_row::options::SortField +pub trait vortex_row::RowEncodeKernel: vortex_array::array::vtable::VTable + +pub fn vortex_row::RowEncodeKernel::row_encode_into(vortex_array::array::view::ArrayView<'_, Self>, vortex_row::options::SortField, &[u32], &mut [u32], &mut [u8], &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult> + pub trait vortex_row::RowSizeKernel: vortex_array::array::vtable::VTable pub fn vortex_row::RowSizeKernel::row_size_contribution(vortex_array::array::view::ArrayView<'_, Self>, vortex_row::options::SortField, &mut [u32], &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult> diff --git a/vortex-row/src/encode.rs b/vortex-row/src/encode.rs new file mode 100644 index 00000000000..1b77d955964 --- /dev/null +++ b/vortex-row/src/encode.rs @@ -0,0 +1,238 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![allow( + clippy::cast_possible_truncation, + reason = "row encoding indexes into u32-sized buffers; lengths are validated to fit in u32" +)] + +//! `RowEncode` variadic scalar function: encode N input columns into a single `ListView`. +//! +//! The output's `(elements, offsets, sizes)` triple is built up in a single left-to-right +//! pass over the input columns. The `sizes` array doubles as the per-row write cursor, so +//! when the last column finishes encoding, the accumulator is the final array - no separate +//! conversion step is needed. + +use std::sync::Arc; + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; + +use crate::codec; +use crate::options::RowEncodeOptions; +use crate::options::SortField; +use crate::options::deserialize_row_encode_options; +use crate::options::serialize_row_encode_options; +use crate::size::compute_sizes; + +/// Variadic scalar function that encodes N input columns into a single `List` +/// [`ListViewArray`] where row `i` contains the row-encoded bytes for column values +/// `cols[0][i], cols[1][i], ...` concatenated left-to-right. +#[derive(Clone, Debug)] +pub struct RowEncode; + +impl ScalarFnVTable for RowEncode { + type Options = RowEncodeOptions; + + fn id(&self) -> ScalarFnId { + ScalarFnId::from("vortex.row_encode") + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(serialize_row_encode_options(options))) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + deserialize_row_encode_options(metadata) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Variadic { min: 1, max: None } + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + ChildName::from(Arc::from(format!("col_{}", child_idx))) + } + + fn return_dtype(&self, _options: &Self::Options, _args: &[DType]) -> VortexResult { + Ok(DType::List( + Arc::new(DType::Primitive(PType::U8, Nullability::NonNullable)), + Nullability::NonNullable, + )) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + execute_row_encode(options, args, ctx) + } + + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +fn execute_row_encode( + options: &RowEncodeOptions, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let nrows = args.row_count(); + + // ===== Phase 1: classify + size pass ===== + let crate::size::SizePassResult { + fixed_per_row, + var_lengths, + col_kinds: _, + first_varlen_idx: _, + columns, + } = compute_sizes(options, args, ctx, "RowEncode")?; + + // ===== Phase 2: totals + buffer ===== + let var_total: u64 = var_lengths + .as_ref() + .map_or(0, |v| v.iter().map(|&x| u64::from(x)).sum()); + let total: u64 = (nrows as u64) + .checked_mul(u64::from(fixed_per_row)) + .and_then(|t| t.checked_add(var_total)) + .vortex_expect("row-encoded total bytes overflow"); + if total > u32::MAX as u64 { + vortex_bail!("row-encoded output size {} bytes exceeds u32::MAX", total); + } + let total_len = total as usize; + + // Allocate the elements buffer (zero-initialized). The zero-init lets every encoder + // assume previously-untouched bytes are zero, simplifying the null-row fill paths. + // PR 2 skips this memset because every byte in the output range is written by some + // encoder. + let mut out_buf: BufferMut = BufferMut::with_capacity(total_len); + out_buf.push_n(0u8, total_len); + + // ===== Phase 3: per-row offsets ===== + // listview_offsets[i] is the absolute byte offset where row `i` begins. + // For pure-fixed: i * fixed_per_row. + // For mixed: i * fixed_per_row + exclusive prefix sum of var_lengths. + let mut listview_offsets: Vec = Vec::with_capacity(nrows); + match var_lengths.as_ref() { + None => { + for i in 0..nrows { + listview_offsets.push( + (i as u32) + .checked_mul(fixed_per_row) + .vortex_expect("row offset overflow (already validated total fits in u32)"), + ); + } + } + Some(v) => { + let mut acc: u32 = 0; + for (i, &l) in v.iter().enumerate() { + let off = (i as u32) + .checked_mul(fixed_per_row) + .and_then(|t| t.checked_add(acc)) + .vortex_expect("row offset overflow"); + listview_offsets.push(off); + acc = acc.checked_add(l).vortex_expect("varlen prefix overflow"); + } + } + } + + // Per-row write cursor (also doubles as the ListView `sizes` slot when done). + let mut row_cursors = vec![0u32; nrows]; + + // ===== Phase 4: encode columns via the cursor path ===== + for (i, col) in columns.iter().enumerate() { + dispatch_encode( + col, + options.fields[i], + &listview_offsets, + &mut row_cursors, + &mut out_buf, + ctx, + )?; + } + + // ===== Phase 5: build ListView output ===== + let elements = PrimitiveArray::new(out_buf.freeze(), Validity::NonNullable).into_array(); + let offsets_arr = PrimitiveArray::new( + Buffer::::copy_from(&listview_offsets), + Validity::NonNullable, + ) + .into_array(); + let sizes_arr = PrimitiveArray::new( + Buffer::::copy_from(&row_cursors), + Validity::NonNullable, + ) + .into_array(); + Ok( + ListViewArray::try_new(elements, offsets_arr, sizes_arr, Validity::NonNullable)? + .into_array(), + ) +} + +/// Dispatch a single column's encoding into the shared `out` buffer. +/// +/// For PR 1 this is just the canonicalize-then-`codec::field_encode` fallback path. +/// In-crate fast paths for `Constant`/`Dict`/`Patched` and the inventory-based registry +/// for downstream encodings are added in PR 3. +pub fn dispatch_encode( + col: &ArrayRef, + field: SortField, + offsets: &[u32], + cursors: &mut [u32], + out: &mut [u8], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let canonical = col.clone().execute::(ctx)?; + codec::field_encode(&canonical, field, offsets, cursors, out, ctx) +} + +/// Mutate-buffer kernel: write this column's per-row bytes into `out` at +/// `offsets[i] + cursors[i]`, advancing `cursors[i]` by the bytes written. +/// +/// Return `Ok(None)` to decline and fall back to the canonical path. +/// +/// Trait is defined now; per-encoding impls and dispatch wiring land in PR 3. +pub trait RowEncodeKernel: VTable { + /// Write this column's per-row bytes into `out` at `offsets[i] + cursors[i]`, advancing + /// `cursors[i]` by the bytes written. + fn row_encode_into( + column: ArrayView<'_, Self>, + field: SortField, + offsets: &[u32], + cursors: &mut [u32], + out: &mut [u8], + ctx: &mut ExecutionCtx, + ) -> VortexResult>; +} diff --git a/vortex-row/src/lib.rs b/vortex-row/src/lib.rs index 6f1d8fbeab3..ef0209f3d9c 100644 --- a/vortex-row/src/lib.rs +++ b/vortex-row/src/lib.rs @@ -7,9 +7,12 @@ //! This commit only establishes the crate skeleton and an `initialize` stub. pub mod codec; +pub mod encode; pub mod options; pub mod size; +pub use encode::RowEncode; +pub use encode::RowEncodeKernel; pub use options::RowEncodeOptions; pub use options::SortField; pub use size::RowSize; diff --git a/vortex-row/src/size.rs b/vortex-row/src/size.rs index fbde52e1863..7148a2a21d8 100644 --- a/vortex-row/src/size.rs +++ b/vortex-row/src/size.rs @@ -72,7 +72,15 @@ pub(crate) enum ColKind { pub(crate) struct SizePassResult { pub fixed_per_row: u32, pub var_lengths: Option>, + #[allow( + dead_code, + reason = "consumed by the arithmetic-write fast path added in PR 2" + )] pub col_kinds: Vec, + #[allow( + dead_code, + reason = "consumed by the arithmetic-write fast path added in PR 2" + )] pub first_varlen_idx: Option, pub columns: Vec, }