Skip to content

Canonicalize chunked nested types through the builder [builders-child-stack] #8967

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
robert3005 wants to merge 1 commit into from
+54 −231
Open

Canonicalize chunked nested types through the builder [builders-child-stack] #8967

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
285 changes: 54 additions & 231 deletions vortex-array/src/arrays/chunked/vtable/canonical.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,6 @@
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use itertools::Itertools as _;
use vortex_buffer::Buffer;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_ensure;
use vortex_error::vortex_err;
Expand All @@ -15,23 +13,11 @@ use crate::IntoArray;
use crate::array::ArrayView;
use crate::arrays::Chunked;
use crate::arrays::ChunkedArray;
use crate::arrays::FixedSizeListArray;
use crate::arrays::ListViewArray;
use crate::arrays::PrimitiveArray;
use crate::arrays::StructArray;
use crate::arrays::VariantArray;
use crate::arrays::chunked::ChunkedArrayExt;
use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt;
use crate::arrays::listview::ListViewArraySlotsExt;
use crate::arrays::listview::ListViewRebuildMode;
use crate::arrays::variant::VariantArraySlotsExt;
use crate::builders::builder_with_capacity_in;
use crate::builtins::ArrayBuiltins;
use crate::dtype::DType;
use crate::dtype::Nullability;
use crate::dtype::PType;
use crate::memory::HostAllocatorExt;
use crate::validity::Validity;

pub(super) fn _canonicalize(
array: ArrayView<'_, Chunked>,
Expand All @@ -48,28 +34,11 @@ pub(super) fn _canonicalize(
return array.chunk(0).clone().execute::<Canonical>(ctx);
}

let owned_chunks: Vec<ArrayRef> = array.iter_chunks().cloned().collect();
Ok(match array.dtype() {
DType::Struct(..) => {
let struct_array = pack_struct_chunks(owned_chunks, ctx)?;
Canonical::Struct(struct_array)
DType::Variant(_) => {
let owned_chunks: Vec<ArrayRef> = array.iter_chunks().cloned().collect();
Canonical::Variant(pack_variant_chunks(owned_chunks, ctx)?)
}
DType::List(elem_dtype, _) => Canonical::List(swizzle_list_chunks(
&owned_chunks,
array.array().validity()?,
elem_dtype,
ctx,
)?),
DType::FixedSizeList(elem_dtype, list_size, _) => {
Canonical::FixedSizeList(swizzle_fixed_size_list_chunks(
&owned_chunks,
array.array().validity()?,
elem_dtype,
*list_size,
ctx,
)?)
}
DType::Variant(_) => Canonical::Variant(pack_variant_chunks(owned_chunks, ctx)?),
_ => {
let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len());
array.array().append_to_builder(builder.as_mut(), ctx)?;
Expand All @@ -78,17 +47,6 @@ pub(super) fn _canonicalize(
})
}

/// Packs many [`StructArray`]s to instead be a single [`StructArray`], where the [`DynArrayData`](crate::array::DynArrayData) for each
/// field is a [`ChunkedArray`].
///
/// The caller guarantees there are at least 2 chunks.
fn pack_struct_chunks(chunks: Vec<ArrayRef>, ctx: &mut ExecutionCtx) -> VortexResult<StructArray> {
chunks
.into_iter()
.map(|c| c.execute::<StructArray>(ctx))
.process_results(|iter| StructArray::try_concat(iter))?
}

/// Packs many [`VariantArray`]s into one [`VariantArray`] with chunked children.
///
/// The caller guarantees there are at least 2 chunks.
Expand Down Expand Up @@ -147,151 +105,14 @@ fn pack_variant_chunks(
VariantArray::try_new(core_storage, shredded)
}

/// Packs [`ListViewArray`]s together into a chunked `ListViewArray`.
///
/// We use the existing arrays (chunks) to form a chunked array of `elements` (the child array).
///
/// The caller guarantees there are at least 2 chunks.
fn swizzle_list_chunks(
chunks: &[ArrayRef],
validity: Validity,
elem_dtype: &DType,
ctx: &mut ExecutionCtx,
) -> VortexResult<ListViewArray> {
let len: usize = chunks.iter().map(|c| c.len()).sum();

assert_eq!(
chunks[0]
.dtype()
.as_list_element_opt()
.vortex_expect("DType was somehow not a list")
.as_ref(),
elem_dtype
);

// Since each list array in `chunks` has offsets local to each array, we can reuse the existing
// array's child `elements` as the chunks and recompute offsets.

let mut list_elements_chunks = Vec::with_capacity(chunks.len());
let mut num_elements = 0;

// TODO(connor)[ListView]: We could potentially choose a smaller type here, but that would make
// this much more complicated.
// We (somewhat arbitrarily) choose `u64` for our offsets and sizes here. These can always be
// narrowed later by the compressor.
let allocator = ctx.allocator();
let mut offsets = allocator.allocate_typed::<u64>(len)?;
let mut sizes = allocator.allocate_typed::<u64>(len)?;
let offsets_out: &mut [u64] = offsets.as_mut_slice_typed::<u64>()?;
let sizes_slice_out: &mut [u64] = sizes.as_mut_slice_typed::<u64>()?;
let mut next_list = 0usize;

for chunk in chunks {
let chunk_array = chunk.clone().execute::<ListViewArray>(ctx)?;
// By rebuilding as zero-copy to `List` and trimming all elements (to prevent gaps), we make
// the final output `ListView` also zero-copyable to `List`.
let chunk_array = chunk_array.rebuild(ListViewRebuildMode::MakeExact, ctx)?;

// Add the `elements` of the current array as a new chunk.
list_elements_chunks.push(chunk_array.elements().clone());

// Cast offsets and sizes to `u64`.
let offsets_arr = chunk_array
.offsets()
.clone()
.cast(DType::Primitive(PType::U64, Nullability::NonNullable))
.vortex_expect("Must be able to fit array offsets in u64")
.execute::<PrimitiveArray>(ctx)?;

let sizes_arr = chunk_array
.sizes()
.clone()
.cast(DType::Primitive(PType::U64, Nullability::NonNullable))
.vortex_expect("Must be able to fit array offsets in u64")
.execute::<PrimitiveArray>(ctx)?;

let offsets_slice = offsets_arr.as_slice::<u64>();
let sizes_slice = sizes_arr.as_slice::<u64>();

// Append offsets and sizes, adjusting offsets to point into the combined array.
for (&offset, &size) in offsets_slice.iter().zip(sizes_slice.iter()) {
offsets_out[next_list] = offset + num_elements;
sizes_slice_out[next_list] = size;
next_list += 1;
}

num_elements += chunk_array.elements().len() as u64;
}
debug_assert_eq!(next_list, len);

// SAFETY: elements are sliced from valid `ListViewArray`s (from `to_listview()`).
let chunked_elements =
unsafe { ChunkedArray::new_unchecked(list_elements_chunks, elem_dtype.clone()) }
.into_array();

let offsets = PrimitiveArray::new(
Buffer::<u64>::from_byte_buffer(offsets.freeze()),
Validity::NonNullable,
)
.into_array();
let sizes = PrimitiveArray::new(
Buffer::<u64>::from_byte_buffer(sizes.freeze()),
Validity::NonNullable,
)
.into_array();

// SAFETY:
// - `offsets` and `sizes` are non-nullable u64 arrays of the same length
// - Each `offset[i] + size[i]` list view is within bounds of elements array because it came
// from valid chunks
// - Validity came from the outer chunked array so it must have the same length
// - Since we made sure that all chunks were zero-copyable to a list above, we know that the
// final concatenated output is also zero-copyable to a list.
Ok(unsafe {
ListViewArray::new_unchecked(chunked_elements, offsets, sizes, validity)
.with_zero_copy_to_list(true)
})
}

/// Packs [`FixedSizeListArray`]s together into a single [`FixedSizeListArray`] whose `elements`
/// child is a [`ChunkedArray`].
///
/// Every chunk shares the same `list_size`, and each chunk's `elements` child is exactly
/// `list_size * chunk.len()` long and starts at the first list, so we can reuse the chunks'
/// `elements` children directly as the chunks of a combined `elements` array without copying.
///
/// The caller guarantees there are at least 2 chunks.
fn swizzle_fixed_size_list_chunks(
chunks: &[ArrayRef],
validity: Validity,
elem_dtype: &DType,
list_size: u32,
ctx: &mut ExecutionCtx,
) -> VortexResult<FixedSizeListArray> {
let len: usize = chunks.iter().map(|c| c.len()).sum();

let mut element_chunks = Vec::with_capacity(chunks.len());
for chunk in chunks {
let chunk_array = chunk.clone().execute::<FixedSizeListArray>(ctx)?;
// A canonical `FixedSizeListArray` keeps its `elements` child trimmed to exactly
// `list_size * chunk.len()` starting at the first list, so the children concatenate
// cleanly into the combined `elements` array.
element_chunks.push(chunk_array.elements().clone());
}

let chunked_elements = ChunkedArray::try_new(element_chunks, elem_dtype.clone())?.into_array();

FixedSizeListArray::try_new(chunked_elements, list_size, validity, len)
}

#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::LazyLock;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

use rstest::rstest;
use vortex_buffer::buffer;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_err;
Expand All @@ -301,15 +122,22 @@ mod tests {
use crate::Canonical;
use crate::IntoArray;
use crate::VortexSessionExecute;
use crate::arrays::Chunked;
use crate::arrays::ChunkedArray;
use crate::arrays::ConstantArray;
use crate::arrays::FixedSizeList;
use crate::arrays::FixedSizeListArray;
use crate::arrays::ListArray;
use crate::arrays::ListView;
use crate::arrays::ListViewArray;
use crate::arrays::PrimitiveArray;
use crate::arrays::Struct;
use crate::arrays::StructArray;
use crate::arrays::VarBinViewArray;
use crate::arrays::VariantArray;
use crate::arrays::chunked::ChunkedArrayExt;
use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt;
use crate::arrays::listview::ListViewArraySlotsExt;
use crate::arrays::struct_::StructArrayExt;
use crate::arrays::variant::VariantArraySlotsExt;
use crate::assert_arrays_eq;
Expand All @@ -318,32 +146,12 @@ mod tests {
use crate::dtype::DType::Variant as VariantDType;
use crate::dtype::Nullability::NonNullable;
use crate::dtype::PType::I32;
use crate::memory::DefaultHostAllocator;
use crate::memory::HostAllocator;
use crate::memory::MemorySessionExt;
use crate::memory::WritableHostBuffer;
use crate::scalar::Scalar;
use crate::validity::Validity;

/// A shared session for these chunked-array tests, used to create execution contexts.
static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);

#[derive(Debug)]
struct CountingAllocator {
allocations: Arc<AtomicUsize>,
}

impl HostAllocator for CountingAllocator {
fn allocate(
&self,
len: usize,
alignment: vortex_buffer::Alignment,
) -> VortexResult<WritableHostBuffer> {
self.allocations.fetch_add(1, Ordering::Relaxed);
DefaultHostAllocator.allocate(len, alignment)
}
}

fn variant_scalar(value: i32) -> Scalar {
Scalar::variant(Scalar::primitive(value, NonNullable))
}
Expand Down Expand Up @@ -660,38 +468,53 @@ mod tests {
Ok(())
}

#[test]
fn list_canonicalize_uses_memory_session_allocator() {
let allocations = Arc::new(AtomicUsize::new(0));
let session = crate::array_session().with_allocator(Arc::new(CountingAllocator {
allocations: Arc::clone(&allocations),
}));
let mut ctx = session.create_execution_ctx();

let l1 = ListArray::try_new(
buffer![1, 2, 3, 4].into_array(),
buffer![0, 3].into_array(),
Validity::NonNullable,
)
.unwrap();
let l2 = ListArray::try_new(
buffer![5, 6].into_array(),
/// Canonicalizing a `ChunkedArray` reuses each chunk's children instead of concatenating
/// them: a nested builder keeps a child chunked on exactly the boundaries it was appended on,
/// however short the appended chunk is.
#[rstest]
#[case::struct_(
StructArray::try_from_iter([("a", buffer![1i32, 2])])
.vortex_expect("struct array")
.into_array(),
|array: &ArrayRef| array.as_::<Struct>().unmasked_field(0).clone()
)]
#[case::fixed_size_list(
FixedSizeListArray::new(buffer![1i32, 2].into_array(), 2, Validity::NonNullable, 1)
.into_array(),
|array: &ArrayRef| array.as_::<FixedSizeList>().elements().clone()
)]
#[case::list(
ListArray::try_new(
buffer![1i32, 2].into_array(),
buffer![0, 2].into_array(),
Validity::NonNullable,
)
.unwrap();
.vortex_expect("list array")
.into_array(),
|array: &ArrayRef| array.as_::<ListView>().elements().clone()
)]
fn canonicalize_reuses_chunk_children(
#[case] chunk: ArrayRef,
#[case] child_of: fn(&ArrayRef) -> ArrayRef,
) -> VortexResult<()> {
let mut ctx = SESSION.create_execution_ctx();

let chunked_list = ChunkedArray::try_new(
vec![l1.into_array(), l2.into_array()],
List(Arc::new(Primitive(I32, NonNullable)), NonNullable),
)
.unwrap()
.into_array();
let chunked =
ChunkedArray::try_new(vec![chunk.clone(), chunk.clone()], chunk.dtype().clone())?
.into_array();
let canonical = chunked.execute::<Canonical>(&mut ctx)?.into_array();

drop(chunked_list.execute::<Canonical>(&mut ctx).unwrap());
assert!(
allocations.load(Ordering::Relaxed) >= 2,
"expected offset+size allocations through MemorySession"
let child = child_of(&canonical);
assert_eq!(
child.as_::<Chunked>().nchunks(),
2,
"each chunk's child should have become a chunk of the combined child",
);

let expected =
ChunkedArray::try_new(vec![chunk.clone(), chunk], canonical.dtype().clone())?;
assert_arrays_eq!(&canonical, &expected, &mut ctx);

Ok(())
}
}
Loading