DuckDB chunk exporter

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<usize>,
+    chunks: Vec<Box<dyn ColumnExporter>>,
+}
+
+pub(crate) fn new_exporter_with_flatten(
+    array: ChunkedArray,
+    cache: &ConversionCache,
+    ctx: &mut ExecutionCtx,
+    flatten: bool,
+) -> VortexResult<Box<dyn ColumnExporter>> {
+    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::<VortexResult<Vec<_>>>()?;
+
+    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(())
+    }
+}

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<I: IntegerPType> {
     // 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(),

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::<Chunked>() {
+        Ok(array) => return chunked::new_exporter_with_flatten(array, cache, ctx, flatten),
+        Err(array) => array,
+    };
+
     let array = match array.try_downcast::<List>() {
         Ok(array) => return list::new_exporter(array, cache, ctx),
         Err(array) => array,

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::<PrimitiveArray>(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<E: IntegerPType> ColumnExporter for RunEndExporter<E> {
         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(())
+    }
+}

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,

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,