diff --git a/vortex-duckdb/src/exporter/mod.rs b/vortex-duckdb/src/exporter/mod.rs index 539a6d994a2..0714f6e8f4f 100644 --- a/vortex-duckdb/src/exporter/mod.rs +++ b/vortex-duckdb/src/exporter/mod.rs @@ -32,9 +32,12 @@ use vortex::array::ExecutionCtx; use vortex::array::arrays::Chunked; use vortex::array::arrays::Constant; use vortex::array::arrays::Dict; +use vortex::array::arrays::Filter; use vortex::array::arrays::List; +use vortex::array::arrays::Slice; use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; +use vortex::array::matcher::Matcher; use vortex::buffer::BitChunks; use vortex::encodings::fastlanes::RLE; use vortex::encodings::runend::RunEnd; @@ -220,6 +223,23 @@ fn cached_values_dict( Ok(dict) } +/// Matches arrays that are not lazy filters or slices. A scan leaves these over its columns, and +/// executing them lets the encodings below them, run-end for example, be exported as they are +/// instead of decoded. +struct NotSliceOrFilter; + +impl Matcher for NotSliceOrFilter { + type Match<'a> = &'a ArrayRef; + + fn matches(array: &ArrayRef) -> bool { + !array.is::() && !array.is::() + } + + fn try_match(array: &ArrayRef) -> Option> { + Self::matches(array).then_some(array) + } +} + /// Create a DuckDB exporter for the given Vortex array. fn new_array_exporter_with_flatten( array: ArrayRef, @@ -227,6 +247,8 @@ fn new_array_exporter_with_flatten( ctx: &mut ExecutionCtx, flatten: bool, ) -> VortexResult> { + let array = array.execute_until::(ctx)?; + let array = match array.try_downcast::() { Ok(array) => return constant::new_exporter_with_flatten(array, cache, ctx, flatten), Err(array) => array, diff --git a/vortex-duckdb/src/exporter/primitive.rs b/vortex-duckdb/src/exporter/primitive.rs index 515ed3bbe24..d023b4616f1 100644 --- a/vortex-duckdb/src/exporter/primitive.rs +++ b/vortex-duckdb/src/exporter/primitive.rs @@ -62,6 +62,9 @@ impl ColumnExporter for PrimitiveExporter { #[cfg(test)] mod tests { + use vortex::array::validity::Validity; + use vortex::buffer::BitBuffer; + use vortex::buffer::buffer; use vortex_array::VortexSessionExecute; use super::*; @@ -113,6 +116,31 @@ mod tests { ); } + /// A validity bitmap that starts inside its first byte cannot be shared with DuckDB, so its + /// bits are copied from the bitmap's own offset. + #[test] + fn test_primitive_exporter_with_offset_validity() -> VortexResult<()> { + let validity = [true, true, true, false, true, false, false, true] + .into_iter() + .collect::() + .slice(3..); + let arr = PrimitiveArray::new(buffer![3i32, 4, 5, 6, 7], Validity::from(validity)); + + let mut chunk = DataChunk::new([LogicalType::new(cpp::duckdb_type::DUCKDB_TYPE_INTEGER)]); + let mut ctx = SESSION.create_execution_ctx(); + + new_exporter(arr, &mut ctx)?.export(0, 5, chunk.get_vector_mut(0), &mut ctx)?; + chunk.set_len(5); + + assert_eq!( + String::try_from(&*chunk)?, + r#"Chunk - [1 Columns] +- FLAT INTEGER: 5 = [ NULL, 4, NULL, NULL, 7] +"# + ); + Ok(()) + } + /// Export a large nullable primitive array over many chunks to exercise the /// zero-copy value and validity paths. #[test] diff --git a/vortex-duckdb/src/exporter/rle.rs b/vortex-duckdb/src/exporter/rle.rs index 483c79f289c..fe3fbf64a59 100644 --- a/vortex-duckdb/src/exporter/rle.rs +++ b/vortex-duckdb/src/exporter/rle.rs @@ -254,7 +254,7 @@ mod tests { assert_eq!( chunk_str, r#"Chunk - [1 Columns] -- FLAT INTEGER: 10 = [ 20, 20, 20, 20, 20, 20, 20, 20, 20, 20] +- DICTIONARY INTEGER: 10 = [ 20, 20, 20, 20, 20, 20, 20, 20, 20, 20] "# ); Ok(()) diff --git a/vortex-duckdb/src/exporter/vector.rs b/vortex-duckdb/src/exporter/vector.rs index 1e457efed51..5cf4c7adad3 100644 --- a/vortex-duckdb/src/exporter/vector.rs +++ b/vortex-duckdb/src/exporter/vector.rs @@ -60,11 +60,11 @@ impl VectorRef { unsafe { self.set_validity_data(u64_offset, len, validity_data) }; } else { debug_assert!(zero_copy.is_none() || !offset.is_multiple_of(64)); - let source = values.bit_buffer().inner().as_slice(); + let bit_buffer = values.bit_buffer(); copy_from_slice( unsafe { self.ensure_validity_slice(len) }, - source, - offset, + bit_buffer.inner().as_slice(), + bit_buffer.offset() + offset, len, ); } diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 38a9bf0218d..2469aaa3dff 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -145,6 +145,7 @@ impl LayoutReader for FlatReader { // to evaluating the expression. let mut array = array.clone().await?; let mask = mask.await?; + let mut ctx = session.create_execution_ctx(); // Slice the array based on the row mask. if row_range.start > 0 || row_range.end < array.len() { @@ -158,14 +159,12 @@ impl LayoutReader for FlatReader { // after this. let array = array.apply_bound(&expr)?; let array = array.filter(mask.clone())?; - let mut ctx = session.create_execution_ctx(); let array_mask = array.fill_null(false)?.execute::(&mut ctx)?; mask.intersect_by_rank(&array_mask) } else { // Run over the full array, with a simpler bitand at the end. let array = array.apply_bound(&expr)?; - let mut ctx = session.create_execution_ctx(); let array_mask = array.fill_null(false)?.execute::(&mut ctx)?; mask.bitand(&array_mask)