diff --git a/vortex-duckdb/src/exporter/bitpacked.rs b/vortex-duckdb/src/exporter/bitpacked.rs new file mode 100644 index 00000000000..53bbd90cec9 --- /dev/null +++ b/vortex-duckdb/src/exporter/bitpacked.rs @@ -0,0 +1,239 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Exports bit-packed integers by unpacking each exported range straight into the vector. + +use std::mem::MaybeUninit; + +use num_traits::WrappingAdd; +use vortex::array::ArrayRef; +use vortex::array::ExecutionCtx; +use vortex::array::match_each_integer_ptype; +use vortex::buffer::BufferMut; +use vortex::dtype::NativePType; +use vortex::dtype::PhysicalPType; +use vortex::encodings::fastlanes::BitPacked; +use vortex::encodings::fastlanes::BitPackedArrayExt; +use vortex::encodings::fastlanes::BitWidthsView; +use vortex::encodings::fastlanes::FL_CHUNK_SIZE; +use vortex::encodings::fastlanes::FoR; +use vortex::encodings::fastlanes::FoRArrayExt; +use vortex::encodings::fastlanes::FoRArraySlotsExt; +use vortex::encodings::fastlanes::FoRReferences; +use vortex::encodings::fastlanes::unpack_iter::BitPacked as UnpackType; +use vortex::encodings::fastlanes::unpack_iter::BitPackingStrategy; +use vortex::encodings::fastlanes::unpack_iter::UnpackedChunks; +use vortex::error::VortexResult; + +use crate::duckdb::VectorBuffer; +use crate::duckdb::VectorRef; +use crate::exporter::ColumnExporter; +use crate::exporter::validity; + +/// Exports a bit-packed array, possibly under a frame of reference, without decoding it into a +/// buffer first. +/// +/// Decoding a whole array allocates and writes a buffer as large as the array, which DuckDB then +/// reads a vector at a time. Unpacking each vector's rows into the vector itself touches only +/// memory that stays in cache. +struct BitPackedExporter { + /// Keeps the packed buffer alive. + _array: ArrayRef, + packed: *const T::Physical, + packed_len: usize, + bit_width: usize, + /// The position of the array's first row within its first chunk. + offset: usize, + len: usize, + reference: T, +} + +// SAFETY: the exporter only reads the packed buffer, which `_array` keeps alive. +unsafe impl Send for BitPackedExporter {} + +/// An exporter for `array` if it is bit-packed at one bit width without patches, possibly under a +/// frame of reference with one reference. +pub(crate) fn try_new_exporter( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + let (packed_array, reference) = match array.as_opt::() { + Some(for_array) => { + let FoRReferences::Global(reference) = for_array.references() else { + return Ok(None); + }; + (for_array.encoded().clone(), Some(reference.clone())) + } + None => (array.clone(), None), + }; + let Some(bitpacked) = packed_array.as_opt::() else { + return Ok(None); + }; + let BitWidthsView::Global(bit_width) = bitpacked.bit_widths() else { + return Ok(None); + }; + // A frame of reference may pack unsigned offsets of signed values, so the values take the + // type of the exported array, which must have the width of the packed values. + if bitpacked.patches().is_some() + || !array.dtype().is_int() + || array.dtype().as_ptype().byte_width() != bitpacked.dtype().as_ptype().byte_width() + { + return Ok(None); + } + let validity = array.validity()?; + if validity.definitely_all_null() { + return Ok(None); + } + + let len = array.len(); + let offset = usize::from(bitpacked.offset()); + match_each_integer_ptype!(array.dtype().as_ptype(), |T| { + let reference = match reference { + Some(reference) => match reference.as_primitive().typed_value::() { + Some(reference) => reference, + None => return Ok(None), + }, + None => T::default(), + }; + let packed = bitpacked.packed_slice::<::Physical>(); + let exporter = Box::new(BitPackedExporter:: { + packed: packed.as_ptr(), + packed_len: packed.len(), + _array: packed_array.clone(), + bit_width: usize::from(bit_width), + offset, + len, + reference, + }); + validity::new_exporter(validity, len, exporter, ctx).map(Some) + }) +} + +impl ColumnExporter for BitPackedExporter { + fn export( + &self, + offset: usize, + len: usize, + vector: &mut VectorRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + assert!(self.len >= offset + len); + let start = self.offset + offset; + let elems_per_chunk = FL_CHUNK_SIZE * self.bit_width / (8 * size_of::()); + let first_chunk = start / FL_CHUNK_SIZE; + // SAFETY: `packed` and `packed_len` describe the packed buffer, which `_array` keeps alive. + let packed = unsafe { std::slice::from_raw_parts(self.packed, self.packed_len) }; + let num_chunks = (start % FL_CHUNK_SIZE + len).div_ceil(FL_CHUNK_SIZE); + let packed = &packed[first_chunk * elems_per_chunk..][..num_chunks * elems_per_chunk]; + + let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; + let mut chunks = UnpackedChunks::try_new_with_strategy( + BitPackingStrategy, + packed, + self.bit_width, + start % FL_CHUNK_SIZE, + len, + &mut scratch, + )?; + // The vector may still point at the memory of an array exported before, so the values + // are unpacked into a buffer of their own that the vector then shares. + let mut values = BufferMut::::with_capacity(len); + chunks.decode_into(&mut values.spare_capacity_mut()[..len]); + // SAFETY: `decode_into` wrote all `len` values. + unsafe { values.set_len(len) }; + if self.reference != T::default() { + for value in values.iter_mut() { + *value = value.wrapping_add(&self.reference); + } + } + let values = values.freeze(); + let data = values.as_ptr(); + // SAFETY: the vector buffer keeps `values` alive for as long as the vector uses them. + unsafe { + vector.set_vector_buffer(&VectorBuffer::new(values)); + // While we are setting a *mut T this is an artifact of the C API, this is in fact const. + vector.set_data_ptr(data.cast_mut()); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex::array::IntoArray; + use vortex::array::VortexSessionExecute; + use vortex::array::arrays::PrimitiveArray; + use vortex::encodings::fastlanes::BitPackedData; + use vortex::encodings::fastlanes::FoR; + use vortex::encodings::fastlanes::FoRArraySlotsExt; + use vortex::encodings::fastlanes::FoRData; + use vortex::error::VortexResult; + use vortex::scalar::Scalar; + + use super::try_new_exporter; + use crate::SESSION; + use crate::cpp; + use crate::duckdb::DataChunk; + use crate::duckdb::LogicalType; + + fn exported( + array: &vortex::array::ArrayRef, + offset: usize, + len: usize, + ) -> VortexResult> { + let mut ctx = SESSION.create_execution_ctx(); + let mut chunk = DataChunk::new([LogicalType::new(cpp::duckdb_type::DUCKDB_TYPE_BIGINT)]); + try_new_exporter(array, &mut ctx)? + .expect("bit-packed arrays have an exporter") + .export(offset, len, chunk.get_vector_mut(0), &mut ctx)?; + chunk.set_len(len); + Ok(chunk.get_vector(0).as_slice_with_len::(len).to_vec()) + } + + #[rstest] + #[case::aligned(0, 2048)] + #[case::unaligned(1000, 2048)] + #[case::tail(4000, 1000)] + fn exports_bit_packed_ranges(#[case] offset: usize, #[case] len: usize) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter((0..5000i64).map(|i| (i * 7) % 1000)); + let packed = BitPackedData::encode(&values.into_array(), 10, &mut ctx)?.into_array(); + let expected: Vec = (offset as i64..(offset + len) as i64) + .map(|i| (i * 7) % 1000) + .collect(); + assert_eq!(exported(&packed, offset, len)?, expected); + + let sliced = packed.slice(300..5000)?; + let expected: Vec = ((offset + 300) as i64 + ..(offset + 300 + len.min(4700 - offset)) as i64) + .map(|i| (i * 7) % 1000) + .collect(); + assert_eq!(exported(&sliced, offset, expected.len())?, expected); + Ok(()) + } + + #[test] + fn exports_frame_of_reference_of_signed_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..3000i64).map(|i| -500 + (i * 13) % 1000); + let for_array = FoRData::encode(PrimitiveArray::from_iter(values.clone()), &mut ctx)?; + let encoded = for_array.encoded().clone(); + let encoded = BitPackedData::encode(&encoded, 10, &mut ctx)?.into_array(); + let for_array = FoR::try_new(encoded, Scalar::from(-500i64))?.into_array(); + let expected: Vec = values.skip(700).take(2048).collect(); + assert_eq!(exported(&for_array, 700, 2048)?, expected); + Ok(()) + } + + #[test] + fn exports_frame_of_reference() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let encoded = PrimitiveArray::from_iter((0..3000i64).map(|i| i % 500)).into_array(); + let encoded = BitPackedData::encode(&encoded, 9, &mut ctx)?.into_array(); + let for_array = FoR::try_new(encoded, Scalar::from(1_000_000i64))?.into_array(); + let expected: Vec = (100..2148i64).map(|i| 1_000_000 + i % 500).collect(); + assert_eq!(exported(&for_array, 100, 2048)?, expected); + Ok(()) + } +} diff --git a/vortex-duckdb/src/exporter/mod.rs b/vortex-duckdb/src/exporter/mod.rs index bb6cac76970..c5b56b298e6 100644 --- a/vortex-duckdb/src/exporter/mod.rs +++ b/vortex-duckdb/src/exporter/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod all_invalid; +mod bitpacked; mod bool; mod cache; mod canonical; @@ -32,9 +33,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 +224,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 +248,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, @@ -262,6 +285,10 @@ fn new_array_exporter_with_flatten( Err(array) => array, }; + if let Some(exporter) = bitpacked::try_new_exporter(&array, ctx)? { + return Ok(exporter); + } + canonical::new_exporter(array, cache, ctx) } diff --git a/vortex-duckdb/src/exporter/primitive.rs b/vortex-duckdb/src/exporter/primitive.rs index ce67d66df4f..9efa313075a 100644 --- a/vortex-duckdb/src/exporter/primitive.rs +++ b/vortex-duckdb/src/exporter/primitive.rs @@ -64,6 +64,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::*; @@ -115,6 +118,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 67b3e7883c3..9b07c9b6567 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/run_end.rs b/vortex-duckdb/src/exporter/run_end.rs index f252866ef3c..7712073a7f7 100644 --- a/vortex-duckdb/src/exporter/run_end.rs +++ b/vortex-duckdb/src/exporter/run_end.rs @@ -36,6 +36,9 @@ struct RunEndExporter { run_end_offset: usize, } +/// Arrays whose runs are shorter than this on average are exported as flat values. +const MIN_AVERAGE_RUN_LEN: usize = 8; + pub(crate) fn new_exporter_with_flatten( array: RunEndArray, cache: &ConversionCache, @@ -44,7 +47,8 @@ pub(crate) fn new_exporter_with_flatten( ) -> VortexResult> { // Our canonicalization is faster than creating a dictionary vector and // letting duckdb flatten it for us. - if flatten { + // A dictionary vector of short runs costs DuckDB more than flat values. + if flatten || array.len() < MIN_AVERAGE_RUN_LEN * array.ends().len() { return canonical::new_exporter(array.into_array(), cache, ctx); } @@ -143,13 +147,48 @@ mod tests { use crate::exporter::ConversionCache; use crate::exporter::new_array_exporter; + /// Formats `values`, each repeated over a run of eight rows, as DuckDB prints a vector. + fn runs_of_eight(values: &[&str]) -> String { + values + .iter() + .flat_map(|value| [*value; 8]) + .collect::>() + .join(", ") + } + #[test] fn test_one_chunk_null() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let source = vec![Some(0u32), Some(1), None, Some(3), None]; + let source = [Some(0u32), Some(1), None, Some(3), None] + .into_iter() + .flat_map(|value| [value; 8]); let array = PrimitiveArray::from_option_iter(source); let array = RunEnd::encode(array.into_array(), &mut ctx)?; + let mut chunk = DataChunk::new([LogicalType::new(DUCKDB_TYPE_INTEGER)]); + new_array_exporter(array.into_array(), &ConversionCache::default(), &mut ctx)?.export( + 0, + 40, + chunk.get_vector_mut(0), + &mut ctx, + )?; + chunk.set_len(40); + assert_eq!( + String::try_from(&*chunk)?, + format!( + "Chunk - [1 Columns]\n- DICTIONARY INTEGER: 40 = [ {}]\n", + runs_of_eight(&["0", "1", "NULL", "3", "NULL"]) + ) + ); + Ok(()) + } + + #[test] + fn short_runs_export_flat() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = PrimitiveArray::from_option_iter([Some(0u32), Some(1), None, Some(3), None]); + let array = RunEnd::encode(array.into_array(), &mut ctx)?; + let mut chunk = DataChunk::new([LogicalType::new(DUCKDB_TYPE_INTEGER)]); new_array_exporter(array.into_array(), &ConversionCache::default(), &mut ctx)?.export( 0, @@ -158,11 +197,10 @@ mod tests { &mut ctx, )?; chunk.set_len(5); - let chunk_str = String::try_from(&*chunk)?; assert_eq!( - chunk_str, + String::try_from(&*chunk)?, r#"Chunk - [1 Columns] -- DICTIONARY INTEGER: 5 = [ 0, 1, NULL, 3, NULL] +- FLAT INTEGER: 5 = [ 0, 1, NULL, 3, NULL] "# ); Ok(()) @@ -175,7 +213,7 @@ mod tests { 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 field = RunEnd::try_new(buffer![8u32, 16].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()]); @@ -183,9 +221,10 @@ mod tests { assert!(exporter.export(&mut chunk, None)?); assert_eq!( String::try_from(&*chunk)?, - r#"Chunk - [1 Columns] -- DICTIONARY INTEGER: 2 = [ 10, 20] -"# + format!( + "Chunk - [1 Columns]\n- DICTIONARY INTEGER: 16 = [ {}]\n", + runs_of_eight(&["10", "20"]) + ) ); assert!(!exporter.export(&mut chunk, None)?); 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)