Support varbinview buffer passing to_duckdb (#2791)

Author: joseph-isaacs

Cargo.lock

@@ -8,7 +8,7 @@ class VortexExtension : public Extension {
 public:
 	void Load(DuckDB &db) override;
 	std::string Name() override;
-        std::string Version() const override;
+	std::string Version() const override;
 };
 
 } // namespace duckdb

duckdb-vortex/duckdb

@@ -11,11 +11,8 @@
 
 #include "vortex.h"
 
-#include <expr/expr.hpp>
-
-extern "C" {
-const char *vortex_duckdb_hello();
-}
+#include "expr/expr.hpp"
+#include "rust_vector_buffer.hpp"
 
 #ifndef DUCKDB_EXTENSION_MAIN
 #error DUCKDB_EXTENSION_MAIN not defined

duckdb-vortex/duckdb-rs

@@ -24,6 +24,7 @@ crate-type = ["staticlib", "cdylib", "rlib"]
 arrow-array = { workspace = true }
 itertools = { workspace = true }
 vortex-array = { workspace = true }
+vortex-buffer = { workspace = true }
 vortex-dict = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }

duckdb-vortex/src/include/vortex_extension.hpp

@@ -0,0 +1,82 @@
+use std::os::raw::c_void;
+
+use duckdb::ffi::duckdb_vector;
+use vortex_buffer::ByteBuffer;
+
+#[derive(Clone)]
+#[repr(C)]
+pub struct FFIDuckDBBufferInternal {
+    pub inner: Box<ByteBuffer>,
+}
+
+// This CANNOT be copied or cloned since the void* inner is actually an Arc, and the ref counter
+// will be incorrect.
+#[repr(C)]
+pub struct ExternalBuffer {
+    pub inner: *mut c_void,
+}
+
+impl From<FFIDuckDBBufferInternal> for ExternalBuffer {
+    fn from(buffer: FFIDuckDBBufferInternal) -> Self {
+        let ptr = Box::into_raw(buffer.inner) as *mut c_void;
+        ExternalBuffer { inner: ptr }
+    }
+}
+
+impl From<ExternalBuffer> for FFIDuckDBBufferInternal {
+    fn from(buffer: ExternalBuffer) -> Self {
+        let inner: Box<ByteBuffer> = unsafe { Box::from_raw(buffer.inner.cast()) };
+        FFIDuckDBBufferInternal { inner }
+    }
+}
+
+// This will free a single FFIDuckDBBuffer, however due to cloning there might be more
+// references to the underlying ByteBuffer that will not be freed in this call.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn ExternalBuffer_free(buffer: *mut ExternalBuffer) {
+    let internal: Box<FFIDuckDBBufferInternal> = unsafe { Box::from_raw(buffer.cast()) };
+    drop(internal)
+}
+
+#[repr(C)]
+#[allow(dead_code)]
+pub struct CppVectorBuffer {
+    pub ptr: *mut c_void,
+}
+
+#[allow(dead_code)]
+unsafe extern "C" {
+    pub fn NewCppVectorBuffer(
+        buffer: *mut ExternalBuffer,
+        free: unsafe extern "C" fn(*mut ExternalBuffer),
+    ) -> *mut CppVectorBuffer;
+
+    pub fn AssignBufferToVec(vector: duckdb_vector, buffer: *mut CppVectorBuffer);
+}
+
+pub unsafe fn new_cpp_vector_buffer(buffer: *mut ExternalBuffer) -> *mut CppVectorBuffer {
+    unsafe { NewCppVectorBuffer(buffer, ExternalBuffer_free) }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use vortex_buffer::ByteBuffer;
+
+    use crate::buffer::{ExternalBuffer, FFIDuckDBBufferInternal};
+
+    #[test]
+    fn test_buff_drop() {
+        let buffer = FFIDuckDBBufferInternal {
+            inner: Box::new(ByteBuffer::from(vec![1, 2, 3])),
+        };
+
+        assert!(buffer.inner.inner().is_unique());
+
+        let buffer_er: ExternalBuffer = buffer.clone().into();
+        let buffer_back: FFIDuckDBBufferInternal = buffer_er.into();
+
+        assert!(!buffer_back.inner.inner().is_unique());
+        assert!(!buffer.inner.inner().is_unique());
+    }
+}

duckdb-vortex/src/vortex_extension.cpp

@@ -1,11 +1,14 @@
 mod data_chunk_adaptor;
+mod varbinview;
 
 use arrow_array::ArrayRef as ArrowArrayRef;
 use duckdb::core::{DataChunkHandle, SelectionVector};
 use duckdb::vtab::arrow::{
     WritableVector, flat_vector_to_arrow_array, write_arrow_array_to_vector,
 };
-use vortex_array::arrays::StructArray;
+use vortex_array::arrays::{
+    ChunkedArray, ChunkedEncoding, StructArray, VarBinViewArray, VarBinViewEncoding,
+};
 use vortex_array::arrow::FromArrowArray;
 use vortex_array::compute::{take, to_arrow_preferred, try_cast};
 use vortex_array::validity::Validity;
@@ -32,6 +35,18 @@ pub fn to_duckdb(array: ArrayRef, chunk: &mut dyn WritableVector) -> VortexResul
         let value = constant.to_duckdb_scalar();
         chunk.flat_vector().assign_to_constant(&value);
         Ok(())
+    } else if array.is_encoding(ChunkedEncoding.id()) {
+        array
+            .as_any()
+            .downcast_ref::<ChunkedArray>()
+            .vortex_expect("chunk checked")
+            .to_duckdb(chunk)
+    } else if array.is_encoding(VarBinViewEncoding.id()) {
+        array
+            .as_any()
+            .downcast_ref::<VarBinViewArray>()
+            .vortex_expect("varbinview id checked")
+            .to_duckdb(chunk)
     } else if array.is_encoding(DictEncoding.id()) {
         array
             .as_any()
@@ -43,6 +58,17 @@ pub fn to_duckdb(array: ArrayRef, chunk: &mut dyn WritableVector) -> VortexResul
     }
 }
 
+impl ToDuckDB for ChunkedArray {
+    fn to_duckdb(&self, chunk: &mut dyn WritableVector) -> VortexResult<()> {
+        // TODO(joe): support multi-chunk arrays without canonical.
+        if self.chunks().len() > 1 {
+            to_arrow_preferred(self)?.to_duckdb(chunk)
+        } else {
+            to_duckdb(self.chunks()[0].clone(), chunk)
+        }
+    }
+}
+
 impl ToDuckDB for DictArray {
     fn to_duckdb(&self, chunk: &mut dyn WritableVector) -> VortexResult<()> {
         // If the values fit into a single vector, we can efficiently delay the take operation.
@@ -133,7 +159,6 @@ impl FromDuckDB<SizedFlatVector> for ArrayRef {
 
 #[cfg(test)]
 mod tests {
-
     use duckdb::core::{DataChunkHandle, LogicalTypeHandle, LogicalTypeId};
     use vortex_array::arrays::{
         BoolArray, ConstantArray, PrimitiveArray, StructArray, VarBinArray,

vortex-duckdb/Cargo.toml

@@ -0,0 +1,116 @@
+use std::ffi::c_char;
+
+use duckdb::vtab::arrow::WritableVector;
+use itertools::Itertools;
+use vortex_array::arrays::{BinaryView, Inlined, VarBinViewArray};
+use vortex_buffer::ByteBuffer;
+use vortex_error::VortexResult;
+
+use crate::ToDuckDB;
+use crate::buffer::{
+    AssignBufferToVec, ExternalBuffer, FFIDuckDBBufferInternal, new_cpp_vector_buffer,
+};
+
+// This is the C++ string view struct
+// private:
+// 	union {
+// 		struct {
+// 			uint32_t length;
+// 			char prefix[4];
+// 			char *ptr;
+// 		} pointer;
+// 		struct {
+// 			uint32_t length;
+// 			char inlined[12];
+// 		} inlined;
+// 	} value;
+// };
+
+#[derive(Clone, Copy)]
+#[repr(C, align(16))]
+// See `BinaryView`
+pub union PtrBinaryView {
+    // Numeric representation. This is logically `u128`, but we split it into the high and low
+    // bits to preserve the alignment.
+    le_bytes: [u8; 16],
+
+    // Inlined representation: strings <= 12 bytes
+    inlined: Inlined,
+
+    // Reference type: strings > 12 bytes.
+    _ref: PtrRef,
+}
+
+#[derive(Clone, Copy, Debug)]
+#[repr(C, align(8))]
+pub struct PtrRef {
+    size: u32,
+    prefix: [u8; 4],
+    ptr: *const c_char,
+}
+
+fn binary_view_to_ptr_binary_view<'a>(
+    view: impl Iterator<Item = &'a BinaryView>,
+    buffers: &[ByteBuffer],
+    used_buffers: &mut [bool],
+) -> Vec<PtrBinaryView> {
+    view.map(|v| {
+        if v.is_inlined() {
+            PtrBinaryView {
+                inlined: *v.as_inlined(),
+            }
+        } else {
+            let view = v.as_view();
+            used_buffers[view.buffer_index() as usize] = true;
+            PtrBinaryView {
+                _ref: PtrRef {
+                    size: v.len(),
+                    prefix: *view.prefix(),
+                    // TODO(joe) verify this.
+                    ptr: unsafe {
+                        buffers[view.buffer_index() as usize]
+                            .as_ptr()
+                            .add(view.offset() as usize)
+                            .cast()
+                    },
+                },
+            }
+        }
+    })
+    .collect_vec()
+}
+
+impl ToDuckDB for VarBinViewArray {
+    fn to_duckdb(&self, chunk: &mut dyn WritableVector) -> VortexResult<()> {
+        let buffers = self.buffers();
+        let mut buffer_used = vec![false; buffers.len()];
+
+        let views: Vec<PtrBinaryView> = binary_view_to_ptr_binary_view(
+            self.views().iter(),
+            buffers,
+            buffer_used.as_mut_slice(),
+        );
+
+        let vec = chunk.flat_vector();
+        buffers
+            .iter()
+            .enumerate()
+            .filter(|&(idx, _buf)| buffer_used[idx])
+            .map(|(_idx, buf)| buf.clone())
+            .for_each(|b| {
+                // Each buffer is wrapped with a C++ VectorBuffer wrapper which will
+                // in turn call `FFIDuckDBBuffer_free` when it is cleaned up in C++ land.
+                // Once all ptrs to the bytes are free the bytes can be freed.
+                let buffer: *mut ExternalBuffer = Box::into_raw(Box::new(
+                    FFIDuckDBBufferInternal { inner: Box::new(b) }.into(),
+                ));
+                let extern_buf = unsafe { new_cpp_vector_buffer(buffer) };
+                // Adds an extra ref to the buffer which will outlive the `views`
+                unsafe { AssignBufferToVec(vec.unowned_ptr(), extern_buf) };
+            });
+
+        chunk.flat_vector().copy(views.as_slice());
+
+        Ok(())
+    }
+}

vortex-duckdb/src/buffer/mod.rs

@@ -3,17 +3,28 @@
 
 /// This is the default chunk size for duckdb.
 /// It is best to return data chunks of this size to duckdb.
+/// 2048 is the default chunk size for duckdb.
 pub const DUCKDB_STANDARD_VECTOR_SIZE: usize = 2048;
 
+mod buffer;
 mod convert;
 
-use std::ffi::c_char;
-
 pub use convert::{FromDuckDB, FromDuckDBType, ToDuckDB, ToDuckDBType, to_duckdb_chunk};
 
-// To generate C decls to include in vortex_duckdb_extension.cpp,
+// Note: To generate C decls to include in vortex_duckdb_extension.cpp,
 // call `cbindgen` from `vortex/vortex-duckdb`.
-#[unsafe(no_mangle)]
-pub unsafe extern "C" fn vortex_duckdb_hello() -> *const c_char {
-    c"Hello, world from Rust! ".as_ptr()
+
+#[cfg(test)]
+mod tests {
+    use duckdb::ffi::duckdb_vector_size;
+
+    use crate::DUCKDB_STANDARD_VECTOR_SIZE;
+
+    #[test]
+    fn assert_duckdb_vector_size_matches() {
+        assert_eq!(
+            Ok(DUCKDB_STANDARD_VECTOR_SIZE),
+            usize::try_from(unsafe { duckdb_vector_size() })
+        );
+    }
 }