diff --git a/encodings/sparse/public-api.lock b/encodings/sparse/public-api.lock index e590c9eec40..e9c5ff6e3b2 100644 --- a/encodings/sparse/public-api.lock +++ b/encodings/sparse/public-api.lock @@ -136,6 +136,74 @@ pub fn vortex_sparse::SparseMetadata::clear(&mut self) pub fn vortex_sparse::SparseMetadata::encoded_len(&self) -> usize +pub struct vortex_sparse::SparseSlots + +pub vortex_sparse::SparseSlots::patch_chunk_offsets: core::option::Option + +pub vortex_sparse::SparseSlots::patch_indices: vortex_array::array::erased::ArrayRef + +pub vortex_sparse::SparseSlots::patch_values: vortex_array::array::erased::ArrayRef + +impl vortex_sparse::SparseSlots + +pub const vortex_sparse::SparseSlots::COUNT: usize + +pub const vortex_sparse::SparseSlots::NAMES: [&'static str; 3] + +pub const vortex_sparse::SparseSlots::PATCH_CHUNK_OFFSETS: usize + +pub const vortex_sparse::SparseSlots::PATCH_INDICES: usize + +pub const vortex_sparse::SparseSlots::PATCH_VALUES: usize + +pub fn vortex_sparse::SparseSlots::from_slots(alloc::vec::Vec>) -> Self + +pub fn vortex_sparse::SparseSlots::into_slots(self) -> alloc::vec::Vec> + +pub struct vortex_sparse::SparseSlotsView<'a> + +pub vortex_sparse::SparseSlotsView::patch_chunk_offsets: core::option::Option<&'a vortex_array::array::erased::ArrayRef> + +pub vortex_sparse::SparseSlotsView::patch_indices: &'a vortex_array::array::erased::ArrayRef + +pub vortex_sparse::SparseSlotsView::patch_values: &'a vortex_array::array::erased::ArrayRef + +impl<'a> vortex_sparse::SparseSlotsView<'a> + +pub fn vortex_sparse::SparseSlotsView<'a>::from_slots(&'a [core::option::Option]) -> Self + +pub fn vortex_sparse::SparseSlotsView<'a>::to_owned(&self) -> vortex_sparse::SparseSlots + +impl<'a> core::clone::Clone for vortex_sparse::SparseSlotsView<'a> + +pub fn vortex_sparse::SparseSlotsView<'a>::clone(&self) -> vortex_sparse::SparseSlotsView<'a> + +impl<'a> core::fmt::Debug for vortex_sparse::SparseSlotsView<'a> + +pub fn vortex_sparse::SparseSlotsView<'a>::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result + +impl<'a> core::marker::Copy for vortex_sparse::SparseSlotsView<'a> + +pub trait vortex_sparse::SparseArraySlotsExt: vortex_array::array::typed::TypedArrayRef + +pub fn vortex_sparse::SparseArraySlotsExt::patch_chunk_offsets(&self) -> core::option::Option<&vortex_array::array::erased::ArrayRef> + +pub fn vortex_sparse::SparseArraySlotsExt::patch_indices(&self) -> &vortex_array::array::erased::ArrayRef + +pub fn vortex_sparse::SparseArraySlotsExt::patch_values(&self) -> &vortex_array::array::erased::ArrayRef + +pub fn vortex_sparse::SparseArraySlotsExt::slots_view(&self) -> vortex_sparse::SparseSlotsView<'_> + +impl> vortex_sparse::SparseArraySlotsExt for T + +pub fn T::patch_chunk_offsets(&self) -> core::option::Option<&vortex_array::array::erased::ArrayRef> + +pub fn T::patch_indices(&self) -> &vortex_array::array::erased::ArrayRef + +pub fn T::patch_values(&self) -> &vortex_array::array::erased::ArrayRef + +pub fn T::slots_view(&self) -> vortex_sparse::SparseSlotsView<'_> + pub trait vortex_sparse::SparseExt pub fn vortex_sparse::SparseExt::patches(&self) -> vortex_array::patches::Patches diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 165b09df83c..ede83e2cdd2 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -9,11 +9,16 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; +use vortex_array::arrays::FixedSizeList; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::ListView; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::NullArray; +use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Struct; use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinView; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::listview::ListViewArrayExt; @@ -56,68 +61,101 @@ use vortex_error::vortex_bail; use crate::ConstantArray; use crate::Sparse; -use crate::SparseArray; -use crate::SparseExt as _; -pub(super) fn execute_sparse( - array: &SparseArray, +use crate::SparseParts; + +fn sparse_validity( + patches: &Patches, + fill_value: &Scalar, + nullability: Nullability, + len: usize, ctx: &mut ExecutionCtx, -) -> VortexResult { - if array.patches().num_patches() == 0 { - return Ok(ConstantArray::new(array.fill_scalar().clone(), array.len()).into_array()); +) -> VortexResult { + if nullability == Nullability::NonNullable { + return Ok(Validity::NonNullable); } - Ok(match array.dtype() { + let fill_validity = if fill_value.is_valid() { + Validity::AllValid + } else { + Validity::AllInvalid + }; + let patch_validity = patches.values().validity()?.into_nullable(); + + fill_validity.patch( + len, + patches.offset(), + patches.indices(), + &patch_validity, + ctx, + ) +} + +pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> VortexResult { + let SparseParts { + patches, + fill_value, + dtype, + len, + } = parts; + + if patches.num_patches() == 0 { + return Ok(ConstantArray::new(fill_value, len).into_array()); + } + + // Patches are already resolved (offset subtracted) by SparseParts::resolve_patches(). + Ok(match &dtype { DType::Null => { - assert!(array.fill_scalar().is_null()); - NullArray::new(array.len()).into_array() - } - DType::Bool(..) => { - let resolved_patches = array.resolved_patches()?; - execute_sparse_bools(&resolved_patches, array.fill_scalar(), ctx)? + assert!(fill_value.is_null()); + NullArray::new(len).into_array() } + DType::Bool(..) => execute_sparse_bools(&patches, &fill_value, ctx)?, DType::Primitive(ptype, ..) => { - let resolved_patches = array.resolved_patches()?; match_each_native_ptype!(ptype, |P| { - execute_sparse_primitives::

(&resolved_patches, array.fill_scalar(), ctx)? + execute_sparse_primitives::

(&patches, &fill_value, ctx)? }) } DType::Struct(struct_fields, ..) => execute_sparse_struct( struct_fields, - array.fill_scalar().as_struct(), - array.dtype().nullability(), - &array.patches(), - array.len(), + fill_value.as_struct(), + dtype.nullability(), + &patches, + len, ctx, )?, DType::Decimal(decimal_dtype, nullability) => { let canonical_decimal_value_type = DecimalType::smallest_decimal_value_type(decimal_dtype); - let fill_value = array.fill_scalar().as_decimal(); + let fill_decimal = fill_value.as_decimal(); match_each_decimal_value_type!(canonical_decimal_value_type, |D| { execute_sparse_decimal::( *decimal_dtype, *nullability, - fill_value, - &array.patches(), - array.len(), + fill_decimal, + &patches, + len, ctx, )? }) } dtype @ DType::Utf8(..) => { - let fill_value = array.fill_scalar().as_utf8().value().cloned(); - let fill_value = fill_value.map(BufferString::into_inner); - execute_varbin(array, dtype.clone(), fill_value, ctx)? + let fill = fill_value.as_utf8().value().cloned(); + let fill = fill.map(BufferString::into_inner); + execute_varbin(&patches, &fill_value, dtype.clone(), fill, len, ctx)? } dtype @ DType::Binary(..) => { - let fill_value = array.fill_scalar().as_binary().value().cloned(); - execute_varbin(array, dtype.clone(), fill_value, ctx)? - } - DType::List(values_dtype, nullability) => { - execute_sparse_lists(array, Arc::clone(values_dtype), *nullability, ctx)? + let fill = fill_value.as_binary().value().cloned(); + execute_varbin(&patches, &fill_value, dtype.clone(), fill, len, ctx)? } + DType::List(values_dtype, nullability) => execute_sparse_lists( + &patches, + &fill_value, + Arc::clone(values_dtype), + len, + *nullability, + ctx, + )?, DType::FixedSizeList(.., nullability) => { - execute_sparse_fixed_size_list(array, *nullability, ctx)? + execute_sparse_fixed_size_list(&patches, &fill_value, len, *nullability, ctx)? } DType::Extension(_ext_dtype) => todo!(), DType::Variant(_) => vortex_bail!("Sparse canonicalization does not support Variant"), @@ -129,34 +167,28 @@ pub(super) fn execute_sparse( reason = "complexity is from nested match_smallest_offset_type macro" )] fn execute_sparse_lists( - array: &SparseArray, + resolved: &Patches, + fill_value: &Scalar, values_dtype: Arc, + len: usize, nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { - let resolved_patches = array.resolved_patches()?; - - let indices = resolved_patches - .indices() - .clone() - .execute::(ctx)?; - let values = resolved_patches - .values() - .clone() - .execute::(ctx)?; - let fill_scalar = array.fill_scalar().as_list(); + let indices = resolved.indices().as_::().into_owned(); + let values = resolved.values().as_::().into_owned(); + let fill_list = fill_value.as_list(); - let n_filled = array.len() - resolved_patches.num_patches(); - let total_canonical_values = values.elements().len() + fill_scalar.len() * n_filled; + let n_filled = len - resolved.num_patches(); + let total_canonical_values = values.elements().len() + fill_list.len() * n_filled; Ok(match_each_integer_ptype!(indices.ptype(), |I| { match_smallest_offset_type!(total_canonical_values, |O| { execute_sparse_lists_inner::( indices.as_slice(), values, - fill_scalar, + fill_list, values_dtype, - array.len(), + len, total_canonical_values, nullability, ctx, @@ -224,27 +256,22 @@ fn execute_sparse_lists_inner( /// Canonicalize a sparse [`FixedSizeListArray`] by expanding it into a dense representation. fn execute_sparse_fixed_size_list( - array: &SparseArray, + resolved: &Patches, + fill_value: &Scalar, + len: usize, nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { - let resolved_patches = array.resolved_patches()?; - let indices = resolved_patches - .indices() - .clone() - .execute::(ctx)?; - let values = resolved_patches - .values() - .clone() - .execute::(ctx)?; - let fill_scalar = array.fill_scalar().as_list(); + let indices = resolved.indices().as_::().into_owned(); + let values = resolved.values().as_::().into_owned(); + let fill_scalar = fill_value.as_list(); Ok(match_each_integer_ptype!(indices.ptype(), |I| { execute_sparse_fixed_size_list_inner::( indices.as_slice(), values, fill_scalar, - array.len(), + len, nullability, ctx, ) @@ -431,10 +458,7 @@ fn execute_sparse_struct( Validity::AllInvalid, ), }; - let patch_values_as_struct = unresolved_patches - .values() - .clone() - .execute::(ctx)?; + let patch_values_as_struct = unresolved_patches.values().as_::().into_owned(); let columns_patch_values = patch_values_as_struct.unmasked_fields(); let names = patch_values_as_struct.names(); let validity = top_level_fill_validity.patch( @@ -500,22 +524,16 @@ fn execute_sparse_decimal( } fn execute_varbin( - array: &SparseArray, + resolved: &Patches, + fill_scalar: &Scalar, dtype: DType, fill_value: Option, + len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let patches = array.resolved_patches()?; - let indices = patches.indices().clone().execute::(ctx)?; - let values = patches.values().clone().execute::(ctx)?; - let validity = { - let arr = array.as_array(); - Validity::from_mask( - arr.validity()?.execute_mask(arr.len(), ctx)?, - dtype.nullability(), - ) - }; - let len = array.len(); + let indices = resolved.indices().as_::().into_owned(); + let values = resolved.values().as_::().into_owned(); + let validity = sparse_validity(resolved, fill_scalar, dtype.nullability(), len, ctx)?; Ok(match_each_integer_ptype!(indices.ptype(), |I| { let indices = indices.to_buffer::(); @@ -1219,35 +1237,35 @@ mod test { let elements_slice = elements_array.as_slice::(); // List 0: [1] - let list0_offset = result_listview.offset_at(0); + let list0_offset = result_listview.offset_at(0) as usize; assert_eq!(elements_slice[list0_offset], 1); // List 1: [5,6,7,8] - let list1_offset = result_listview.offset_at(1); - let list1_size = result_listview.size_at(1); + let list1_offset = result_listview.offset_at(1) as usize; + let list1_size = result_listview.size_at(1) as usize; assert_eq!( &elements_slice[list1_offset..list1_offset + list1_size], &[5, 6, 7, 8] ); // List 2: [5,6,7,8] - let list2_offset = result_listview.offset_at(2); - let list2_size = result_listview.size_at(2); + let list2_offset = result_listview.offset_at(2) as usize; + let list2_size = result_listview.size_at(2) as usize; assert_eq!( &elements_slice[list2_offset..list2_offset + list2_size], &[5, 6, 7, 8] ); // List 3: [2] - let list3_offset = result_listview.offset_at(3); + let list3_offset = result_listview.offset_at(3) as usize; assert_eq!(elements_slice[list3_offset], 2); // List 4: [1] - let list4_offset = result_listview.offset_at(4); + let list4_offset = result_listview.offset_at(4) as usize; assert_eq!(elements_slice[list4_offset], 1); // List 5: [2] - let list5_offset = result_listview.offset_at(5); + let list5_offset = result_listview.offset_at(5) as usize; assert_eq!(elements_slice[list5_offset], 2); Ok(()) } diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index cb944178f01..a0a6dc00e09 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -9,6 +9,7 @@ use std::hash::Hasher; use kernel::PARENT_KERNELS; use prost::Message as _; +use vortex_array::AnyCanonical; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -23,6 +24,7 @@ use vortex_array::IntoArray; use vortex_array::Precision; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::buffer::BufferHandle; @@ -33,6 +35,8 @@ use vortex_array::patches::PatchSlotIndices; use vortex_array::patches::Patches; use vortex_array::patches::PatchesData; use vortex_array::patches::PatchesMetadata; +use vortex_array::require_child; +use vortex_array::require_opt_child; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::scalar_fn::fns::operators::Operator; @@ -66,6 +70,47 @@ mod slice; /// A [`Sparse`]-encoded Vortex array. pub type SparseArray = Array; +#[vortex_array::array_slots(Sparse)] +pub struct SparseSlots { + pub patch_indices: ArrayRef, + pub patch_values: ArrayRef, + pub patch_chunk_offsets: Option, +} + +/// Concrete parts of a [`SparseArray`] after iterative execution. +pub(crate) struct SparseParts { + pub patches: Patches, + pub fill_value: Scalar, + pub dtype: DType, + pub len: usize, +} + +pub(crate) trait SparseOwnedExt { + fn into_parts(self) -> VortexResult; +} + +impl SparseOwnedExt for Array { + fn into_parts(self) -> VortexResult { + let patches = Patches::new( + self.len(), + self.patches().offset(), + self.as_ref().slots()[SparseSlots::PATCH_INDICES] + .clone() + .vortex_expect("indices"), + self.as_ref().slots()[SparseSlots::PATCH_VALUES] + .clone() + .vortex_expect("values"), + self.as_ref().slots()[SparseSlots::PATCH_CHUNK_OFFSETS].clone(), + )?; + Ok(SparseParts { + patches, + fill_value: self.fill_scalar().clone(), + dtype: self.dtype().clone(), + len: self.len(), + }) + } +} + #[derive(Clone, prost::Message)] #[repr(C)] pub struct SparseMetadata { @@ -192,7 +237,7 @@ impl VTable for Sparse { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + SparseSlots::NAMES[idx].to_string() } fn reduce_parent( @@ -213,17 +258,60 @@ impl VTable for Sparse { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - execute_sparse(&array, ctx).map(ExecutionResult::done) + // Resolve offset first: wrap indices in Binary(indices, offset, Sub) and + // reassemble with offset=0. Uses slot children (not data) since the executor + // may have updated slots via reduce_parent/execute_parent. + let array = if array.patches().offset() != 0 { + let offset = array.patches().offset(); + let indices = array.patch_indices(); + let values = array.patch_values().clone(); + let len = array.len(); + let offset_scalar = Scalar::from(offset).cast(indices.dtype())?; + let resolved_indices = indices.binary( + ConstantArray::new(offset_scalar, indices.len()).into_array(), + Operator::Sub, + )?; + let patches = Patches::new(len, 0, resolved_indices.clone(), values, None)?; + // Decompose, update in place, and reassemble without re-validation. + match array.try_into_parts() { + Ok(mut parts) => { + parts.data.patches_data = PatchesData::from_patches(&patches); + parts.slots[SparseSlots::PATCH_INDICES] = Some(resolved_indices); + parts.slots[SparseSlots::PATCH_CHUNK_OFFSETS] = None; + unsafe { Array::from_parts_unchecked(parts) } + } + Err(array) => unsafe { + Sparse::new_unchecked(patches, array.fill_scalar().clone()) + }, + } + } else { + array + }; + + // Require children to be executed through the scheduler, + // enabling cross-step optimization via reduce_parent rules. + let array = require_child!( + array, array.patch_indices(), SparseSlots::PATCH_INDICES => Primitive + ); + let array = require_child!( + array, array.patch_values(), SparseSlots::PATCH_VALUES => AnyCanonical + ); + require_opt_child!( + array, + array.patch_chunk_offsets(), + SparseSlots::PATCH_CHUNK_OFFSETS => Primitive + ); + + let parts = array.into_parts()?; + // TODO(joe): remove ctx from execute_sparse since all slots should be canonical. + execute_sparse(parts, ctx).map(ExecutionResult::done) } } -pub(crate) const NUM_SLOTS: usize = 3; -pub(crate) const SLOT_NAMES: [&str; NUM_SLOTS] = - ["patch_indices", "patch_values", "patch_chunk_offsets"]; const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { - indices: 0, - values: 1, - chunk_offsets: 2, + indices: SparseSlots::PATCH_INDICES, + values: SparseSlots::PATCH_VALUES, + chunk_offsets: SparseSlots::PATCH_CHUNK_OFFSETS, }; #[derive(Clone, Debug)] @@ -339,7 +427,7 @@ impl SparseData { } fn make_slots(patches: &Patches) -> Vec> { - let mut slots = Vec::with_capacity(NUM_SLOTS); + let mut slots = Vec::with_capacity(SparseSlots::COUNT); PatchesData::push_slots(&mut slots, Some(patches)); slots }