noarrows

Signed-off-by: Robert Kruszewski <[email protected]>

Author: robert3005

vortex-duckdb/src/convert/vector.rs

@@ -4,34 +4,22 @@
 /// Copied of duckdb-rs (https://github.com/duckdb/duckdb-rs/blob/main/crates/duckdb/src/vtab/arrow.rs)
 use std::sync::Arc;
 
-use arrow_array::builder::GenericBinaryBuilder;
-use arrow_array::types::{
-    Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type,
-    UInt32Type, UInt64Type,
-};
-use arrow_array::{
-    Array, BooleanArray, Date32Array, Decimal128Array, FixedSizeListArray, GenericListViewArray,
-    PrimitiveArray, StringArray, Time64MicrosecondArray, Time64NanosecondArray,
-    TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
-    TimestampSecondArray,
-};
-use arrow_buffer::buffer::BooleanBuffer;
-use arrow_schema::Field;
 use num_traits::AsPrimitive;
-use vortex::ArrayRef;
-use vortex::arrays::StructArray;
+use vortex::arrays::{PrimitiveArray, StructArray, TemporalArray};
 use vortex::arrow::FromArrowArray;
-use vortex::buffer::BufferMut;
+use vortex::buffer::{Buffer, BufferMut};
+use vortex::dtype::datetime::TimeUnit;
 use vortex::dtype::{DType, DecimalDType, FieldNames, Nullability};
 use vortex::error::{VortexExpect, VortexResult, vortex_err};
 use vortex::scalar::DecimalType;
 use vortex::validity::Validity;
+use vortex::{ArrayRef, IntoArray};
 
 use crate::convert::dtype::FromLogicalType;
 use crate::cpp::{
     DUCKDB_TYPE, duckdb_date, duckdb_list_entry, duckdb_string_t, duckdb_string_t_data,
     duckdb_string_t_length, duckdb_time, duckdb_time_ns, duckdb_timestamp, duckdb_timestamp_ms,
-    duckdb_timestamp_s,
+    duckdb_timestamp_ns, duckdb_timestamp_s,
 };
 use crate::duckdb::{DataChunk, Vector};
 use crate::exporter::precision_to_duckdb_storage_size;
@@ -62,75 +50,71 @@ impl<'a> DuckString<'a> {
     }
 }
 
-// FIXME: flat vectors don't have all of thsese types. I think they only
-/// Converts flat vector to an arrow array
-pub fn flat_vector_to_arrow_array(
-    vector: &mut Vector,
-    len: usize,
-) -> Result<Arc<dyn Array>, Box<dyn std::error::Error>> {
+/// Converts flat vector to a vortex array
+pub fn flat_vector_to_arrow_array(vector: &mut Vector, len: usize) -> VortexResult<ArrayRef> {
     let type_id = vector.logical_type().as_type_id();
     match type_id {
         DUCKDB_TYPE::DUCKDB_TYPE_INTEGER => {
             let data = vector.as_slice_with_len::<i32>(len);
 
-            Ok(Arc::new(
-                PrimitiveArray::<Int32Type>::from_iter_values_with_nulls(
-                    data.iter().copied(),
-                    vector.validity_ref(data.len()).to_null_buffer(),
-                ),
-            ))
+            Ok(PrimitiveArray::new(
+                Buffer::<i32>::copy_from(data),
+                vector.validity_ref(data.len()).to_validity(),
+            )
+            .into_array())
         }
         DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP => {
             let data = vector.as_slice_with_len::<duckdb_timestamp>(len);
             let micros = data.iter().map(|duckdb_timestamp { micros }| *micros);
-            let structs = TimestampMicrosecondArray::from_iter_values_with_nulls(
-                micros,
-                vector.validity_ref(data.len()).to_null_buffer(),
-            );
-
-            Ok(Arc::new(structs))
+            let arr = PrimitiveArray::new(
+                Buffer::from_trusted_len_iter(micros),
+                vector.validity_ref(data.len()).to_validity(),
+            )
+            .into_array();
+            Ok(TemporalArray::new_timestamp(arr, TimeUnit::Microseconds, None).into_array())
         }
         DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_S => {
             let data = vector.as_slice_with_len::<duckdb_timestamp_s>(len);
             let seconds = data.iter().map(|duckdb_timestamp_s { seconds }| *seconds);
-            let structs = TimestampSecondArray::from_iter_values_with_nulls(
-                seconds,
-                vector.validity_ref(data.len()).to_null_buffer(),
-            );
-
-            Ok(Arc::new(structs))
+            let arr = PrimitiveArray::new(
+                Buffer::from_trusted_len_iter(seconds),
+                vector.validity_ref(data.len()).to_validity(),
+            )
+            .into_array();
+            Ok(TemporalArray::new_timestamp(arr, TimeUnit::Seconds, None).into_array())
         }
         DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_MS => {
             let data = vector.as_slice_with_len::<duckdb_timestamp_ms>(len);
             let millis = data.iter().map(|duckdb_timestamp_ms { millis }| *millis);
-            let structs = TimestampMillisecondArray::from_iter_values_with_nulls(
-                millis,
-                vector.validity_ref(data.len()).to_null_buffer(),
-            );
-
-            Ok(Arc::new(structs))
+            let arr = PrimitiveArray::new(
+                Buffer::from_trusted_len_iter(millis),
+                vector.validity_ref(data.len()).to_validity(),
+            )
+            .into_array();
+            Ok(TemporalArray::new_timestamp(arr, TimeUnit::Milliseconds, None).into_array())
         }
         DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_NS => {
-            let data = vector.as_slice_with_len::<duckdb_timestamp>(len);
-            let nanos = data
-                .iter()
-                .map(|duckdb_timestamp { micros }| *micros * 1000);
-            let structs = TimestampNanosecondArray::from_iter_values_with_nulls(
-                nanos,
-                vector.validity_ref(data.len()).to_null_buffer(),
-            );
-
-            Ok(Arc::new(structs))
+            let data = vector.as_slice_with_len::<duckdb_timestamp_ns>(len);
+            let nanos = data.iter().map(|duckdb_timestamp_ns { nanos }| *nanos);
+            let arr = PrimitiveArray::new(
+                Buffer::from_trusted_len_iter(nanos),
+                vector.validity_ref(data.len()).to_validity(),
+            )
+            .into_array();
+            Ok(TemporalArray::new_timestamp(arr, TimeUnit::Nanoseconds, None).into_array())
         }
         DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_TZ => {
             let data = vector.as_slice_with_len::<duckdb_timestamp>(len);
-            let structs = TimestampMicrosecondArray::from_iter_values_with_nulls(
-                data.iter().map(|duckdb_timestamp { micros }| *micros),
-                vector.validity_ref(data.len()).to_null_buffer(),
+            let micros = data.iter().map(|duckdb_timestamp { micros }| *micros);
+            let arr = PrimitiveArray::new(
+                Buffer::from_trusted_len_iter(micros),
+                vector.validity_ref(data.len()).to_validity(),
+            )
+            .into_array();
+            Ok(
+                TemporalArray::new_timestamp(arr, TimeUnit::Nanoseconds, Some("UTC".to_string()))
+                    .into_array(),
             )
-            .with_timezone("UTC");
-
-            Ok(Arc::new(structs))
         }
         DUCKDB_TYPE::DUCKDB_TYPE_VARCHAR => {
             let data = vector.as_slice_with_len::<duckdb_string_t>(len);

vortex-duckdb/src/duckdb/vector.rs

@@ -4,12 +4,12 @@
 use std::ffi::{CStr, CString, c_void};
 use std::ptr;
 
-use arrow_buffer;
-use arrow_buffer::Buffer;
 use bitvec::macros::internal::funty::Fundamental;
 use bitvec::slice::BitSlice;
 use bitvec::view::BitView;
+use vortex::buffer::{BitBuffer, Buffer};
 use vortex::error::{VortexResult, VortexUnwrap, vortex_bail, vortex_err};
+use vortex::validity::Validity;
 
 use crate::cpp::{duckdb_vx_error, idx_t};
 use crate::duckdb::vector_buffer::VectorBuffer;
@@ -302,20 +302,17 @@ impl ValidityRef<'_> {
         (validity_entry & (1u64 << idx_in_entry)) != 0
     }
 
-    /// Creates a NullBuffer directly from the DuckDB validity mask for optimal performance.
-    ///
-    /// Returns None if all values are valid (no null buffer needed).
-    pub fn to_null_buffer(&self) -> Option<arrow_buffer::NullBuffer> {
+    /// Creates a Validity directly from the DuckDB validity mask for optimal performance.
+    pub fn to_validity(&self) -> Validity {
         let Some(validity) = self.validity else {
-            // All values are valid - no null buffer needed
-            return None;
+            // All values are valid
+            return Validity::AllValid;
         };
 
-        // Create copy of the buffer from the DuckDB validity mask.
-        let buffer = Buffer::from_iter(validity.iter().cloned());
-
-        let boolean_buffer = arrow_buffer::BooleanBuffer::new(buffer, 0, self.len);
-        Some(arrow_buffer::NullBuffer::new(boolean_buffer))
+        Validity::from(BitBuffer::new(
+            Buffer::<u64>::copy_from(validity).into_byte_buffer(),
+            self.len,
+        ))
     }
 }
 
@@ -325,19 +322,23 @@ mod tests {
     use crate::cpp::DUCKDB_TYPE;
 
     #[test]
-    fn test_create_null_buffer_all_valid() {
+    fn test_create_validity_all_valid() {
         // Test case where all values are valid - should return None
         let len = 10;
         let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_INTEGER);
         let vector = Vector::with_capacity(logical_type, len);
 
         let validity = vector.validity_ref(len);
-        let null_buffer = validity.to_null_buffer();
-        assert!(null_buffer.is_none(), "Expected None for all-valid vector");
+        let validity = validity.to_validity();
+        assert_eq!(
+            validity,
+            Validity::AllValid,
+            "Expected None for all-valid vector"
+        );
     }
 
     #[test]
-    fn test_create_null_buffer_with_nulls() {
+    fn test_create_validity_with_nulls() {
         // Test case with some null values
         let len = 10;
         let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_INTEGER);
@@ -351,30 +352,24 @@ mod tests {
         validity_slice.set(7, false); // null at position 7
 
         let validity = vector.validity_ref(len);
-        let null_buffer = validity.to_null_buffer();
-        assert!(
-            null_buffer.is_some(),
-            "Expected Some(NullBuffer) for vector with nulls"
-        );
-
-        let null_buffer = null_buffer.unwrap();
-        assert_eq!(null_buffer.len(), len);
+        let validity = validity.to_validity();
+        assert_eq!(validity.maybe_len(), Some(len));
 
         // Check that the right positions are null
-        assert!(null_buffer.is_valid(0));
-        assert!(null_buffer.is_null(1));
-        assert!(null_buffer.is_valid(2));
-        assert!(null_buffer.is_null(3));
-        assert!(null_buffer.is_valid(4));
-        assert!(null_buffer.is_valid(5));
-        assert!(null_buffer.is_valid(6));
-        assert!(null_buffer.is_null(7));
-        assert!(null_buffer.is_valid(8));
-        assert!(null_buffer.is_valid(9));
+        assert!(validity.is_valid(0));
+        assert!(validity.is_null(1));
+        assert!(validity.is_valid(2));
+        assert!(validity.is_null(3));
+        assert!(validity.is_valid(4));
+        assert!(validity.is_valid(5));
+        assert!(validity.is_valid(6));
+        assert!(validity.is_null(7));
+        assert!(validity.is_valid(8));
+        assert!(validity.is_valid(9));
     }
 
     #[test]
-    fn test_create_null_buffer_single_element() {
+    fn test_create_validity_single_element() {
         // Test with a single element that is null
         let len = 1;
         let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_INTEGER);
@@ -385,16 +380,13 @@ mod tests {
         validity_slice.set(0, false); // null at position 0
 
         let validity = vector.validity_ref(len);
-        let null_buffer = validity.to_null_buffer();
-        assert!(null_buffer.is_some());
-
-        let null_buffer = null_buffer.unwrap();
-        assert_eq!(null_buffer.len(), 1);
-        assert!(null_buffer.is_null(0));
+        let validity = validity.to_validity();
+        assert_eq!(validity.maybe_len(), Some(1));
+        assert!(validity.is_null(0));
     }
 
     #[test]
-    fn test_create_null_buffer_single_element_valid() {
+    fn test_create_validity_single_element_valid() {
         // Test with a single valid element
         let len = 1;
         let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_INTEGER);
@@ -405,29 +397,26 @@ mod tests {
         let _validity_slice = unsafe { vector.ensure_validity_bitslice(len) };
 
         let validity = vector.validity_ref(len);
-        let null_buffer = validity.to_null_buffer();
-        assert!(null_buffer.is_some());
-
-        let null_buffer = null_buffer.unwrap();
-        assert_eq!(null_buffer.len(), 1);
-        assert!(null_buffer.is_valid(0));
+        let validity = validity.to_validity();
+        assert_eq!(validity.maybe_len(), Some(1));
+        assert!(validity.is_valid(0));
     }
 
     #[test]
-    fn test_create_null_buffer_empty() {
+    fn test_create_validity_empty() {
         // Test with zero length
         let len = 0;
         let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_INTEGER);
         let vector = Vector::with_capacity(logical_type, len);
 
         let validity = vector.validity_ref(len);
-        let null_buffer = validity.to_null_buffer();
+        let validity = validity.to_validity();
         // Even with zero length, if validity mask doesn't exist, should return None
-        assert!(null_buffer.is_none());
+        assert_eq!(validity, Validity::AllValid);
     }
 
     #[test]
-    fn test_create_null_buffer_all_nulls() {
+    fn test_create_validity_all_nulls() {
         // Test case where all values are null
         let len = 10;
         let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_INTEGER);
@@ -441,16 +430,10 @@ mod tests {
         }
 
         let validity = vector.validity_ref(len);
-        let null_buffer = validity.to_null_buffer();
-        assert!(null_buffer.is_some());
-
-        let null_buffer = null_buffer.unwrap();
-        assert_eq!(null_buffer.len(), len);
-
+        let validity = validity.to_validity();
+        assert_eq!(validity.maybe_len(), Some(len));
         // Check that all positions are null
-        for i in 0..len {
-            assert!(null_buffer.is_null(i), "Element {i} should be null");
-        }
+        assert_eq!(validity, Validity::AllInvalid);
     }
 
     #[test]