diff --git a/Cargo.lock b/Cargo.lock index a92a0f5be59..ff0734f31aa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10319,7 +10319,6 @@ dependencies = [ name = "vortex-layout" version = "0.1.0" dependencies = [ - "arcref", "arrow-array 58.4.0", "arrow-schema 58.4.0", "async-stream", @@ -10335,13 +10334,10 @@ dependencies = [ "oneshot", "parking_lot", "paste", - "pin-project-lite", "prost 0.14.4", "rstest", "rustc-hash", "sketches-ddsketch", - "temp-env", - "termtree", "tokio", "tracing", "vortex-array", @@ -10351,6 +10347,7 @@ dependencies = [ "vortex-error", "vortex-flatbuffers", "vortex-io", + "vortex-layout-commons", "vortex-mask", "vortex-metrics", "vortex-runend", @@ -10360,6 +10357,29 @@ dependencies = [ "vortex-utils", ] +[[package]] +name = "vortex-layout-commons" +version = "0.1.0" +dependencies = [ + "arcref", + "async-trait", + "flatbuffers", + "futures", + "itertools 0.14.0", + "once_cell", + "parking_lot", + "pin-project-lite", + "rstest", + "termtree", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-flatbuffers", + "vortex-mask", + "vortex-session", + "vortex-utils", +] + [[package]] name = "vortex-mask" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 65452f18300..19d3b30f4b4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ members = [ "vortex-json", "vortex-compressor", "vortex-btrblocks", + "vortex-layout-commons", "vortex-layout", "vortex-scan", "vortex-file", @@ -316,6 +317,7 @@ vortex-io = { version = "0.1.0", path = "./vortex-io", default-features = false vortex-ipc = { version = "0.1.0", path = "./vortex-ipc", default-features = false } vortex-json = { version = "0.1.0", path = "./vortex-json", default-features = false } vortex-layout = { version = "0.1.0", path = "./vortex-layout", default-features = false } +vortex-layout-commons = { version = "0.1.0", path = "./vortex-layout-commons", default-features = false } vortex-mask = { version = "0.1.0", path = "./vortex-mask", default-features = false } vortex-metrics = { version = "0.1.0", path = "./vortex-metrics", default-features = false } vortex-onpair = { version = "0.1.0", path = "./encodings/onpair", default-features = false } diff --git a/vortex-cuda/src/file.rs b/vortex-cuda/src/file.rs index 2a9806729f7..263a51f748a 100644 --- a/vortex-cuda/src/file.rs +++ b/vortex-cuda/src/file.rs @@ -174,7 +174,9 @@ mod tests { use vortex::array::dtype::StructFields; use vortex::buffer::ByteBuffer; use vortex::layout::layouts::flat::FlatLayout; + use vortex::layout::layouts::flat::FlatLayoutExt; use vortex::layout::layouts::zoned::ZonedLayout; + use vortex::layout::layouts::zoned::ZonedLayoutExt; use vortex::session::registry::ReadContext; use super::*; diff --git a/vortex-layout-commons/Cargo.toml b/vortex-layout-commons/Cargo.toml new file mode 100644 index 00000000000..ba23b5170d5 --- /dev/null +++ b/vortex-layout-commons/Cargo.toml @@ -0,0 +1,41 @@ +[package] +name = "vortex-layout-commons" +authors = { workspace = true } +categories = { workspace = true } +description = "Core traits and types for Vortex layouts" +edition = { workspace = true } +homepage = { workspace = true } +include = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +readme = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[package.metadata.docs.rs] +all-features = true + +[dependencies] +arcref = { workspace = true } +async-trait = { workspace = true } +flatbuffers = { workspace = true } +futures = { workspace = true, features = ["alloc", "async-await"] } +itertools = { workspace = true } +once_cell = { workspace = true, features = ["parking_lot"] } +parking_lot = { workspace = true } +pin-project-lite = { workspace = true } +termtree = { workspace = true } +vortex-array = { workspace = true } +vortex-buffer = { workspace = true } +vortex-error = { workspace = true } +vortex-flatbuffers = { workspace = true, features = ["layout"] } +vortex-mask = { workspace = true } +vortex-session = { workspace = true } +vortex-utils = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } + +[lints] +workspace = true diff --git a/vortex-layout-commons/src/children.rs b/vortex-layout-commons/src/children.rs new file mode 100644 index 00000000000..15324132825 --- /dev/null +++ b/vortex-layout-commons/src/children.rs @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Debug; +use std::fmt::Formatter; +use std::sync::Arc; + +use vortex_array::dtype::DType; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use crate::LayoutRef; + +/// Abstract way of accessing the children of a layout. +/// +/// This allows layout trees to use lazy serialized children as well as in-memory children. +pub trait LayoutChildren: 'static + Send + Sync { + fn to_arc(&self) -> Arc; + + fn child(&self, idx: usize, dtype: &DType) -> VortexResult; + + fn child_row_count(&self, idx: usize) -> u64; + + fn nchildren(&self) -> usize; + + /// Returns `true` if the child at `idx` is known, without materializing it, to be indivisible. + /// + /// Implementations must conservatively return `false` when answering would require + /// materializing the child. + fn child_is_indivisible(&self, _idx: usize) -> bool { + false + } +} + +impl Debug for dyn LayoutChildren { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LayoutChildren") + .field("nchildren", &self.nchildren()) + .finish() + } +} + +impl LayoutChildren for Arc { + fn to_arc(&self) -> Arc { + Arc::clone(self) + } + + fn child(&self, idx: usize, dtype: &DType) -> VortexResult { + self.as_ref().child(idx, dtype) + } + + fn child_row_count(&self, idx: usize) -> u64 { + self.as_ref().child_row_count(idx) + } + + fn nchildren(&self) -> usize { + self.as_ref().nchildren() + } + + fn child_is_indivisible(&self, idx: usize) -> bool { + self.as_ref().child_is_indivisible(idx) + } +} + +/// In-memory owned layout children. +#[derive(Clone)] +pub struct OwnedLayoutChildren(Vec); + +impl OwnedLayoutChildren { + pub fn layout_children(children: Vec) -> Arc { + Arc::new(Self(children)) + } +} + +/// Create an in-memory child adapter from owned layout references. +pub fn layout_children(children: Vec) -> Arc { + OwnedLayoutChildren::layout_children(children) +} + +impl LayoutChildren for OwnedLayoutChildren { + fn to_arc(&self) -> Arc { + Arc::new(self.clone()) + } + + fn child(&self, idx: usize, dtype: &DType) -> VortexResult { + if idx >= self.0.len() { + vortex_bail!("Child index out of bounds: {} of {}", idx, self.0.len()); + } + let child = &self.0[idx]; + if child.dtype() != dtype { + vortex_bail!("Child dtype mismatch: {} != {}", child.dtype(), dtype); + } + Ok(Arc::clone(child)) + } + + fn child_row_count(&self, idx: usize) -> u64 { + self.0[idx].row_count() + } + + fn nchildren(&self) -> usize { + self.0.len() + } + + fn child_is_indivisible(&self, idx: usize) -> bool { + self.0[idx].dyn_is_indivisible() + } +} diff --git a/vortex-layout-commons/src/display.rs b/vortex-layout-commons/src/display.rs new file mode 100644 index 00000000000..6e77e46bf89 --- /dev/null +++ b/vortex-layout-commons/src/display.rs @@ -0,0 +1,173 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use futures::future::try_join_all; +use termtree::Tree; +use vortex_array::serde::SerializedArray; +use vortex_error::VortexResult; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::LayoutRef; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; + +/// Display the layout as a tree, fetching buffer sizes from the segment source. +/// +/// # Warning +/// +/// This function performs IO to fetch each referenced segment. +pub(crate) async fn display_tree_with_segment_sizes( + layout: LayoutRef, + segment_source: Arc, +) -> VortexResult { + let mut segment_ids = Vec::new(); + let mut segment_buffer_sizes = HashMap::new(); + collect_segment_ids(&layout, &mut segment_ids, &mut segment_buffer_sizes)?; + segment_ids.sort_unstable(); + segment_ids.dedup(); + + let fetches = segment_ids.into_iter().map(|segment_id| { + let segment_source = Arc::clone(&segment_source); + async move { + let buffer = segment_source.request(segment_id).await?; + let parts = SerializedArray::try_from(buffer)?; + VortexResult::Ok((segment_id, parts.buffer_lengths())) + } + }); + segment_buffer_sizes.extend(try_join_all(fetches).await?); + + Ok(DisplayLayoutTree { + layout, + segment_buffer_sizes: Some(segment_buffer_sizes), + verbose: true, + }) +} + +fn collect_segment_ids( + layout: &LayoutRef, + segment_ids: &mut Vec, + segment_buffer_sizes: &mut HashMap>, +) -> VortexResult<()> { + let inlined_sizes = layout.inlined_segment_buffer_sizes(); + segment_buffer_sizes.extend(inlined_sizes.iter().cloned()); + segment_ids.extend( + layout + .segment_ids() + .into_iter() + .filter(|id| !inlined_sizes.iter().any(|(inlined_id, _)| inlined_id == id)), + ); + for child in layout.children()? { + collect_segment_ids(&child, segment_ids, segment_buffer_sizes)?; + } + Ok(()) +} + +/// Display wrapper for a layout tree. +pub struct DisplayLayoutTree { + layout: LayoutRef, + segment_buffer_sizes: Option>>, + verbose: bool, +} + +impl DisplayLayoutTree { + /// Create a layout tree display without fetching segment data. + pub fn new(layout: LayoutRef, verbose: bool) -> Self { + Self { + layout, + segment_buffer_sizes: None, + verbose, + } + } + + fn make_tree(&self, layout: LayoutRef) -> VortexResult> { + let mut node_parts = vec![ + layout.encoding_id().to_string(), + format!("dtype: {}", layout.dtype()), + ]; + + if layout.nchildren() > 0 { + node_parts.push(format!("children: {}", layout.nchildren())); + } + + if self.verbose { + let metadata = layout.metadata(); + if !metadata.is_empty() { + node_parts.push(format!("metadata: {} bytes", metadata.len())); + } + node_parts.push(format!("rows: {}", layout.row_count())); + } + + let segments = layout.segment_ids(); + if segments.len() == 1 { + let segment_id = segments[0]; + let inlined_sizes = layout.inlined_segment_buffer_sizes(); + if let Some(buffer_sizes) = self + .segment_buffer_sizes + .as_ref() + .and_then(|sizes| sizes.get(&segment_id)) + .or_else(|| { + inlined_sizes + .iter() + .find_map(|(id, sizes)| (*id == segment_id).then_some(sizes)) + }) + { + node_parts.push(format_buffer_sizes(buffer_sizes, *segment_id)); + } else if !self.verbose { + node_parts.push(format!("segment: {}", *segment_id)); + } else { + node_parts.push(format!("segments: [{}]", *segment_id)); + } + } else if !segments.is_empty() && self.verbose { + node_parts.push(format!( + "segments: [{}]", + segments + .iter() + .map(|segment| (**segment).to_string()) + .collect::>() + .join(", ") + )); + } + + let children = layout.children()?; + let child_names = layout.child_names().collect::>(); + let child_trees = if child_names.len() == children.len() { + children + .into_iter() + .zip(child_names) + .map(|(child, name)| { + let child_tree = self.make_tree(child)?; + Ok(Tree::new(format!("{name}: {}", child_tree.root)) + .with_leaves(child_tree.leaves)) + }) + .collect::>>()? + } else { + children + .into_iter() + .map(|child| self.make_tree(child)) + .collect::>>()? + }; + + Ok(Tree::new(node_parts.join(", ")).with_leaves(child_trees)) + } +} + +fn format_buffer_sizes(buffer_sizes: &[usize], segment_id: u32) -> String { + let sizes = buffer_sizes + .iter() + .map(|size| format!("{size}B")) + .collect::>() + .join(", "); + let total = buffer_sizes.iter().sum::(); + format!("segment {segment_id}, buffers=[{sizes}], total={total}B") +} + +impl std::fmt::Display for DisplayLayoutTree { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self.make_tree(Arc::clone(&self.layout)) { + Ok(tree) => write!(f, "{tree}"), + Err(error) => write!(f, "Error building layout tree: {error}"), + } + } +} diff --git a/vortex-layout/src/encoding.rs b/vortex-layout-commons/src/encoding.rs similarity index 100% rename from vortex-layout/src/encoding.rs rename to vortex-layout-commons/src/encoding.rs diff --git a/vortex-layout-commons/src/flatbuffers.rs b/vortex-layout-commons/src/flatbuffers.rs new file mode 100644 index 00000000000..8db3de757bf --- /dev/null +++ b/vortex-layout-commons/src/flatbuffers.rs @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use flatbuffers::FlatBufferBuilder; +use flatbuffers::WIPOffset; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_flatbuffers::FlatBufferRoot; +use vortex_flatbuffers::WriteFlatBuffer; +use vortex_flatbuffers::layout; + +use crate::DynLayout; +use crate::LayoutContext; + +impl dyn DynLayout + '_ { + /// Serialize the layout into a [`FlatBufferBuilder`]. + pub fn flatbuffer_writer<'a>( + &'a self, + ctx: &'a LayoutContext, + ) -> impl WriteFlatBuffer = layout::Layout<'a>> + FlatBufferRoot + 'a { + LayoutFlatBufferWriter { layout: self, ctx } + } +} + +/// An adapter struct for writing a layout to a FlatBuffer. +struct LayoutFlatBufferWriter<'a> { + layout: &'a dyn DynLayout, + ctx: &'a LayoutContext, +} + +impl FlatBufferRoot for LayoutFlatBufferWriter<'_> {} + +impl WriteFlatBuffer for LayoutFlatBufferWriter<'_> { + type Target<'fb> = layout::Layout<'fb>; + + fn write_flatbuffer<'fb>( + &self, + fbb: &mut FlatBufferBuilder<'fb>, + ) -> VortexResult>> { + // First we recurse into the children and write them out + let child_layouts = self.layout.children()?; + let children = child_layouts + .iter() + .map(|layout| { + LayoutFlatBufferWriter { + layout: layout.as_ref(), + ctx: self.ctx, + } + .write_flatbuffer(fbb) + }) + .collect::>>()?; + let children = (!children.is_empty()).then(|| fbb.create_vector(&children)); + + // Next we write out the metadata if it's non-empty. + let metadata = self.layout.metadata(); + let metadata = (!metadata.is_empty()).then(|| fbb.create_vector(&metadata)); + + let segments = self + .layout + .segment_ids() + .into_iter() + .map(|s| *s) + .collect::>(); + let segments = (!segments.is_empty()).then(|| fbb.create_vector(&segments)); + + // Dictionary-encode the layout ID + let encoding = self.ctx.intern(&self.layout.encoding_id()).ok_or_else(|| { + vortex_err!( + "Layout encoding {} not permitted by ctx", + self.layout.encoding_id() + ) + })?; + + Ok(layout::Layout::create( + fbb, + &layout::LayoutArgs { + encoding, + row_count: self.layout.row_count(), + metadata, + children, + segments, + }, + )) + } +} diff --git a/vortex-layout/src/layout.rs b/vortex-layout-commons/src/layout.rs similarity index 96% rename from vortex-layout/src/layout.rs rename to vortex-layout-commons/src/layout.rs index 5ed2f344847..a1485acd9d1 100644 --- a/vortex-layout/src/layout.rs +++ b/vortex-layout-commons/src/layout.rs @@ -259,6 +259,11 @@ pub trait DynLayout: 'static + Send + Sync + Debug { /// Returns directly referenced segment IDs. fn dyn_segment_ids(&self) -> Vec; + /// Returns buffer sizes available from inline layout metadata. + fn dyn_inlined_segment_buffer_sizes(&self) -> Vec<(SegmentId, Vec)> { + Vec::new() + } + /// Constructs a reader. fn dyn_new_reader( &self, @@ -320,6 +325,10 @@ impl DynLayout for Layout { self.inner.segment_ids.clone() } + fn dyn_inlined_segment_buffer_sizes(&self) -> Vec<(SegmentId, Vec)> { + V::inlined_segment_buffer_sizes(self) + } + fn dyn_new_reader( &self, name: Arc, @@ -419,6 +428,11 @@ impl dyn DynLayout + '_ { self.dyn_segment_ids() } + /// Returns buffer sizes available from inline layout metadata. + pub fn inlined_segment_buffer_sizes(&self) -> Vec<(SegmentId, Vec)> { + self.dyn_inlined_segment_buffer_sizes() + } + /// Constructs a reader for this layout. pub fn new_reader( &self, diff --git a/vortex-layout-commons/src/lib.rs b/vortex-layout-commons/src/lib.rs new file mode 100644 index 00000000000..a921cdd5533 --- /dev/null +++ b/vortex-layout-commons/src/lib.rs @@ -0,0 +1,33 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Core traits and types shared by Vortex layout implementations and consumers. +//! +//! Layout extension crates can depend on this crate for the typed [`Layout`] model, [`VTable`], +//! reader and writer traits, and segment IO abstractions without depending on the built-in layouts +//! and scan planner in `vortex-layout`. + +pub mod display; +pub mod segments; +pub mod sequence; + +pub use children::*; +pub use encoding::*; +pub use layout::*; +pub use reader::*; +pub use reader_context::*; +pub use strategy::*; +use vortex_session::registry::Interner; +pub use vtable::*; + +mod children; +mod encoding; +mod flatbuffers; +mod layout; +mod reader; +mod reader_context; +mod strategy; +mod vtable; + +/// Registry context used when serializing layouts. +pub type LayoutContext = Interner; diff --git a/vortex-layout/src/reader.rs b/vortex-layout-commons/src/reader.rs similarity index 97% rename from vortex-layout/src/reader.rs rename to vortex-layout-commons/src/reader.rs index 91dda0e23c5..fb493d65ff0 100644 --- a/vortex-layout/src/reader.rs +++ b/vortex-layout-commons/src/reader.rs @@ -162,8 +162,8 @@ impl RowSplits { self.splits.reserve(additional); } - /// Create a new RowSplits with preallocated "capacity" - pub(crate) fn new_capacity(capacity: usize) -> Self { + /// Creates an empty collection with space for at least `capacity` split points. + pub fn with_capacity(capacity: usize) -> Self { Self { splits: Vec::with_capacity(capacity), run_start: 0, @@ -171,7 +171,8 @@ impl RowSplits { } } - pub(crate) fn into_sorted_deduped(mut self) -> Vec { + /// Returns the split points in ascending order with duplicates removed. + pub fn into_sorted_deduped(mut self) -> Vec { let final_run_dropped = self.drop_repeated_run(); // Surviving runs always have a descent between them, so the boundaries are ascending // iff a single run survived: no run before the final one (`prev_run_start == 0`) and @@ -413,7 +414,7 @@ mod tests { // No pushes at all. #[case(vec![])] fn into_sorted_deduped_matches_model(#[case] runs: Vec>) { - let mut splits = RowSplits::new_capacity(16); + let mut splits = RowSplits::with_capacity(16); let mut model = Vec::new(); for run in &runs { for &row in run { diff --git a/vortex-layout/src/reader_context.rs b/vortex-layout-commons/src/reader_context.rs similarity index 100% rename from vortex-layout/src/reader_context.rs rename to vortex-layout-commons/src/reader_context.rs diff --git a/vortex-layout-commons/src/segments/mod.rs b/vortex-layout-commons/src/segments/mod.rs new file mode 100644 index 00000000000..5f59d6ab182 --- /dev/null +++ b/vortex-layout-commons/src/segments/mod.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Segment access contracts for layouts. +//! +//! Layouts refer to byte ranges by [`SegmentId`]. A [`SegmentSource`] resolves those ids to buffer +//! handles for readers, while a [`SegmentSink`] assigns ids when writers emit buffers. Cache and +//! request-sharing policies belong to the runtime using these contracts. + +mod sink; +mod source; + +use std::fmt::Display; +use std::ops::Deref; + +pub use sink::*; +pub use source::*; +use vortex_error::VortexError; + +/// Identifier for a single physical segment referenced by a layout. +/// +/// Segment ids are local to a file or segment source. The file footer maps ids to physical offsets; +/// custom storage systems may map them to object-store keys or other random-access locations. +// TODO(ngates): should this be a `[u8]` instead? Allowing for arbitrary segment identifiers? +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct SegmentId(u32); + +impl From for SegmentId { + fn from(value: u32) -> Self { + Self(value) + } +} + +impl TryFrom for SegmentId { + type Error = VortexError; + + fn try_from(value: usize) -> Result { + Ok(Self::from(u32::try_from(value)?)) + } +} + +impl Deref for SegmentId { + type Target = u32; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl Display for SegmentId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "SegmentId({})", self.0) + } +} diff --git a/vortex-layout/src/segments/sink.rs b/vortex-layout-commons/src/segments/sink.rs similarity index 100% rename from vortex-layout/src/segments/sink.rs rename to vortex-layout-commons/src/segments/sink.rs diff --git a/vortex-layout/src/segments/source.rs b/vortex-layout-commons/src/segments/source.rs similarity index 100% rename from vortex-layout/src/segments/source.rs rename to vortex-layout-commons/src/segments/source.rs diff --git a/vortex-layout/src/sequence.rs b/vortex-layout-commons/src/sequence.rs similarity index 100% rename from vortex-layout/src/sequence.rs rename to vortex-layout-commons/src/sequence.rs diff --git a/vortex-layout-commons/src/strategy.rs b/vortex-layout-commons/src/strategy.rs new file mode 100644 index 00000000000..4ff4b748936 --- /dev/null +++ b/vortex-layout-commons/src/strategy.rs @@ -0,0 +1,244 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use async_trait::async_trait; +use vortex_array::ArrayContext; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; + +use crate::LayoutRef; +use crate::segments::SegmentSinkRef; +use crate::sequence::SendableSequentialStream; +use crate::sequence::SequencePointer; + +/// A shared counter of the bytes that layout strategies are holding but have not yet emitted. +/// +/// Clones share the same counter, so a tracker can be handed to a writer before the write begins +/// and polled while it runs. Strategies report their own retained bytes with +/// [`Self::reserve`], which releases the reservation on drop. +#[derive(Clone, Debug, Default)] +pub struct BufferedBytesTracker(Arc); + +impl BufferedBytesTracker { + /// Creates a tracker with a zeroed counter. + pub fn new() -> Self { + Self::default() + } + + /// Returns the number of bytes currently retained by layout strategies. + pub fn buffered_bytes(&self) -> u64 { + self.0.load(Ordering::Relaxed) + } + + /// Records `bytes` as buffered until the returned reservation is dropped. + pub fn reserve(&self, bytes: u64) -> BufferedBytesReservation { + self.0.fetch_add(bytes, Ordering::Relaxed); + BufferedBytesReservation { + tracker: self.clone(), + bytes, + } + } +} + +/// An outstanding claim on a [`BufferedBytesTracker`], released when dropped. +#[derive(Debug)] +pub struct BufferedBytesReservation { + tracker: BufferedBytesTracker, + bytes: u64, +} + +impl BufferedBytesReservation { + /// Returns the number of bytes held by this reservation. + pub fn bytes(&self) -> u64 { + self.bytes + } +} + +impl Drop for BufferedBytesReservation { + fn drop(&mut self) { + self.tracker.0.fetch_sub(self.bytes, Ordering::Relaxed); + } +} + +/// State shared by every strategy participating in a single layout write. +/// +/// Clones share the [`BufferedBytesTracker`] while retaining the array serialization context. +/// Passing this context through the strategy tree keeps writer-scoped state independent of the +/// strategy instances, which may be shared by multiple leaves or writers. +#[derive(Clone)] +pub struct LayoutWriterContext { + array_ctx: ArrayContext, + allowed_aggregates: Option>>, + buffered_bytes: BufferedBytesTracker, +} + +impl LayoutWriterContext { + /// Creates a context for a layout write with a fresh buffered bytes tracker. + pub fn new(array_ctx: ArrayContext) -> Self { + Self { + array_ctx, + allowed_aggregates: None, + buffered_bytes: BufferedBytesTracker::new(), + } + } + + /// Restrict the aggregate functions this write may record, e.g. in a zone map. + /// + /// A write that would record an aggregate outside `allowed` fails, matching the array and + /// layout contexts: a silently thinner zone map is a file that prunes worse than the + /// caller asked for, with nothing in the output saying so. The id set is a plain set of + /// ids — callers that source it from editions resolve it themselves. + pub fn with_allowed_aggregates(mut self, allowed: HashSet) -> Self { + self.allowed_aggregates = Some(Arc::new(allowed)); + self + } + + /// Returns whether `aggregate` may be recorded by this write. Unrestricted contexts + /// permit every aggregate. + pub fn allows_aggregate(&self, aggregate: &AggregateFnId) -> bool { + self.allowed_aggregates + .as_ref() + .is_none_or(|allowed| allowed.contains(aggregate)) + } + + /// Replaces the buffered bytes tracker, so callers can observe the counter from outside the + /// strategy tree. + pub fn with_buffered_bytes_tracker(mut self, tracker: BufferedBytesTracker) -> Self { + self.buffered_bytes = tracker; + self + } + + /// Returns the array serialization context. + pub fn array_ctx(&self) -> &ArrayContext { + &self.array_ctx + } + + /// Returns the tracker that accounts for bytes retained by layout strategies. + pub fn buffered_bytes_tracker(&self) -> &BufferedBytesTracker { + &self.buffered_bytes + } + + /// Returns the number of bytes currently retained by layout strategies. + pub fn buffered_bytes(&self) -> u64 { + self.buffered_bytes.buffered_bytes() + } + + /// Records `bytes` as retained by this write until the returned reservation is dropped. + pub fn reserve_buffered_bytes(&self, bytes: u64) -> BufferedBytesReservation { + self.buffered_bytes.reserve(bytes) + } +} + +impl From for LayoutWriterContext { + fn from(array_ctx: ArrayContext) -> Self { + Self::new(array_ctx) + } +} + +/// Writes an ordered array stream into a layout tree and segment sink. +/// +/// Layout strategies are writer-side extension points. Strategies may repartition, buffer, +/// collect columns, compute statistics, compress arrays, or delegate to child strategies before +/// finally emitting segments. They must preserve the logical row order represented by the +/// [`SequencePointer`]s in the input stream. +#[async_trait] +pub trait LayoutStrategy: 'static + Send + Sync { + /// Asynchronously process an ordered stream of array chunks, emitting them into a sink and + /// returning the [`Layout`][crate::Layout] instance that can be parsed to retrieve the data + /// from rest. + /// + /// This trait uses the `#[async_trait]` attribute to denote that trait objects of this type + /// can be `Box`ed or `Arc`ed and shared around. Commonly, these strategies are composed to + /// form a operator of operations, each of which modifies the chunk stream in some way before + /// passing the data on to a downstream writer. + /// + /// # Sequencing and EOF + /// + /// The `stream` parameter is a stream of ordered array chunks, each of which is associated + /// with a sequence pointer that indicates its position in the overall array. By passing + /// around these pointers (essentially vector clocks), the writer can support concurrent + /// and parallel processing while maintaining a deterministic order of data in the file. + /// The `ctx` parameter carries both array serialization state and writer-scoped accounting + /// through every child strategy. + /// + /// The `eof` parameter is a guaranteed to be greater than all sequence pointers in the stream. + /// + /// Because child strategies can write to the end-of-file pointer, it is very important that + /// **all strategies must await all children concurrently**. Otherwise it is possible to + /// deadlock if one child is waiting to write to EOF while your strategy is preventing the + /// stream from progressing to completion. + /// + /// # Blocking operations + /// + /// This is an async trait method, which will return a `BoxFuture` that you can await from + /// any runtime. Implementations should avoid directly performing blocking work within the + /// `write_stream`, and should instead spawn it onto an appropriate runtime or threadpool + /// dedicated to such work. + /// + /// Such operations are common, and include things like compression and parsing large blobs + /// of data, or serializing very large messages to flatbuffers. + async fn write_stream( + &self, + ctx: LayoutWriterContext, + segment_sink: SegmentSinkRef, + stream: SendableSequentialStream, + eof: SequencePointer, + session: &VortexSession, + ) -> VortexResult; +} + +#[async_trait] +impl LayoutStrategy for Arc { + async fn write_stream( + &self, + ctx: LayoutWriterContext, + segment_sink: SegmentSinkRef, + stream: SendableSequentialStream, + eof: SequencePointer, + session: &VortexSession, + ) -> VortexResult { + (**self) + .write_stream(ctx, segment_sink, stream, eof, session) + .await + } +} + +#[cfg(test)] +mod tests { + use crate::strategy::BufferedBytesTracker; + + #[test] + fn reservations_accumulate_and_release() { + let tracker = BufferedBytesTracker::new(); + assert_eq!(tracker.buffered_bytes(), 0); + + let first = tracker.reserve(16); + let second = tracker.reserve(32); + assert_eq!(tracker.buffered_bytes(), 48); + assert_eq!(first.bytes(), 16); + + drop(first); + assert_eq!(tracker.buffered_bytes(), 32); + + drop(second); + assert_eq!(tracker.buffered_bytes(), 0); + } + + #[test] + fn clones_share_the_same_counter() { + let tracker = BufferedBytesTracker::new(); + let observer = tracker.clone(); + + let reservation = tracker.reserve(8); + assert_eq!(observer.buffered_bytes(), 8); + + drop(reservation); + assert_eq!(observer.buffered_bytes(), 0); + } +} diff --git a/vortex-layout/src/vtable.rs b/vortex-layout-commons/src/vtable.rs similarity index 93% rename from vortex-layout/src/vtable.rs rename to vortex-layout-commons/src/vtable.rs index 7ec27b67046..509123ea2c4 100644 --- a/vortex-layout/src/vtable.rs +++ b/vortex-layout-commons/src/vtable.rs @@ -105,6 +105,14 @@ pub trait VTable: 'static + Clone + Send + Sync + Debug { /// Returns the relationship between the child in logical `slot` and its parent. fn child_type(layout: &Layout, slot: usize) -> LayoutChildType; + /// Returns buffer sizes available from inline layout metadata. + /// + /// This is used by layout tree displays to avoid fetching segments whose serialized array + /// metadata is already embedded in the layout. + fn inlined_segment_buffer_sizes(_layout: &Layout) -> Vec<(SegmentId, Vec)> { + Vec::new() + } + /// Construct a reader for this layout. fn new_reader( layout: &Layout, diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index f772b9ab639..7ba5441fd52 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -17,7 +17,6 @@ version = { workspace = true } all-features = true [dependencies] -arcref = { workspace = true } arrow-array = { workspace = true } arrow-schema = { workspace = true } async-stream = { workspace = true } @@ -32,11 +31,9 @@ once_cell = { workspace = true, features = ["parking_lot"] } oneshot = { workspace = true } parking_lot = { workspace = true } paste = { workspace = true } -pin-project-lite = { workspace = true } prost = { workspace = true } rustc-hash = { workspace = true } sketches-ddsketch = { workspace = true } -termtree = { workspace = true } tokio = { workspace = true, features = ["rt"], optional = true } tracing = { workspace = true } vortex-array = { workspace = true } @@ -46,6 +43,7 @@ vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-flatbuffers = { workspace = true, features = ["layout"] } vortex-io = { workspace = true } +vortex-layout-commons = { workspace = true } vortex-mask = { workspace = true } vortex-metrics = { workspace = true } vortex-runend = { workspace = true } @@ -58,7 +56,6 @@ vortex-utils = { workspace = true, features = ["dashmap"] } futures = { workspace = true, features = ["executor"] } insta = { workspace = true } rstest = { workspace = true } -temp-env = { workspace = true } tokio = { workspace = true, features = ["rt", "macros"] } vortex-array = { path = "../vortex-array", features = ["_test-harness"] } vortex-io = { path = "../vortex-io", features = ["tokio"] } diff --git a/vortex-layout/src/children.rs b/vortex-layout/src/children.rs index 274ef124a60..fd79eac1b78 100644 --- a/vortex-layout/src/children.rs +++ b/vortex-layout/src/children.rs @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::fmt::Debug; -use std::fmt::Formatter; use std::sync::Arc; use flatbuffers::Follow; @@ -14,114 +12,17 @@ use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_flatbuffers::FlatBuffer; use vortex_flatbuffers::layout as fbl; +use vortex_layout_commons::LayoutBuildContext; +use vortex_layout_commons::LayoutChildren; +use vortex_layout_commons::LayoutRef; +pub(crate) use vortex_layout_commons::OwnedLayoutChildren; +use vortex_layout_commons::segments::SegmentId; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; -use crate::LayoutBuildContext; -use crate::LayoutRef; use crate::layouts::foreign::new_foreign_layout; -use crate::segments::SegmentId; use crate::session::LayoutRegistry; -/// Abstract way of accessing the children of a layout. -/// -/// This allows us to abstract over the lazy flatbuffer-based layouts, as well as the in-memory -/// layout trees. -pub trait LayoutChildren: 'static + Send + Sync { - fn to_arc(&self) -> Arc; - - fn child(&self, idx: usize, dtype: &DType) -> VortexResult; - - fn child_row_count(&self, idx: usize) -> u64; - - fn nchildren(&self) -> usize; - - /// Returns `true` if the child at `idx` is known — without materializing it — to be - /// indivisible: it registers no split boundaries strictly inside its row range (see - /// [`VTable::is_indivisible`](crate::VTable::is_indivisible)). - /// - /// Implementations must conservatively return `false` when answering would require - /// materializing the child. - fn child_is_indivisible(&self, _idx: usize) -> bool { - false - } -} - -impl Debug for dyn LayoutChildren { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.debug_struct("LayoutChildren") - .field("nchildren", &self.nchildren()) - .finish() - } -} - -impl LayoutChildren for Arc { - fn to_arc(&self) -> Arc { - Arc::clone(self) - } - - fn child(&self, idx: usize, dtype: &DType) -> VortexResult { - self.as_ref().child(idx, dtype) - } - - fn child_row_count(&self, idx: usize) -> u64 { - self.as_ref().child_row_count(idx) - } - - fn nchildren(&self) -> usize { - self.as_ref().nchildren() - } - - fn child_is_indivisible(&self, idx: usize) -> bool { - self.as_ref().child_is_indivisible(idx) - } -} - -/// An implementation of [`LayoutChildren`] for in-memory owned children. -#[derive(Clone)] -pub(crate) struct OwnedLayoutChildren(Vec); - -impl OwnedLayoutChildren { - pub fn layout_children(children: Vec) -> Arc { - Arc::new(Self(children)) - } -} - -/// Create an in-memory child adapter from owned layout references. -pub fn layout_children(children: Vec) -> Arc { - OwnedLayoutChildren::layout_children(children) -} - -/// In-memory implementation of [`LayoutChildren`]. -impl LayoutChildren for OwnedLayoutChildren { - fn to_arc(&self) -> Arc { - Arc::new(self.clone()) - } - - fn child(&self, idx: usize, dtype: &DType) -> VortexResult { - if idx >= self.0.len() { - vortex_bail!("Child index out of bounds: {} of {}", idx, self.0.len()); - } - let child = &self.0[idx]; - if child.dtype() != dtype { - vortex_bail!("Child dtype mismatch: {} != {}", child.dtype(), dtype); - } - Ok(Arc::clone(child)) - } - - fn child_row_count(&self, idx: usize) -> u64 { - self.0[idx].row_count() - } - - fn nchildren(&self) -> usize { - self.0.len() - } - - fn child_is_indivisible(&self, idx: usize) -> bool { - self.0[idx].dyn_is_indivisible() - } -} - #[derive(Clone)] pub(crate) struct ViewedLayoutChildren { flatbuffer: FlatBuffer, diff --git a/vortex-layout/src/display.rs b/vortex-layout/src/display.rs deleted file mode 100644 index c3743a1df45..00000000000 --- a/vortex-layout/src/display.rs +++ /dev/null @@ -1,489 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; - -use futures::future::try_join_all; -use termtree::Tree; -use vortex_array::serde::SerializedArray; -use vortex_error::VortexResult; -use vortex_utils::aliases::hash_map::HashMap; - -use crate::LayoutRef; -use crate::layouts::flat::Flat; -use crate::layouts::flat::FlatLayout; -use crate::segments::SegmentId; -use crate::segments::SegmentSource; - -/// Display the layout as a tree, fetching segment sizes from the segment source. -/// -/// # Warning -/// -/// This function performs IO to fetch each segment's buffer. For layouts with -/// many segments, this may result in significant IO overhead. -pub(super) async fn display_tree_with_segment_sizes( - layout: LayoutRef, - segment_source: Arc, -) -> VortexResult { - // First, collect all segment IDs from the layout tree (excluding those with inline array_tree) - let mut segments_to_fetch = Vec::new(); - collect_segments_to_fetch(&layout, &mut segments_to_fetch)?; - segments_to_fetch.dedup(); - - // Fetch segments in parallel and parse buffer info - let fetch_futures = segments_to_fetch.iter().map(|&segment_id| { - let segment_source = Arc::clone(&segment_source); - async move { - let buffer = segment_source.request(segment_id).await?; - let parts = SerializedArray::try_from(buffer)?; - VortexResult::Ok((segment_id, parts.buffer_lengths())) - } - }); - let results = try_join_all(fetch_futures).await?; - let segment_buffer_sizes: HashMap> = results.into_iter().collect(); - - Ok(DisplayLayoutTree { - layout, - segment_buffer_sizes: Some(segment_buffer_sizes), - verbose: true, - }) -} - -/// Collect segment IDs that need to be fetched (those without inline array_tree). -fn collect_segments_to_fetch( - layout: &LayoutRef, - segment_ids: &mut Vec, -) -> VortexResult<()> { - // For FlatLayout, only add if there's no inline array_tree - if let Some(flat_layout) = layout.as_opt::() { - if flat_layout.array_tree().is_none() { - segment_ids.push(flat_layout.segment_id()); - } - } else { - // For other layouts, add all segment IDs - segment_ids.extend(layout.segment_ids()); - } - - // Recurse into children - for child in layout.children()? { - collect_segments_to_fetch(&child, segment_ids)?; - } - Ok(()) -} - -/// Build a tree node for a FlatLayout, showing buffer sizes. -fn format_flat_layout_buffers( - flat_layout: &FlatLayout, - segment_buffer_sizes: Option<&HashMap>>, -) -> String { - let segment_id = flat_layout.segment_id(); - - // First, try to get buffer info from inline array_tree - if let Some(array_tree) = flat_layout.array_tree() - && let Ok(parts) = SerializedArray::from_array_tree(array_tree.as_ref().to_vec()) - { - return format_buffer_sizes(&parts.buffer_lengths(), *segment_id); - } - - // Otherwise, try to get from fetched segment info - if let Some(sizes_map) = segment_buffer_sizes - && let Some(buffer_sizes) = sizes_map.get(&segment_id) - { - return format_buffer_sizes(buffer_sizes, *segment_id); - } - - // Fallback: just show segment ID - format!("segment: {}", *segment_id) -} - -fn format_buffer_sizes(buffer_sizes: &[usize], segment_id: u32) -> String { - let buffer_sizes_str: Vec = buffer_sizes.iter().map(|s| format!("{}B", s)).collect(); - let total: usize = buffer_sizes.iter().sum(); - format!( - "segment {}, buffers=[{}], total={}B", - segment_id, - buffer_sizes_str.join(", "), - total - ) -} - -/// Display wrapper for layout tree visualization. -pub struct DisplayLayoutTree { - layout: LayoutRef, - segment_buffer_sizes: Option>>, - verbose: bool, -} - -impl DisplayLayoutTree { - /// Create a new display tree without pre-fetched segment buffer sizes. - pub fn new(layout: LayoutRef, verbose: bool) -> Self { - Self { - layout, - segment_buffer_sizes: None, - verbose, - } - } - - fn make_tree(&self, layout: LayoutRef) -> VortexResult> { - // Build the node label with encoding, dtype, and metadata - let mut node_parts = vec![ - format!("{}", layout.encoding_id()), - format!("dtype: {}", layout.dtype()), - ]; - - // Add child count if there are children - let nchildren = layout.nchildren(); - if nchildren > 0 { - node_parts.push(format!("children: {}", nchildren)); - } - - // Add metadata and row count if verbose - if self.verbose { - let metadata = layout.metadata(); - if !metadata.is_empty() { - node_parts.push(format!("metadata: {} bytes", metadata.len())); - } - node_parts.push(format!("rows: {}", layout.row_count())); - } - - // For FlatLayout, show buffer info - if let Some(flat_layout) = layout.as_opt::() { - node_parts.push(format_flat_layout_buffers( - flat_layout, - self.segment_buffer_sizes.as_ref(), - )); - } else { - // Not a FlatLayout - show segment IDs if any (for verbose mode) - if self.verbose { - let segment_ids = layout.segment_ids(); - if !segment_ids.is_empty() { - node_parts.push(format!( - "segments: [{}]", - segment_ids - .iter() - .map(|s| format!("{}", **s)) - .collect::>() - .join(", ") - )); - } - } - } - - let node_name = node_parts.join(", "); - - // Get children and child names directly from the layout - let children = layout.children()?; - let child_names: Vec<_> = layout.child_names().collect(); - - // Build child trees - let child_trees: VortexResult>> = - if !children.is_empty() && child_names.len() == children.len() { - // If we have names for all children, use them - children - .into_iter() - .zip(child_names.iter()) - .map(|(child, name)| { - let child_tree = self.make_tree(child)?; - Ok(Tree::new(format!("{}: {}", name, child_tree.root)) - .with_leaves(child_tree.leaves)) - }) - .collect() - } else if !children.is_empty() { - // No names available, just show children - children.into_iter().map(|c| self.make_tree(c)).collect() - } else { - // Leaf node - no children - Ok(Vec::new()) - }; - - Ok(Tree::new(node_name).with_leaves(child_trees?)) - } -} - -impl std::fmt::Display for DisplayLayoutTree { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self.make_tree(Arc::clone(&self.layout)) { - Ok(tree) => write!(f, "{}", tree), - Err(e) => write!(f, "Error building layout tree: {}", e), - } - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use vortex_array::ArrayContext; - use vortex_array::IntoArray; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::builders::ArrayBuilder; - use vortex_array::builders::VarBinViewBuilder; - use vortex_array::dtype::DType; - use vortex_array::dtype::FieldName; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::Nullability::NonNullable; - use vortex_array::dtype::PType; - use vortex_array::dtype::StructFields; - use vortex_array::serde::SerializedArray; - use vortex_array::validity::Validity; - use vortex_buffer::BitBufferMut; - use vortex_buffer::buffer; - use vortex_io::runtime::single::block_on; - use vortex_io::session::RuntimeSessionExt; - - use crate::OwnedLayoutChildren; - use crate::layouts::chunked::ChunkedLayout; - use crate::layouts::flat::Flat; - use crate::layouts::flat::writer::FlatLayoutStrategy; - use crate::layouts::struct_::StructLayout; - use crate::segments::TestSegments; - use crate::sequence::SequenceId; - use crate::sequence::SequentialArrayStreamExt; - use crate::strategy::LayoutStrategy; - use crate::test::new_session; - - /// Test display_tree with inline array_tree metadata (no segment source needed). - #[test] - fn test_display_tree_inline_array_tree() { - // LazyLock caches the env var on first read, so only nextest (separate processes) can isolate it. - if std::env::var("NEXTEST_RUN_ID").is_ok() { - temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", Some("1"), || { - block_on(|handle| async move { - let session = new_session().with_handle(handle); - let ctx = ArrayContext::empty(); - let segments = Arc::new(TestSegments::default()); - - // Create nullable i64 array (2 buffers: data + validity) - let (ptr1, eof1) = SequenceId::root().split(); - let mut validity_builder = BitBufferMut::with_capacity(5); - for b in [true, false, true, true, false] { - validity_builder.append(b); - } - let validity = Validity::Array( - BoolArray::new(validity_builder.freeze(), Validity::NonNullable) - .into_array(), - ); - let array1 = PrimitiveArray::new(buffer![1i64, 2, 3, 4, 5], validity); - let layout1 = FlatLayoutStrategy::default() - .write_stream( - ctx.clone().into(), - Arc::::clone(&segments), - array1.into_array().to_array_stream().sequenced(ptr1), - eof1, - &session, - ) - .await - .unwrap(); - - // Create utf8 array (2 buffers: views + data) - let (ptr2, eof2) = SequenceId::root().split(); - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(NonNullable), 5); - for s in [ - "hello world this is long", - "another long string", - "short", - "medium str", - "x", - ] { - builder.append_value(s); - } - let layout2 = FlatLayoutStrategy::default() - .write_stream( - ctx.clone().into(), - Arc::::clone(&segments), - builder - .finish() - .into_array() - .to_array_stream() - .sequenced(ptr2), - eof2, - &session, - ) - .await - .unwrap(); - - // Create struct layout - let struct_layout = StructLayout::new( - 5, - DType::Struct( - StructFields::new( - vec![FieldName::from("numbers"), FieldName::from("strings")].into(), - vec![ - DType::Primitive(PType::I64, Nullability::Nullable), - DType::Utf8(NonNullable), - ], - ), - NonNullable, - ), - vec![ - ChunkedLayout::new( - 5, - DType::Primitive(PType::I64, Nullability::Nullable), - OwnedLayoutChildren::layout_children(vec![layout1]), - ) - .into_layout(), - layout2, - ], - ) - .into_layout(); - - let output = format!("{}", struct_layout.display_tree_verbose(true)); - - let expected = "\ -vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5 -├── numbers: vortex.chunked, dtype: i64?, children: 1, rows: 5 -│ └── [0]: vortex.flat, dtype: i64?, metadata: 171 bytes, rows: 5, segment 0, buffers=[40B, 1B], total=41B -└── strings: vortex.flat, dtype: utf8, metadata: 110 bytes, rows: 5, segment 1, buffers=[43B, 80B], total=123B -"; - assert_eq!(output, expected); - }) - }) - } - } - - /// Test display_tree_with_segments using async segment source to fetch buffer sizes. - #[test] - fn test_display_tree_with_segment_source() { - if std::env::var("NEXTEST_RUN_ID").is_ok() { - temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", None::<&str>, || { - block_on(|handle| async move { - let session = new_session().with_handle(handle); - let ctx = ArrayContext::empty(); - let segments = Arc::new(TestSegments::default()); - - // Create simple i32 array - let (ptr1, eof1) = SequenceId::root().split(); - let array1 = - PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::NonNullable); - let layout1 = FlatLayoutStrategy::default() - .write_stream( - ctx.clone().into(), - Arc::::clone(&segments), - array1.into_array().to_array_stream().sequenced(ptr1), - eof1, - &session, - ) - .await - .unwrap(); - - // Create another i32 array - let (ptr2, eof2) = SequenceId::root().split(); - let array2 = - PrimitiveArray::new(buffer![6i32, 7, 8, 9, 10], Validity::NonNullable); - let layout2 = FlatLayoutStrategy::default() - .write_stream( - ctx.clone().into(), - Arc::::clone(&segments), - array2.into_array().to_array_stream().sequenced(ptr2), - eof2, - &session, - ) - .await - .unwrap(); - - // Create chunked layout - let chunked_layout = ChunkedLayout::new( - 10, - DType::Primitive(PType::I32, NonNullable), - OwnedLayoutChildren::layout_children(vec![layout1, layout2]), - ) - .into_layout(); - - let output = chunked_layout - .display_tree_with_segments(segments) - .await - .unwrap(); - - let expected = "\ -vortex.chunked, dtype: i32, children: 2, rows: 10 -├── [0]: vortex.flat, dtype: i32, rows: 5, segment 0, buffers=[20B], total=20B -└── [1]: vortex.flat, dtype: i32, rows: 5, segment 1, buffers=[20B], total=20B -"; - assert_eq!(output.to_string(), expected); - }) - }) - } - } - - /// Test display_array_tree with inline array node metadata. - #[test] - fn test_display_array_tree_with_inline_node() { - if std::env::var("NEXTEST_RUN_ID").is_ok() { - temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", Some("1"), || { - let ctx = ArrayContext::empty(); - let segments = Arc::new(TestSegments::default()); - let (ptr, eof) = SequenceId::root().split(); - - // Create a simple primitive array - let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::AllValid); - let layout = block_on(|handle| async { - let session = new_session().with_handle(handle); - FlatLayoutStrategy::default() - .write_stream( - ctx.clone().into(), - Arc::::clone(&segments), - array.into_array().to_array_stream().sequenced(ptr), - eof, - &session, - ) - .await - .unwrap() - }); - - let flat_layout = layout.as_::(); - - let array_tree = flat_layout.array_tree().expect( - "array_tree should be populated when FLAT_LAYOUT_INLINE_ARRAY_NODE is set", - ); - - let parts = SerializedArray::from_array_tree(array_tree.as_ref().to_vec()) - .expect("should parse array_tree"); - assert_eq!(parts.buffer_lengths(), vec![20]); // 5 i32 values = 20 bytes - - assert_eq!( - layout.display_tree().to_string(), - "\ -vortex.flat, dtype: i32?, segment 0, buffers=[20B], total=20B -" - ); - }) - } - } - - /// Test display_tree without inline array node (shows segment ID). - #[test] - fn test_display_tree_without_inline_node() { - if std::env::var("NEXTEST_RUN_ID").is_ok() { - temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", Some("1"), || { - let ctx = ArrayContext::empty(); - let segments = Arc::new(TestSegments::default()); - let (ptr, eof) = SequenceId::root().split(); - - // Create a simple primitive array - let array = PrimitiveArray::new(buffer![10i64, 20, 30], Validity::NonNullable); - let layout = block_on(|handle| async { - let session = new_session().with_handle(handle); - FlatLayoutStrategy::default() - .write_stream( - ctx.into(), - Arc::::clone(&segments), - array.into_array().to_array_stream().sequenced(ptr), - eof, - &session, - ) - .await - .unwrap() - }); - - // Test display_tree exact output (with inline array_tree enabled by env var from other test) - assert_eq!( - layout.display_tree().to_string(), - "\ -vortex.flat, dtype: i64, segment 0, buffers=[24B], total=24B -" - ); - }) - } - } -} diff --git a/vortex-layout/src/flatbuffers.rs b/vortex-layout/src/flatbuffers.rs index a807db578af..e2009f19008 100644 --- a/vortex-layout/src/flatbuffers.rs +++ b/vortex-layout/src/flatbuffers.rs @@ -4,23 +4,17 @@ use std::env; use std::sync::LazyLock; -use flatbuffers::FlatBufferBuilder; use flatbuffers::VerifierOptions; -use flatbuffers::WIPOffset; use flatbuffers::root_with_opts; use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_flatbuffers::FlatBuffer; -use vortex_flatbuffers::FlatBufferRoot; -use vortex_flatbuffers::WriteFlatBuffer; use vortex_flatbuffers::layout; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; -use crate::DynLayout; use crate::LayoutBuildContext; -use crate::LayoutContext; use crate::LayoutRef; use crate::children::ViewedLayoutChildren; use crate::layouts::foreign::new_foreign_layout; @@ -149,78 +143,6 @@ fn foreign_layout_from_fb( )) } -impl dyn DynLayout + '_ { - /// Serialize the layout into a [`FlatBufferBuilder`]. - pub fn flatbuffer_writer<'a>( - &'a self, - ctx: &'a LayoutContext, - ) -> impl WriteFlatBuffer = layout::Layout<'a>> + FlatBufferRoot + 'a { - LayoutFlatBufferWriter { layout: self, ctx } - } -} - -/// An adapter struct for writing a layout to a FlatBuffer. -struct LayoutFlatBufferWriter<'a> { - layout: &'a dyn DynLayout, - ctx: &'a LayoutContext, -} - -impl FlatBufferRoot for LayoutFlatBufferWriter<'_> {} - -impl WriteFlatBuffer for LayoutFlatBufferWriter<'_> { - type Target<'fb> = layout::Layout<'fb>; - - fn write_flatbuffer<'fb>( - &self, - fbb: &mut FlatBufferBuilder<'fb>, - ) -> VortexResult>> { - // First we recurse into the children and write them out - let child_layouts = self.layout.children()?; - let children = child_layouts - .iter() - .map(|layout| { - LayoutFlatBufferWriter { - layout: layout.as_ref(), - ctx: self.ctx, - } - .write_flatbuffer(fbb) - }) - .collect::>>()?; - let children = (!children.is_empty()).then(|| fbb.create_vector(&children)); - - // Next we write out the metadata if it's non-empty. - let metadata = self.layout.metadata(); - let metadata = (!metadata.is_empty()).then(|| fbb.create_vector(&metadata)); - - let segments = self - .layout - .segment_ids() - .into_iter() - .map(|s| *s) - .collect::>(); - let segments = (!segments.is_empty()).then(|| fbb.create_vector(&segments)); - - // Dictionary-encode the layout ID - let encoding = self.ctx.intern(&self.layout.encoding_id()).ok_or_else(|| { - vortex_err!( - "Layout encoding {} not permitted by ctx", - self.layout.encoding_id() - ) - })?; - - Ok(layout::Layout::create( - fbb, - &layout::LayoutArgs { - encoding, - row_count: self.layout.row_count(), - metadata, - children, - segments, - }, - )) - } -} - #[cfg(test)] mod tests { use flatbuffers::FlatBufferBuilder; diff --git a/vortex-layout/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs index 459429046c2..eafe6cc4953 100644 --- a/vortex-layout/src/layouts/chunked/mod.rs +++ b/vortex-layout/src/layouts/chunked/mod.rs @@ -95,9 +95,17 @@ impl VTable for Chunked { } } -impl Layout { +/// Convenience methods for [`ChunkedLayout`]. +pub trait ChunkedLayoutExt: Sized { /// Construct a chunked layout. - pub fn new(row_count: u64, dtype: DType, children: Arc) -> Self { + fn new(row_count: u64, dtype: DType, children: Arc) -> Self; + + /// Rebuild this layout with owned children. + fn with_children(&self, children: Vec) -> Self; +} + +impl ChunkedLayoutExt for ChunkedLayout { + fn new(row_count: u64, dtype: DType, children: Arc) -> Self { let offsets = chunk_offsets(children.as_ref()).vortex_expect("chunk row counts overflow"); assert_eq!( offsets.last().copied(), @@ -117,8 +125,7 @@ impl Layout { .into_typed() } - /// Rebuild this layout with owned children. - pub fn with_children(&self, children: Vec) -> Self { + fn with_children(&self, children: Vec) -> Self { Self::new( self.row_count(), self.dtype().clone(), diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index b5f4f5b438d..df6ed00616f 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -448,8 +448,10 @@ mod test { use crate::LayoutStrategy; use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; + use crate::layouts::chunked::ChunkedLayoutExt; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::flat::FlatLayout; + use crate::layouts::flat::FlatLayoutExt; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::scan::split_by::SplitBy; use crate::segments::SegmentId; diff --git a/vortex-layout/src/layouts/chunked/writer.rs b/vortex-layout/src/layouts/chunked/writer.rs index 8de7c30fe47..f75917c6c60 100644 --- a/vortex-layout/src/layouts/chunked/writer.rs +++ b/vortex-layout/src/layouts/chunked/writer.rs @@ -18,6 +18,7 @@ use crate::LayoutStrategy; use crate::LayoutWriterContext; use crate::children::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::chunked::ChunkedLayoutExt; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; diff --git a/vortex-layout/src/layouts/dict/mod.rs b/vortex-layout/src/layouts/dict/mod.rs index 84771304690..ace9d0f7015 100644 --- a/vortex-layout/src/layouts/dict/mod.rs +++ b/vortex-layout/src/layouts/dict/mod.rs @@ -125,8 +125,18 @@ impl VTable for Dict { } } -impl Layout { - pub(crate) fn new(values: LayoutRef, codes: LayoutRef) -> Self { +pub(crate) trait DictLayoutExt: Sized { + fn new(values: LayoutRef, codes: LayoutRef) -> Self; + + fn new_with_all_values_referenced( + values: LayoutRef, + codes: LayoutRef, + all_values_referenced: bool, + ) -> Self; +} + +impl DictLayoutExt for DictLayout { + fn new(values: LayoutRef, codes: LayoutRef) -> Self { Self::new_with_all_values_referenced(values, codes, false) } @@ -151,7 +161,9 @@ impl Layout { ) .into_typed() } +} +impl DictData { /// Returns whether every dictionary value is known to be referenced. pub fn has_all_values_referenced(&self) -> bool { self.all_values_referenced diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index b7c93a992fb..a988c469f70 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -42,8 +42,10 @@ use crate::LayoutStrategy; use crate::LayoutWriterContext; use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::chunked::ChunkedLayoutExt; use crate::layouts::compressed::CompressorPlugin; use crate::layouts::dict::DictLayout; +use crate::layouts::dict::DictLayoutExt; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequenceId; diff --git a/vortex-layout/src/layouts/flat/mod.rs b/vortex-layout/src/layouts/flat/mod.rs index 6913d2fc0c4..c111faea5c9 100644 --- a/vortex-layout/src/layouts/flat/mod.rs +++ b/vortex-layout/src/layouts/flat/mod.rs @@ -10,6 +10,7 @@ use std::sync::LazyLock; use vortex_array::ProstMetadata; use vortex_array::dtype::DType; +use vortex_array::serde::SerializedArray; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -106,6 +107,16 @@ impl VTable for Flat { vortex_panic!("Flat layout has no child {idx}") } + fn inlined_segment_buffer_sizes(layout: &Layout) -> Vec<(SegmentId, Vec)> { + layout + .array_tree() + .and_then(|array_tree| { + SerializedArray::from_array_tree(array_tree.as_ref().to_vec()).ok() + }) + .map(|parts| vec![(layout.segment_id(), parts.buffer_lengths())]) + .unwrap_or_default() + } + fn new_reader( layout: &Layout, name: Arc, @@ -122,14 +133,27 @@ impl VTable for Flat { } } -impl Layout { +/// Constructors for [`FlatLayout`]. +pub trait FlatLayoutExt: Sized { /// Construct a flat layout without an inline array encoding tree. - pub fn new(row_count: u64, dtype: DType, segment_id: SegmentId, ctx: ReadContext) -> Self { + fn new(row_count: u64, dtype: DType, segment_id: SegmentId, ctx: ReadContext) -> Self; + + /// Construct a flat layout with optional inline array metadata. + fn new_with_metadata( + row_count: u64, + dtype: DType, + segment_id: SegmentId, + ctx: ReadContext, + array_tree: Option, + ) -> Self; +} + +impl FlatLayoutExt for FlatLayout { + fn new(row_count: u64, dtype: DType, segment_id: SegmentId, ctx: ReadContext) -> Self { Self::new_with_metadata(row_count, dtype, segment_id, ctx, None) } - /// Construct a flat layout with optional inline array metadata. - pub fn new_with_metadata( + fn new_with_metadata( row_count: u64, dtype: DType, segment_id: SegmentId, @@ -150,7 +174,9 @@ impl Layout { ) .into_typed() } +} +impl FlatData { /// Returns the serialized array segment ID. pub fn segment_id(&self) -> SegmentId { self.segment_id diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index 9761c71f9ae..67dae461cdd 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -26,7 +26,9 @@ use crate::LayoutStrategy; use crate::LayoutWriterContext; use crate::children::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::chunked::ChunkedLayoutExt; use crate::layouts::flat::FlatLayout; +use crate::layouts::flat::FlatLayoutExt; use crate::layouts::flat::flat_layout_inline_array_node; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; @@ -196,6 +198,7 @@ mod tests { use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; use vortex_array::expr::stats::StatsProviderExt; + use vortex_array::serde::SerializeOptions; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_buffer::BitBufferMut; @@ -206,17 +209,72 @@ mod tests { use vortex_io::session::RuntimeSessionExt; use vortex_mask::AllOr; use vortex_mask::Mask; + use vortex_session::registry::ReadContext; use vortex_utils::aliases::hash_set::HashSet; use crate::LayoutStrategy; use crate::LayoutStrategyEncodingValidator; + use crate::layouts::flat::FlatLayout; + use crate::layouts::flat::FlatLayoutExt; use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::segments::SegmentId; use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; use crate::test::SESSION; use crate::test::new_session; + #[test] + fn inline_array_tree_display_preserves_buffer_sizes() -> VortexResult<()> { + let session = new_session(); + let ctx = ArrayContext::empty(); + let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable).into_array(); + let buffers = array.serialize(&ctx, &session, &SerializeOptions::default())?; + let array_tree = buffers[buffers.len() - 2].clone(); + let layout = FlatLayout::new_with_metadata( + 3, + array.dtype().clone(), + SegmentId::from(0), + ReadContext::new(ctx.to_ids()), + Some(array_tree), + ); + + assert_eq!( + layout.to_layout().display_tree().to_string(), + "vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B\n" + ); + Ok(()) + } + + #[test] + fn segment_display_fetches_buffer_sizes() { + block_on(|handle| async move { + let session = new_session().with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let (pointer, eof) = SequenceId::root().split(); + let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); + let layout = FlatLayoutStrategy::default() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + array.into_array().to_array_stream().sequenced(pointer), + eof, + &session, + ) + .await + .unwrap(); + + assert_eq!( + layout + .display_tree_with_segments(segments) + .await + .unwrap() + .to_string(), + "vortex.flat, dtype: i32, rows: 3, segment 0, buffers=[12B], total=12B\n" + ); + }) + } + // Currently, flat layouts do not force compute stats during write, they only retain // pre-computed stats. #[should_panic] diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index 0ba5be17dbc..aff194bad37 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -78,7 +78,7 @@ impl VTable for List { args: &LayoutDeserializeArgs<'_>, metadata: &ListLayoutMetadata, ) -> VortexResult { - ListLayout::validate_children(args.dtype, args.children.nchildren())?; + validate_children(args.dtype, args.children.nchildren())?; let elements_dtype = args .dtype .as_list_element_opt() @@ -163,9 +163,31 @@ impl VTable for List { } } -impl Layout { +/// Convenience methods for [`ListLayout`]. +pub trait ListLayoutExt: Sized { /// Construct a list layout from its children. - pub fn new( + fn new( + dtype: DType, + elements: LayoutRef, + offsets: LayoutRef, + validity: Option, + ) -> Self; + + /// Returns the elements child. + fn elements(&self) -> VortexResult; + + /// Returns the offsets child. + fn offsets(&self) -> VortexResult; + + /// Returns the optional validity child. + fn validity(&self) -> VortexResult>; + + /// Returns the list element dtype. + fn elements_dtype(&self) -> &DType; +} + +impl ListLayoutExt for ListLayout { + fn new( dtype: DType, elements: LayoutRef, offsets: LayoutRef, @@ -175,7 +197,7 @@ impl Layout { let offsets_ptype = offsets.dtype().as_ptype(); let mut children = vec![elements, offsets]; children.extend(validity); - Self::validate_children(&dtype, children.len()).vortex_expect("invalid list children"); + validate_children(&dtype, children.len()).vortex_expect("invalid list children"); LayoutParts::new( List, dtype, @@ -188,41 +210,42 @@ impl Layout { } /// Returns the elements child. - pub fn elements(&self) -> VortexResult { + fn elements(&self) -> VortexResult { self.slot(ELEMENTS_CHILD_INDEX)? .ok_or_else(|| vortex_err!("ListLayout elements slot is absent")) } /// Returns the offsets child. - pub fn offsets(&self) -> VortexResult { + fn offsets(&self) -> VortexResult { self.slot(OFFSETS_CHILD_INDEX)? .ok_or_else(|| vortex_err!("ListLayout offsets slot is absent")) } /// Returns the optional validity child. - pub fn validity(&self) -> VortexResult> { + fn validity(&self) -> VortexResult> { self.slot(VALIDITY_CHILD_INDEX) } - /// Returns the integer ptype used by offsets. - pub fn offsets_ptype(&self) -> PType { - self.offsets_ptype - } - - /// Returns the list element dtype. - pub fn elements_dtype(&self) -> &DType { + fn elements_dtype(&self) -> &DType { self.dtype() .as_list_element_opt() .vortex_expect("ListLayout dtype must be a List") } +} - fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> { - let expected = NUM_CHILDREN_NON_NULLABLE + usize::from(dtype.is_nullable()); - vortex_ensure_eq!(nchildren, expected); - Ok(()) +impl ListData { + /// Returns the integer ptype used by offsets. + pub fn offsets_ptype(&self) -> PType { + self.offsets_ptype } } +fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> { + let expected = NUM_CHILDREN_NON_NULLABLE + usize::from(dtype.is_nullable()); + vortex_ensure_eq!(nchildren, expected); + Ok(()) +} + #[derive(prost::Message)] pub struct ListLayoutMetadata { #[prost(enumeration = "PType", tag = "1")] diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 53227635de6..bf93c14653b 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -35,6 +35,7 @@ use crate::LayoutReaderRef; use crate::RowSplits; use crate::SplitRange; use crate::layouts::list::ListLayout; +use crate::layouts::list::ListLayoutExt; use crate::layouts::list::expr::ListChildrenNeeded; use crate::layouts::list::expr::get_necessary_bound_list_children; use crate::layouts::list::expr::rewrite_offsets_expr; @@ -368,7 +369,7 @@ impl LayoutReader for ListReader { let element_row_count = self.elements.row_count(); if element_row_count != 0 { - let mut element_splits = RowSplits::new_capacity(128); + let mut element_splits = RowSplits::with_capacity(128); self.elements.register_splits( &[FieldMask::All], &SplitRange::root(0..element_row_count)?, diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 4d8565fdd10..17c3f5789a7 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -35,6 +35,7 @@ use crate::LayoutStrategy; use crate::LayoutWriterContext; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::list::ListLayout; +use crate::layouts::list::ListLayoutExt; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequenceId; diff --git a/vortex-layout/src/layouts/struct_/mod.rs b/vortex-layout/src/layouts/struct_/mod.rs index 862b789a349..b5ad2b12c47 100644 --- a/vortex-layout/src/layouts/struct_/mod.rs +++ b/vortex-layout/src/layouts/struct_/mod.rs @@ -63,7 +63,7 @@ impl VTable for Struct { args: &LayoutDeserializeArgs<'_>, _metadata: &EmptyMetadata, ) -> VortexResult { - Layout::::validate_children(args.dtype, args.children.nchildren())?; + validate_children(args.dtype, args.children.nchildren())?; for idx in 0..args.children.nchildren() { let child_row_count = args.children.child_row_count(idx); @@ -90,7 +90,7 @@ impl VTable for Struct { } fn child_dtype(layout: &Layout, slot: usize) -> VortexResult { - StructLayout::slot_dtype(layout.dtype(), slot) + slot_dtype(layout.dtype(), slot) } fn child_type(layout: &Layout, slot: usize) -> LayoutChildType { @@ -124,10 +124,23 @@ impl VTable for Struct { } } -impl Layout { +/// Convenience methods for [`StructLayout`]. +pub trait StructLayoutExt: Sized { /// Construct a struct layout from owned children. - pub fn new(row_count: u64, dtype: DType, children: Vec) -> Self { - Self::validate_children(&dtype, children.len()).vortex_expect("invalid struct children"); + fn new(row_count: u64, dtype: DType, children: Vec) -> Self; + + /// Returns the struct fields. + fn struct_fields(&self) -> &StructFields; + + /// Invokes `per_child` for fields selected by `field_mask`. + fn matching_fields(&self, field_mask: &[FieldMask], per_child: F) -> VortexResult<()> + where + F: FnMut(FieldMask, usize) -> VortexResult<()>; +} + +impl StructLayoutExt for StructLayout { + fn new(row_count: u64, dtype: DType, children: Vec) -> Self { + validate_children(&dtype, children.len()).vortex_expect("invalid struct children"); LayoutParts::new( Struct, dtype, @@ -140,14 +153,14 @@ impl Layout { } /// Returns the struct fields. - pub fn struct_fields(&self) -> &StructFields { + fn struct_fields(&self) -> &StructFields { self.dtype() .as_struct_fields_opt() .vortex_expect("Struct layout dtype must be a struct") } /// Invokes `per_child` for fields selected by `field_mask`. - pub fn matching_fields(&self, field_mask: &[FieldMask], mut per_child: F) -> VortexResult<()> + fn matching_fields(&self, field_mask: &[FieldMask], mut per_child: F) -> VortexResult<()> where F: FnMut(FieldMask, usize) -> VortexResult<()>, { @@ -173,30 +186,28 @@ impl Layout { } Ok(()) } +} - fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> { - let fields = dtype - .as_struct_fields_opt() - .ok_or_else(|| vortex_err!("Expected struct dtype"))?; - let expected = fields.nfields() + usize::from(dtype.is_nullable()); - vortex_ensure!( - nchildren == expected, - "Struct layout has {nchildren} children, expected {expected}" - ); - Ok(()) - } +fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> { + let fields = dtype + .as_struct_fields_opt() + .ok_or_else(|| vortex_err!("Expected struct dtype"))?; + let expected = fields.nfields() + usize::from(dtype.is_nullable()); + vortex_ensure!( + nchildren == expected, + "Struct layout has {nchildren} children, expected {expected}" + ); + Ok(()) +} - /// Returns the dtype of the child in logical `slot`: slot 0 is the non-nullable validity - /// bitmap, and slot `i + 1` is struct field `i`. - fn slot_dtype(dtype: &DType, slot: usize) -> VortexResult { - if slot == 0 { - Ok(DType::Bool(Nullability::NonNullable)) - } else { - dtype - .as_struct_fields_opt() - .and_then(|fields| fields.field_by_index(slot - 1)) - .ok_or_else(|| vortex_err!("Missing field {}", slot - 1)) - } +fn slot_dtype(dtype: &DType, slot: usize) -> VortexResult { + if slot == 0 { + Ok(DType::Bool(Nullability::NonNullable)) + } else { + dtype + .as_struct_fields_opt() + .and_then(|fields| fields.field_by_index(slot - 1)) + .ok_or_else(|| vortex_err!("Missing field {}", slot - 1)) } } @@ -208,6 +219,7 @@ mod tests { use super::*; use crate::layouts::flat::FlatLayout; + use crate::layouts::flat::FlatLayoutExt; use crate::segments::SegmentId; fn flat_child(dtype: DType, segment: u32) -> LayoutRef { diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 005a4e47fed..a6258ae2fb4 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -50,6 +50,7 @@ use crate::RowSplits; use crate::SplitRange; use crate::layouts::partitioned::BoundPartitionedExprEval; use crate::layouts::struct_::StructLayout; +use crate::layouts::struct_::StructLayoutExt; use crate::segments::SegmentSource; pub struct StructReader { diff --git a/vortex-layout/src/layouts/struct_/writer.rs b/vortex-layout/src/layouts/struct_/writer.rs index fee71596fee..38e930b1959 100644 --- a/vortex-layout/src/layouts/struct_/writer.rs +++ b/vortex-layout/src/layouts/struct_/writer.rs @@ -42,6 +42,7 @@ use crate::LayoutRef; use crate::LayoutStrategy; use crate::LayoutWriterContext; use crate::layouts::struct_::StructLayout; +use crate::layouts::struct_::StructLayoutExt; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequenceId; diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 1a3c1adc524..ab14974f9ab 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -360,6 +360,7 @@ mod tests { use crate::layouts::repartition::RepartitionWriterOptions; use crate::layouts::table::TableStrategy; use crate::layouts::zoned::Zoned; + use crate::layouts::zoned::ZonedLayoutExt; use crate::layouts::zoned::writer::ZonedLayoutOptions; use crate::layouts::zoned::writer::ZonedStrategy; use crate::segments::TestSegments; diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index bc3d7d0626e..72641f7ce1f 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -166,7 +166,7 @@ impl VTable for Zoned { session: &VortexSession, ctx: &LayoutReaderContext, ) -> VortexResult { - layout.zoned_reader(name, segment_source, session, ctx) + zoned_reader(layout, name, segment_source, session, ctx) } } @@ -234,43 +234,7 @@ impl VTable for LegacyStats { session: &VortexSession, ctx: &LayoutReaderContext, ) -> VortexResult { - layout.zoned_reader(name, segment_source, session, ctx) - } -} - -impl LegacyStatsLayout { - fn zoned_reader( - &self, - name: Arc, - segment_source: Arc, - session: &VortexSession, - ctx: &LayoutReaderContext, - ) -> VortexResult { - if self.zone_len == 0 { - return self - .slot(0)? - .vortex_expect("ZonedLayout always has a data child") - .new_reader(name, segment_source, session, ctx); - } - - Ok(Arc::new(ZonedReader::try_new( - self.clone(), - self.nzones(), - name, - segment_source, - session.clone(), - ctx.clone(), - )?)) - } - - pub fn nzones(&self) -> usize { - usize::try_from(self.children().child_row_count(1)) - .vortex_expect("Invalid number of zones, cannot handle more than usize zones") - } - - /// Returns display names for the zone-map aggregates stored by this layout. - pub fn present_aggregates(&self) -> Arc<[String]> { - present_aggregates(&self.zone_map_schema) + zoned_reader(layout, name, segment_source, session, ctx) } } @@ -280,10 +244,23 @@ pub(crate) enum ZoneMapSchema { AggregateFns(Arc<[AggregateFnRef]>), } -impl ZonedLayout { +/// Constructor for [`ZonedLayout`]. +pub trait ZonedLayoutExt: Sized { /// Create a zoned layout from a data child, a zone-map child, a zone length, and the aggregate /// functions stored in the zone map. - pub fn try_new( + fn try_new( + data: LayoutRef, + zones: LayoutRef, + zone_len: NonZeroUsize, + aggregate_fns: Arc<[AggregateFnRef]>, + ) -> VortexResult; + + /// Returns the number of zones in the layout. + fn nzones(&self) -> usize; +} + +impl ZonedLayoutExt for ZonedLayout { + fn try_new( data: LayoutRef, zones: LayoutRef, zone_len: NonZeroUsize, @@ -314,6 +291,12 @@ impl ZonedLayout { .into_typed()) } + fn nzones(&self) -> usize { + nzones(self) + } +} + +impl ZonedData { pub fn zone_len(&self) -> usize { self.zone_len } @@ -323,41 +306,6 @@ impl ZonedLayout { present_aggregates(&self.zone_map_schema) } - /// Builds a reader for a zoned layout, bypassing [`ZonedReader`] when the zone map is empty - /// (`zone_len == 0`) and reading the data child directly, since there is nothing to prune with. - /// This covers both legacy zero-length zones and layouts whose aggregates the session cannot - /// reconstruct. - fn zoned_reader( - &self, - name: Arc, - segment_source: Arc, - session: &VortexSession, - ctx: &LayoutReaderContext, - ) -> VortexResult { - if self.zone_len == 0 { - return self - .slot(0)? - .vortex_expect("ZonedLayout always has a data child") - .new_reader(name, segment_source, session, ctx); - } - - Ok(Arc::new(ZonedReader::try_new( - self.clone(), - self.nzones(), - name, - segment_source, - session.clone(), - ctx.clone(), - )?)) - } - - pub fn nzones(&self) -> usize { - usize::try_from(self.children().child_row_count(1)) - .vortex_expect("Invalid number of zones, cannot handle more than usize zones") - } -} - -impl ZonedData { fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> { match &self.zone_map_schema { ZoneMapSchema::LegacyStats(stats) => stats @@ -370,6 +318,38 @@ impl ZonedData { } } +fn zoned_reader( + layout: &Layout, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, +) -> VortexResult +where + V: VTable, +{ + if layout.zone_len == 0 { + return layout + .slot(0)? + .vortex_expect("ZonedLayout always has a data child") + .new_reader(name, segment_source, session, ctx); + } + + Ok(Arc::new(ZonedReader::try_new( + layout.clone(), + nzones(layout), + name, + segment_source, + session.clone(), + ctx.clone(), + )?)) +} + +fn nzones>(layout: &Layout) -> usize { + usize::try_from(layout.children().child_row_count(1)) + .vortex_expect("Invalid number of zones, cannot handle more than usize zones") +} + fn present_aggregates(schema: &ZoneMapSchema) -> Arc<[String]> { match schema { ZoneMapSchema::LegacyStats(stats) => stats @@ -510,6 +490,7 @@ mod tests { use crate::LayoutBuildContext; use crate::children::OwnedLayoutChildren; use crate::layouts::flat::FlatLayout; + use crate::layouts::flat::FlatLayoutExt; use crate::segments::SegmentId; fn aggregate_spec(aggregate_fn: AggregateFnRef) -> AggregateSpecProto { diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 4151679e6c3..c6ae508082f 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -37,6 +37,7 @@ use crate::LayoutStrategy; use crate::LayoutWriterContext; use crate::layouts::zoned::AggregateStatsAccumulator; use crate::layouts::zoned::ZonedLayout; +use crate::layouts::zoned::ZonedLayoutExt; use crate::layouts::zoned::aggregate_partials; use crate::layouts::zoned::schema::default_bounded_stat_max_bytes; use crate::segments::SegmentSinkRef; diff --git a/vortex-layout/src/lib.rs b/vortex-layout/src/lib.rs index 0dd7527ba28..3526f21c94a 100644 --- a/vortex-layout/src/lib.rs +++ b/vortex-layout/src/lib.rs @@ -17,30 +17,32 @@ pub mod layouts; pub mod plan; -pub use children::*; -pub use encoding::*; pub use flatbuffers::*; -pub use layout::*; -pub use reader::*; -pub use reader_context::*; -pub use strategy::*; -use vortex_session::registry::Interner; -pub use vtable::*; +pub use vortex_layout_commons::*; pub mod aliases; mod children; -pub mod display; -mod encoding; mod flatbuffers; -mod layout; -mod reader; -mod reader_context; +mod reader { + pub use vortex_layout_commons::LayoutReader; + pub use vortex_layout_commons::RowSplits; + pub use vortex_layout_commons::SplitRange; +} pub mod scan; -pub mod segments; -pub mod sequence; pub mod session; mod strategy; +pub use strategy::*; #[cfg(test)] mod test; -mod vtable; -pub type LayoutContext = Interner; +/// Layout tree display helpers. +pub mod display { + pub use vortex_layout_commons::display::*; +} + +/// Segment access contracts and runtime implementations used by layout readers and writers. +pub mod segments; + +/// Sequence types used to preserve writer ordering. +pub mod sequence { + pub use vortex_layout_commons::sequence::*; +} diff --git a/vortex-layout/src/plan/lower.rs b/vortex-layout/src/plan/lower.rs index e8911734549..4bbb7190036 100644 --- a/vortex-layout/src/plan/lower.rs +++ b/vortex-layout/src/plan/lower.rs @@ -24,6 +24,7 @@ use crate::layouts::list::OFFSETS_CHILD_INDEX; use crate::layouts::list::VALIDITY_CHILD_INDEX; use crate::layouts::struct_::Struct; use crate::layouts::struct_::StructLayout; +use crate::layouts::struct_::StructLayoutExt; use crate::plan::ConcatPlan; use crate::plan::ListPackPlan; use crate::plan::PackPlan; diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index 019ebbaf1ba..9e42e72012b 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -19,12 +19,17 @@ use super::*; use crate::LayoutRef; use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::chunked::ChunkedLayoutExt; use crate::layouts::dict::DictLayout; +use crate::layouts::dict::DictLayoutExt; use crate::layouts::flat::FlatLayout; +use crate::layouts::flat::FlatLayoutExt; use crate::layouts::foreign::new_foreign_layout; use crate::layouts::list::ListLayout; +use crate::layouts::list::ListLayoutExt; use crate::layouts::row_idx::row_idx; use crate::layouts::struct_::StructLayout; +use crate::layouts::struct_::StructLayoutExt; use crate::segments::SegmentId; fn primitive(ptype: PType, nullability: Nullability) -> DType { diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index 6106d4f671d..7920b1a5e3b 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -47,7 +47,7 @@ impl SplitBy { SplitBy::Layout => { // We usually have under 100 splits so reserving upfront saves // us some allocations - let mut row_splits = RowSplits::new_capacity(128); + let mut row_splits = RowSplits::with_capacity(128); row_splits.push(row_range.start); layout_reader.register_splits( field_mask, diff --git a/vortex-layout/src/segments/mod.rs b/vortex-layout/src/segments/mod.rs index a7ecf26de3b..29b7c259e34 100644 --- a/vortex-layout/src/segments/mod.rs +++ b/vortex-layout/src/segments/mod.rs @@ -1,64 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Segment access abstractions for layouts. +//! Segment access contracts and runtime implementations. //! -//! Layouts refer to byte ranges by [`SegmentId`]. A [`SegmentSource`] resolves those ids to buffer -//! handles for readers, while a [`SegmentSink`] assigns ids when writers emit buffers. -//! [`SegmentCache`] implementations can sit in front of sources to reuse segment bytes across -//! scans. +//! Layouts refer to byte ranges by [`SegmentId`]. The source and sink contracts are re-exported +//! from [`vortex_layout_commons`], while this module provides the cache and request-sharing +//! policies used by Vortex file readers. mod cache; mod shared; -mod sink; -mod source; #[cfg(any(test, feature = "_test-harness"))] mod test; -use std::fmt::Display; -use std::ops::Deref; - pub use cache::*; pub use shared::*; -pub use sink::*; -pub use source::*; #[cfg(any(test, feature = "_test-harness"))] pub use test::*; -use vortex_error::VortexError; - -/// Identifier for a single physical segment referenced by a layout. -/// -/// Segment ids are local to a file or segment source. The file footer maps ids to physical offsets; -/// custom storage systems may map them to object-store keys or other random-access locations. -// TODO(ngates): should this be a `[u8]` instead? Allowing for arbitrary segment identifiers? -#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct SegmentId(u32); - -impl From for SegmentId { - fn from(value: u32) -> Self { - Self(value) - } -} - -impl TryFrom for SegmentId { - type Error = VortexError; - - fn try_from(value: usize) -> Result { - Ok(Self::from(u32::try_from(value)?)) - } -} - -impl Deref for SegmentId { - type Target = u32; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} - -impl Display for SegmentId { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "SegmentId({})", self.0) - } -} +pub use vortex_layout_commons::segments::*; diff --git a/vortex-layout/src/strategy.rs b/vortex-layout/src/strategy.rs index 5a0b1025e4a..e2e8d5cae23 100644 --- a/vortex-layout/src/strategy.rs +++ b/vortex-layout/src/strategy.rs @@ -2,14 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; -use std::sync::atomic::AtomicU64; -use std::sync::atomic::Ordering; use async_trait::async_trait; use futures::StreamExt; -use vortex_array::ArrayContext; use vortex_array::ArrayId; -use vortex_array::aggregate_fn::AggregateFnId; use vortex_array::normalize::NormalizeOptions; use vortex_array::normalize::Operation; use vortex_error::VortexResult; @@ -17,188 +13,14 @@ use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt; -/// A shared counter of the bytes that layout strategies are holding but have not yet emitted. -/// -/// Clones share the same counter, so a tracker can be handed to a writer before the write begins -/// and polled while it runs. Strategies report their own retained bytes with -/// [`Self::reserve`], which releases the reservation on drop. -#[derive(Clone, Debug, Default)] -pub struct BufferedBytesTracker(Arc); - -impl BufferedBytesTracker { - /// Creates a tracker with a zeroed counter. - pub fn new() -> Self { - Self::default() - } - - /// Returns the number of bytes currently retained by layout strategies. - pub fn buffered_bytes(&self) -> u64 { - self.0.load(Ordering::Relaxed) - } - - /// Records `bytes` as buffered until the returned reservation is dropped. - pub fn reserve(&self, bytes: u64) -> BufferedBytesReservation { - self.0.fetch_add(bytes, Ordering::Relaxed); - BufferedBytesReservation { - tracker: self.clone(), - bytes, - } - } -} - -/// An outstanding claim on a [`BufferedBytesTracker`], released when dropped. -#[derive(Debug)] -pub struct BufferedBytesReservation { - tracker: BufferedBytesTracker, - bytes: u64, -} - -impl BufferedBytesReservation { - /// Returns the number of bytes held by this reservation. - pub fn bytes(&self) -> u64 { - self.bytes - } -} - -impl Drop for BufferedBytesReservation { - fn drop(&mut self) { - self.tracker.0.fetch_sub(self.bytes, Ordering::Relaxed); - } -} - -/// State shared by every strategy participating in a single layout write. -/// -/// Clones share the [`BufferedBytesTracker`] while retaining the array serialization context. -/// Passing this context through the strategy tree keeps writer-scoped state independent of the -/// strategy instances, which may be shared by multiple leaves or writers. -#[derive(Clone)] -pub struct LayoutWriterContext { - array_ctx: ArrayContext, - allowed_aggregates: Option>>, - buffered_bytes: BufferedBytesTracker, -} - -impl LayoutWriterContext { - /// Creates a context for a layout write with a fresh buffered bytes tracker. - pub fn new(array_ctx: ArrayContext) -> Self { - Self { - array_ctx, - allowed_aggregates: None, - buffered_bytes: BufferedBytesTracker::new(), - } - } - - /// Restrict the aggregate functions this write may record, e.g. in a zone map. - /// - /// A write that would record an aggregate outside `allowed` fails, matching the array and - /// layout contexts: a silently thinner zone map is a file that prunes worse than the - /// caller asked for, with nothing in the output saying so. The id set is a plain set of - /// ids — callers that source it from editions resolve it themselves. - pub fn with_allowed_aggregates(mut self, allowed: HashSet) -> Self { - self.allowed_aggregates = Some(Arc::new(allowed)); - self - } - - /// Returns whether `aggregate` may be recorded by this write. Unrestricted contexts - /// permit every aggregate. - pub fn allows_aggregate(&self, aggregate: &AggregateFnId) -> bool { - self.allowed_aggregates - .as_ref() - .is_none_or(|allowed| allowed.contains(aggregate)) - } - - /// Replaces the buffered bytes tracker, so callers can observe the counter from outside the - /// strategy tree. - pub fn with_buffered_bytes_tracker(mut self, tracker: BufferedBytesTracker) -> Self { - self.buffered_bytes = tracker; - self - } - - /// Returns the array serialization context. - pub fn array_ctx(&self) -> &ArrayContext { - &self.array_ctx - } - - /// Returns the tracker that accounts for bytes retained by layout strategies. - pub fn buffered_bytes_tracker(&self) -> &BufferedBytesTracker { - &self.buffered_bytes - } - - /// Returns the number of bytes currently retained by layout strategies. - pub fn buffered_bytes(&self) -> u64 { - self.buffered_bytes.buffered_bytes() - } - - /// Records `bytes` as retained by this write until the returned reservation is dropped. - pub fn reserve_buffered_bytes(&self, bytes: u64) -> BufferedBytesReservation { - self.buffered_bytes.reserve(bytes) - } -} - -impl From for LayoutWriterContext { - fn from(array_ctx: ArrayContext) -> Self { - Self::new(array_ctx) - } -} - -/// Writes an ordered array stream into a layout tree and segment sink. -/// -/// Layout strategies are writer-side extension points. Strategies may repartition, buffer, -/// collect columns, compute statistics, compress arrays, or delegate to child strategies before -/// finally emitting segments. They must preserve the logical row order represented by the -/// [`SequencePointer`]s in the input stream. -#[async_trait] -pub trait LayoutStrategy: 'static + Send + Sync { - /// Asynchronously process an ordered stream of array chunks, emitting them into a sink and - /// returning the [`Layout`][crate::Layout] instance that can be parsed to retrieve the data - /// from rest. - /// - /// This trait uses the `#[async_trait]` attribute to denote that trait objects of this type - /// can be `Box`ed or `Arc`ed and shared around. Commonly, these strategies are composed to - /// form a operator of operations, each of which modifies the chunk stream in some way before - /// passing the data on to a downstream writer. - /// - /// # Sequencing and EOF - /// - /// The `stream` parameter is a stream of ordered array chunks, each of which is associated - /// with a sequence pointer that indicates its position in the overall array. By passing - /// around these pointers (essentially vector clocks), the writer can support concurrent - /// and parallel processing while maintaining a deterministic order of data in the file. - /// The `ctx` parameter carries both array serialization state and writer-scoped accounting - /// through every child strategy. - /// - /// The `eof` parameter is a guaranteed to be greater than all sequence pointers in the stream. - /// - /// Because child strategies can write to the end-of-file pointer, it is very important that - /// **all strategies must await all children concurrently**. Otherwise it is possible to - /// deadlock if one child is waiting to write to EOF while your strategy is preventing the - /// stream from progressing to completion. - /// - /// # Blocking operations - /// - /// This is an async trait method, which will return a `BoxFuture` that you can await from - /// any runtime. Implementations should avoid directly performing blocking work within the - /// `write_stream`, and should instead spawn it onto an appropriate runtime or threadpool - /// dedicated to such work. - /// - /// Such operations are common, and include things like compression and parsing large blobs - /// of data, or serializing very large messages to flatbuffers. - async fn write_stream( - &self, - ctx: LayoutWriterContext, - segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - eof: SequencePointer, - session: &VortexSession, - ) -> VortexResult; -} - /// A layout strategy wrapper that rejects arrays containing encodings outside an allow-list. /// /// Canonical encodings are always permitted. Every chunk is recursively validated before it is @@ -251,53 +73,3 @@ impl LayoutStrategy for LayoutStrategyEncodingValidator { .await } } - -#[async_trait] -impl LayoutStrategy for Arc { - async fn write_stream( - &self, - ctx: LayoutWriterContext, - segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - eof: SequencePointer, - session: &VortexSession, - ) -> VortexResult { - (**self) - .write_stream(ctx, segment_sink, stream, eof, session) - .await - } -} - -#[cfg(test)] -mod tests { - use crate::strategy::BufferedBytesTracker; - - #[test] - fn reservations_accumulate_and_release() { - let tracker = BufferedBytesTracker::new(); - assert_eq!(tracker.buffered_bytes(), 0); - - let first = tracker.reserve(16); - let second = tracker.reserve(32); - assert_eq!(tracker.buffered_bytes(), 48); - assert_eq!(first.bytes(), 16); - - drop(first); - assert_eq!(tracker.buffered_bytes(), 32); - - drop(second); - assert_eq!(tracker.buffered_bytes(), 0); - } - - #[test] - fn clones_share_the_same_counter() { - let tracker = BufferedBytesTracker::new(); - let observer = tracker.clone(); - - let reservation = tracker.reserve(8); - assert_eq!(observer.buffered_bytes(), 8); - - drop(reservation); - assert_eq!(observer.buffered_bytes(), 0); - } -}