feat[array]: duckdb execute

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

Cargo.lock

@@ -49,7 +49,7 @@ impl<I: UnsignedPType> Take<[I]> for &FixedSizeListVector {
         }
 
         let element_indices = expand_indices(indices, self.list_size());
-        let taken_elements = self.elements().take(element_indices.as_slice());
+        let taken_elements = self.elements().as_ref().take(element_indices.as_slice());
 
         debug_assert_eq!(taken_elements.len(), indices.len() * list_size);
         debug_assert_eq!(taken_validity.len(), indices.len());
@@ -92,7 +92,10 @@ fn take_nullable<I: UnsignedPType>(
     }
 
     let expanded_nullable_indices = expand_nullable_indices(indices, list_size);
-    let taken_elements = fsl.elements().take(expanded_nullable_indices.as_slice());
+    let taken_elements = fsl
+        .elements()
+        .as_ref()
+        .take(expanded_nullable_indices.as_slice());
 
     debug_assert_eq!(taken_elements.len(), indices.len() * list_size);
     debug_assert_eq!(taken_validity.len(), indices.len());

vortex-compute/src/take/vector/fixed_size_list.rs

@@ -27,6 +27,22 @@ mod struct_;
 #[cfg(test)]
 mod tests;
 
+impl<I: UnsignedPType> Take<PVector<I>> for Vector {
+    type Output = Vector;
+
+    fn take(self, indices: &PVector<I>) -> Vector {
+        (&self).take(indices)
+    }
+}
+
+impl<I: UnsignedPType> Take<[I]> for Vector {
+    type Output = Vector;
+
+    fn take(self, indices: &[I]) -> Vector {
+        (&self).take(indices)
+    }
+}
+
 impl<I: UnsignedPType> Take<PVector<I>> for &Vector {
     type Output = Vector;
 
@@ -43,16 +59,24 @@ impl<I: UnsignedPType> Take<[I]> for &Vector {
     }
 }
 
-impl<T> Take<PrimitiveVector> for &T
+impl Take<PrimitiveVector> for &Vector
 where
-    for<'a> &'a T: Take<PVector<u8>, Output = T>,
-    for<'a> &'a T: Take<PVector<u16>, Output = T>,
-    for<'a> &'a T: Take<PVector<u32>, Output = T>,
-    for<'a> &'a T: Take<PVector<u64>, Output = T>,
+    for<'a> &'a Vector: Take<PVector<u8>, Output = Vector>,
+    for<'a> &'a Vector: Take<PVector<u16>, Output = Vector>,
+    for<'a> &'a Vector: Take<PVector<u32>, Output = Vector>,
+    for<'a> &'a Vector: Take<PVector<u64>, Output = Vector>,
 {
-    type Output = T;
+    type Output = Vector;
 
-    fn take(self, indices: &PrimitiveVector) -> T {
+    fn take(self, indices: &PrimitiveVector) -> Vector {
         match_each_unsigned_pvector!(indices, |iv| { self.take(iv) })
     }
 }
+
+impl Take<PrimitiveVector> for Vector {
+    type Output = Vector;
+
+    fn take(self, indices: &PrimitiveVector) -> Vector {
+        (&self).take(indices)
+    }
+}

vortex-compute/src/take/vector/mod.rs

@@ -402,7 +402,7 @@ macro_rules! match_each_primitive_type {
                 let $type = <$crate::duckdb::Double as $crate::duckdb::PrimitiveType>::NATIVE;
                 $body
             }
-            _ => vortex_panic!(
+            _ => vortex::error::vortex_panic!(
                 "Unexpected type for match_each_primitive_type: {:?}",
                 $self.as_type_id()
             ),

vortex-duckdb/src/duckdb/logical_type.rs

@@ -2,10 +2,15 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use itertools::Itertools;
+use vortex::array::ArrayRef;
+use vortex::array::VectorExecutor;
 use vortex::array::arrays::BoolArray;
+use vortex::buffer::BitBuffer;
 use vortex::error::VortexResult;
 use vortex::mask::Mask;
+use vortex::session::VortexSession;
 
+use crate::LogicalType;
 use crate::duckdb::Vector;
 use crate::exporter::ColumnExporter;
 use crate::exporter::all_invalid;
@@ -52,10 +57,55 @@ impl ColumnExporter for BoolExporter {
     }
 }
 
+struct BoolVectorExporter {
+    buffer: BitBuffer,
+    mask: Mask,
+}
+
+pub(crate) fn new_vector_exporter(
+    array: ArrayRef,
+    session: &VortexSession,
+) -> VortexResult<Box<dyn ColumnExporter>> {
+    let vector = array.execute_vector(session)?.into_bool();
+    let (buffer, mask) = vector.into_parts();
+    if mask.all_false() {
+        return Ok(all_invalid::new_exporter(
+            buffer.len(),
+            &LogicalType::bool(),
+        ));
+    }
+    Ok(Box::new(BoolVectorExporter { buffer, mask }))
+}
+
+impl ColumnExporter for BoolVectorExporter {
+    fn export(&self, offset: usize, len: usize, vector: &mut Vector) -> VortexResult<()> {
+        // Set validity if necessary.
+        if unsafe { vector.set_validity(&self.mask, offset, len) } {
+            // All values are null, so no point copying the data.
+            return Ok(());
+        }
+
+        // DuckDB uses byte bools, not bit bools.
+        // maybe we can convert into these from a compressed array sometimes?.
+        unsafe { vector.as_slice_mut(len) }.copy_from_slice(
+            &self
+                .buffer
+                .slice(offset..(offset + len))
+                .iter()
+                .collect_vec(),
+        );
+
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::iter;
 
+    use vortex::VortexSessionDefault;
+    use vortex::array::IntoArray;
+
     use super::*;
     use crate::cpp;
     use crate::duckdb::DataChunk;
@@ -97,6 +147,29 @@ mod tests {
             format!(
                 r#"Chunk - [1 Columns]
 - FLAT BOOLEAN: 65 = [ {}]
+"#,
+                iter::repeat_n("true", 65).join(", ")
+            )
+        );
+    }
+
+    #[test]
+    fn test_bool_vector_long() {
+        let arr = BoolArray::from_iter([true; 128]);
+
+        let mut chunk = DataChunk::new([LogicalType::new(cpp::duckdb_type::DUCKDB_TYPE_BOOLEAN)]);
+
+        new_vector_exporter(arr.into_array(), &VortexSession::default())
+            .unwrap()
+            .export(1, 66, &mut chunk.get_vector(0))
+            .unwrap();
+        chunk.set_len(65);
+
+        assert_eq!(
+            format!("{}", String::try_from(&chunk).unwrap()),
+            format!(
+                r#"Chunk - [1 Columns]
+- FLAT BOOLEAN: 65 = [ {}]
 "#,
                 iter::repeat_n("true", 65).join(", ")
             )

vortex-duckdb/src/exporter/bool.rs

@@ -6,6 +6,7 @@ use std::sync::Arc;
 use parking_lot::Mutex;
 use vortex::array::ArrayRef;
 use vortex::array::Canonical;
+use vortex::vector::Vector as VxVector;
 use vortex_utils::aliases::dash_map::DashMap;
 
 use crate::duckdb::Vector;
@@ -19,6 +20,8 @@ use crate::duckdb::Vector;
 pub struct ConversionCache {
     pub values_cache: DashMap<usize, (ArrayRef, Arc<Mutex<Vector>>)>,
     pub canonical_cache: DashMap<usize, (ArrayRef, Canonical)>,
+    // Use for `USE_VORTEX_OPERATORS` will replace canonical_cache.
+    pub vector_cache: DashMap<usize, (ArrayRef, VxVector)>,
     // A value which must be unique for a given DuckDB operator.
     instance_id: u64,
 }

vortex-duckdb/src/exporter/cache.rs

@@ -4,6 +4,8 @@
 use std::marker::PhantomData;
 
 use num_traits::ToPrimitive;
+use vortex::array::ArrayRef;
+use vortex::array::VectorExecutor;
 use vortex::array::arrays::DecimalArray;
 use vortex::buffer::Buffer;
 use vortex::dtype::BigCast;
@@ -15,6 +17,7 @@ use vortex::error::VortexResult;
 use vortex::error::vortex_bail;
 use vortex::mask::Mask;
 use vortex::scalar::DecimalType;
+use vortex::session::VortexSession;
 
 use crate::duckdb::Vector;
 use crate::duckdb::VectorBuffer;
@@ -59,18 +62,6 @@ pub(crate) fn new_exporter(array: &DecimalArray) -> VortexResult<Box<dyn ColumnE
     })
 }
 
-/// Maps a decimal precision into the small type that can represent it.
-/// see <https://duckdb.org/docs/stable/sql/data_types/numeric.html#fixed-point-decimals>
-pub fn precision_to_duckdb_storage_size(decimal_dtype: &DecimalDType) -> VortexResult<DecimalType> {
-    Ok(match decimal_dtype.precision() {
-        1..=4 => DecimalType::I16,
-        5..=9 => DecimalType::I32,
-        10..=18 => DecimalType::I64,
-        19..=38 => DecimalType::I128,
-        decimal_dtype => vortex_bail!("cannot represent decimal in ducdkb {decimal_dtype}"),
-    })
-}
-
 impl<D: NativeDecimalType, N: NativeDecimalType> ColumnExporter for DecimalExporter<D, N>
 where
     D: ToPrimitive,
@@ -115,6 +106,112 @@ impl<D: NativeDecimalType> ColumnExporter for DecimalZeroCopyExporter<D> {
     }
 }
 
+struct DecimalVectorExporter<D: NativeDecimalType, N: NativeDecimalType> {
+    values: Buffer<D>,
+    mask: Mask,
+    /// The DecimalType of the DuckDB column.
+    dest_value_type: PhantomData<N>,
+}
+
+struct DecimalVectorZeroCopyExporter<D: NativeDecimalType> {
+    len: usize,
+    begin: *const D,
+    shared_buffer: VectorBuffer,
+    mask: Mask,
+}
+
+pub(crate) fn new_vector_exporter(
+    array: ArrayRef,
+    session: &VortexSession,
+) -> VortexResult<Box<dyn ColumnExporter>> {
+    let vector = array.execute_vector(session)?.into_decimal();
+    let decimal_type = vector.decimal_type();
+
+    let dest_values_type =
+        precision_to_duckdb_storage_size(&DecimalDType::new(vector.precision(), vector.scale()))?;
+
+    if decimal_type == dest_values_type {
+        match_each_decimal_value_type!(decimal_type, |D| {
+            let vector = D::downcast(vector);
+            let (_, buffer, mask) = vector.into_parts();
+            return Ok(Box::new(DecimalVectorZeroCopyExporter {
+                len: buffer.len(),
+                begin: buffer.as_ptr(),
+                shared_buffer: VectorBuffer::new(buffer),
+                mask,
+            }));
+        })
+    }
+
+    match_each_decimal_value_type!(decimal_type, |D| {
+        match_each_decimal_value_type!(decimal_type, |N| {
+            let vector = D::downcast(vector);
+            let (_, buffer, mask) = vector.into_parts();
+            Ok(Box::new(DecimalVectorExporter {
+                values: buffer,
+                mask,
+                dest_value_type: PhantomData::<N>,
+            }))
+        })
+    })
+}
+
+impl<D: NativeDecimalType, N: NativeDecimalType> ColumnExporter for DecimalVectorExporter<D, N>
+where
+    D: ToPrimitive,
+    N: BigCast,
+{
+    fn export(&self, offset: usize, len: usize, vector: &mut Vector) -> VortexResult<()> {
+        // Set validity if necessary.
+        if unsafe { vector.set_validity(&self.mask, offset, len) } {
+            // All values are null, so no point copying the data.
+            return Ok(());
+        }
+
+        // Copy the values from the Vortex array to the DuckDB vector.
+        for (src, dst) in self.values[offset..offset + len]
+            .iter()
+            .zip(unsafe { vector.as_slice_mut(len) })
+        {
+            *dst = <N as BigCast>::from(*src).vortex_expect(
+                "We know all decimals with this scale/precision fit into the target bit width",
+            );
+        }
+
+        Ok(())
+    }
+}
+
+impl<D: NativeDecimalType> ColumnExporter for DecimalVectorZeroCopyExporter<D> {
+    fn export(&self, offset: usize, len: usize, vector: &mut Vector) -> VortexResult<()> {
+        if unsafe { vector.set_validity(&self.mask, offset, len) } {
+            // All values are null, so no point copying the data.
+            return Ok(());
+        }
+
+        assert!(self.len >= offset + len);
+
+        let pos = unsafe { self.begin.add(offset) };
+        unsafe { vector.set_vector_buffer(&self.shared_buffer) };
+        // While we are setting a *mut T this is an artifact of the C API, this is in fact const.
+        unsafe { vector.set_data_ptr(pos as *mut D) };
+
+        Ok(())
+    }
+}
+
+/// Maps a decimal precision into the small type that can represent it.
+/// see <https://duckdb.org/docs/stable/sql/data_types/numeric.html#fixed-point-decimals>
+pub fn precision_to_duckdb_storage_size(decimal_dtype: &DecimalDType) -> VortexResult<DecimalType> {
+    Ok(match decimal_dtype.precision() {
+        1..=4 => DecimalType::I16,
+        5..=9 => DecimalType::I32,
+        10..=18 => DecimalType::I64,
+        19..=38 => DecimalType::I128,
+        decimal_dtype => vortex_bail!("cannot represent decimal in ducdkb {decimal_dtype}"),
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use vortex::array::arrays::DecimalArray;