diff --git a/vortex-duckdb/src/exporter/chunked.rs b/vortex-duckdb/src/exporter/chunked.rs new file mode 100644 index 00000000000..6e6b6642030 --- /dev/null +++ b/vortex-duckdb/src/exporter/chunked.rs @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex::array::ExecutionCtx; +use vortex::array::IntoArray; +use vortex::array::arrays::ChunkedArray; +use vortex::array::arrays::chunked::ChunkedArrayExt; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; + +use crate::duckdb::VectorRef; +use crate::exporter::ColumnExporter; +use crate::exporter::ConversionCache; +use crate::exporter::canonical; +use crate::exporter::new_array_exporter; + +struct ChunkedExporter { + chunk_offsets: Vec, + chunks: Vec>, +} + +pub(crate) fn new_exporter_with_flatten( + array: ChunkedArray, + cache: &ConversionCache, + ctx: &mut ExecutionCtx, + flatten: bool, +) -> VortexResult> { + if flatten { + return canonical::new_exporter(array.into_array(), cache, ctx); + } + + let chunk_offsets = array.chunk_offsets().to_vec(); + let chunks = array + .chunks() + .iter() + .map(|chunk| new_array_exporter(chunk.clone(), cache, ctx)) + .collect::>>()?; + + Ok(Box::new(ChunkedExporter { + chunk_offsets, + chunks, + })) +} + +impl ChunkedExporter { + fn chunk_index(&self, offset: usize) -> usize { + self.chunk_offsets + .partition_point(|&chunk_offset| chunk_offset <= offset) + .saturating_sub(1) + } +} + +impl ColumnExporter for ChunkedExporter { + fn preferred_batch_len(&self, offset: usize, max_len: usize) -> usize { + if max_len == 0 || self.chunks.is_empty() { + return 0; + } + + let chunk_idx = self.chunk_index(offset); + let chunk_start = self.chunk_offsets[chunk_idx]; + let chunk_end = self.chunk_offsets[chunk_idx + 1]; + let len = (chunk_end - offset).min(max_len); + self.chunks[chunk_idx].preferred_batch_len(offset - chunk_start, len) + } + + fn export( + &self, + offset: usize, + len: usize, + vector: &mut VectorRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if len == 0 { + return Ok(()); + } + + let chunk_idx = self.chunk_index(offset); + let chunk_start = self.chunk_offsets[chunk_idx]; + let chunk_end = self.chunk_offsets[chunk_idx + 1]; + let offset_in_chunk = offset - chunk_start; + vortex_ensure!( + offset + len <= chunk_end, + "chunked DuckDB export range {offset}..{} crosses chunk boundary at {chunk_end}", + offset + len + ); + + self.chunks[chunk_idx].export(offset_in_chunk, len, vector, ctx) + } +} + +#[cfg(test)] +mod tests { + use vortex::array::IntoArray; + use vortex::array::VortexSessionExecute; + use vortex::array::arrays::ChunkedArray; + use vortex::array::arrays::DictArray; + use vortex::array::arrays::StructArray; + use vortex::array::arrays::VarBinViewArray; + use vortex::buffer::buffer; + use vortex::error::VortexResult; + + use crate::SESSION; + use crate::duckdb::DataChunk; + use crate::duckdb::LogicalType; + use crate::exporter::ArrayExporter; + use crate::exporter::ConversionCache; + + #[test] + fn chunked_exporter_emits_chunk_aligned_vectors() -> VortexResult<()> { + let values0 = VarBinViewArray::from_iter_str(["a", "b"]).into_array(); + let chunk0 = DictArray::try_new(buffer![0u8, 1].into_array(), values0)?.into_array(); + let dtype = chunk0.dtype().clone(); + + let values1 = VarBinViewArray::from_iter_str(["c", "d", "e"]).into_array(); + let chunk1 = DictArray::try_new(buffer![0u8, 1, 2].into_array(), values1)?.into_array(); + + let field = ChunkedArray::try_new(vec![chunk0, chunk1], dtype)?.into_array(); + let array = StructArray::from_fields(&[("field", field)])?; + let mut exporter = ArrayExporter::try_new( + &array, + &ConversionCache::default(), + SESSION.create_execution_ctx(), + )?; + let mut chunk = DataChunk::new([LogicalType::varchar()]); + + assert!(exporter.export(&mut chunk, None, None)?); + assert_eq!( + format!("{}", String::try_from(&*chunk)?), + r#"Chunk - [1 Columns] +- DICTIONARY VARCHAR: 2 = [ a, b] +"# + ); + + assert!(exporter.export(&mut chunk, None, None)?); + assert_eq!( + format!("{}", String::try_from(&*chunk)?), + r#"Chunk - [1 Columns] +- DICTIONARY VARCHAR: 3 = [ c, d, e] +"# + ); + + assert!(!exporter.export(&mut chunk, None, None)?); + Ok(()) + } +} diff --git a/vortex-duckdb/src/exporter/dict.rs b/vortex-duckdb/src/exporter/dict.rs index cba2f85591f..45aa724dd86 100644 --- a/vortex-duckdb/src/exporter/dict.rs +++ b/vortex-duckdb/src/exporter/dict.rs @@ -25,6 +25,7 @@ use crate::exporter::all_invalid; use crate::exporter::cache::ConversionCache; use crate::exporter::constant; use crate::exporter::new_array_exporter; +use crate::exporter::new_array_exporter_with_flatten; struct DictExporter { // Store the dictionary values once and export the same dictionary with each codes chunk. @@ -104,7 +105,7 @@ pub(crate) fn new_exporter_with_flatten( None => { // Create a new reusable dictionary for the values. let mut reusable_dict = ReusableDict::new(values.dtype().try_into()?, values.len()); - new_array_exporter(values.clone(), cache, ctx)?.export( + new_array_exporter_with_flatten(values.clone(), cache, ctx, true)?.export( 0, values.len(), reusable_dict.vector(), diff --git a/vortex-duckdb/src/exporter/mod.rs b/vortex-duckdb/src/exporter/mod.rs index 7282594ec45..1dad7dc9cda 100644 --- a/vortex-duckdb/src/exporter/mod.rs +++ b/vortex-duckdb/src/exporter/mod.rs @@ -5,6 +5,7 @@ mod all_invalid; mod bool; mod cache; mod canonical; +mod chunked; mod constant; mod decimal; mod dict; @@ -26,6 +27,7 @@ pub use cache::ConversionCache; pub use decimal::precision_to_duckdb_storage_size; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; +use vortex::array::arrays::Chunked; use vortex::array::arrays::Constant; use vortex::array::arrays::Dict; use vortex::array::arrays::List; @@ -37,6 +39,7 @@ use vortex::encodings::sequence::Sequence; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; use crate::duckdb::DataChunkRef; use crate::duckdb::VectorRef; @@ -96,8 +99,18 @@ impl ArrayExporter { vortex_bail!("Expected {expected_cols} columns in output chunk, got {chunk_cols}"); } - let chunk_len = duckdb_vector_size().min(self.remaining); let position = self.array_len - self.remaining; + let mut chunk_len = duckdb_vector_size().min(self.remaining); + if !zero_projection { + for field in &self.fields { + chunk_len = field.preferred_batch_len(position, chunk_len); + } + } + vortex_ensure!( + chunk_len > 0, + "column exporter returned zero rows for non-empty export" + ); + self.remaining -= chunk_len; chunk.set_len(chunk_len); @@ -156,6 +169,14 @@ impl ArrayExporter { /// the offset, len and `WritableVector` as options. Not sure what it should return though? /// This would allow Vortex extension authors to plug into the DuckDB exporter system. pub trait ColumnExporter: 'static { + /// Preferred number of rows to export next, capped by `max_len`. + /// + /// Exporters that preserve physical structure can use this to keep a DuckDB vector inside a + /// natural boundary, such as a chunked-array child boundary. + fn preferred_batch_len(&self, _offset: usize, max_len: usize) -> usize { + max_len + } + /// Export the given range of data from the Vortex array to the DuckDB vector. fn export( &self, @@ -201,6 +222,11 @@ fn new_array_exporter_with_flatten( Err(array) => array, }; + let array = match array.try_downcast::() { + Ok(array) => return chunked::new_exporter_with_flatten(array, cache, ctx, flatten), + Err(array) => array, + }; + let array = match array.try_downcast::() { Ok(array) => return list::new_exporter(array, cache, ctx), Err(array) => array, diff --git a/vortex-duckdb/src/exporter/run_end.rs b/vortex-duckdb/src/exporter/run_end.rs index de3c08c4759..a56ee8879cb 100644 --- a/vortex-duckdb/src/exporter/run_end.rs +++ b/vortex-duckdb/src/exporter/run_end.rs @@ -22,7 +22,7 @@ use crate::duckdb::VectorRef; use crate::exporter::ColumnExporter; use crate::exporter::cache::ConversionCache; use crate::exporter::canonical; -use crate::exporter::new_array_exporter; +use crate::exporter::new_array_exporter_with_flatten; /// We export run-end arrays to a DuckDB dictionary vector, using a selection vector to /// repeat the values in the run-end array. @@ -50,7 +50,9 @@ pub(crate) fn new_exporter_with_flatten( let ends = array.ends().clone(); let values = array.values().clone(); let ends = ends.execute::(ctx)?; - let values_exporter = new_array_exporter(values.clone(), cache, ctx)?; + // REE exports values in run-index space, not outer row space. Materialize the dictionary + // payload so chunked physical boundaries in the values child cannot constrain row batches. + let values_exporter = new_array_exporter_with_flatten(values.clone(), cache, ctx, true)?; match_each_integer_ptype!(ends.ptype(), |E| { Ok(Box::new(RunEndExporter { @@ -137,3 +139,45 @@ impl ColumnExporter for RunEndExporter { Ok(()) } } + +#[cfg(test)] +mod tests { + use vortex::array::IntoArray; + use vortex::array::VortexSessionExecute; + use vortex::array::arrays::ChunkedArray; + use vortex::array::arrays::PrimitiveArray; + use vortex::array::arrays::StructArray; + use vortex::buffer::buffer; + use vortex::encodings::runend::RunEnd; + use vortex::error::VortexResult; + + use crate::SESSION; + use crate::duckdb::DataChunk; + use crate::duckdb::LogicalType; + use crate::exporter::ArrayExporter; + use crate::exporter::ConversionCache; + + #[test] + fn run_end_with_chunked_values_exports_across_value_chunks() -> VortexResult<()> { + let values0 = PrimitiveArray::from_iter([10i32]).into_array(); + let dtype = values0.dtype().clone(); + let values1 = PrimitiveArray::from_iter([20i32]).into_array(); + let values = ChunkedArray::try_new(vec![values0, values1], dtype)?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let field = RunEnd::try_new(buffer![1u32, 2].into_array(), values, &mut ctx)?.into_array(); + let array = StructArray::from_fields(&[("field", field)])?; + let mut exporter = ArrayExporter::try_new(&array, &ConversionCache::default(), ctx)?; + let mut chunk = DataChunk::new([LogicalType::int32()]); + + assert!(exporter.export(&mut chunk, None, None)?); + assert_eq!( + format!("{}", String::try_from(&*chunk)?), + r#"Chunk - [1 Columns] +- DICTIONARY INTEGER: 2 = [ 10, 20] +"# + ); + + assert!(!exporter.export(&mut chunk, None, None)?); + Ok(()) + } +} diff --git a/vortex-duckdb/src/exporter/struct_.rs b/vortex-duckdb/src/exporter/struct_.rs index e9aca03a3bb..dbe65dd90d0 100644 --- a/vortex-duckdb/src/exporter/struct_.rs +++ b/vortex-duckdb/src/exporter/struct_.rs @@ -61,6 +61,12 @@ pub(crate) fn new_exporter( } impl ColumnExporter for StructExporter { + fn preferred_batch_len(&self, offset: usize, max_len: usize) -> usize { + self.children + .iter() + .fold(max_len, |len, child| child.preferred_batch_len(offset, len)) + } + fn export( &self, offset: usize, diff --git a/vortex-duckdb/src/exporter/validity.rs b/vortex-duckdb/src/exporter/validity.rs index 429699e9f55..13774b33f27 100644 --- a/vortex-duckdb/src/exporter/validity.rs +++ b/vortex-duckdb/src/exporter/validity.rs @@ -70,6 +70,10 @@ pub(crate) fn new_exporter( } impl ColumnExporter for ValidityExporter { + fn preferred_batch_len(&self, offset: usize, max_len: usize) -> usize { + self.exporter.preferred_batch_len(offset, max_len) + } + fn export( &self, offset: usize,