fixup

Author: joseph-isaacs

vortex-duckdb/src/convert/array/mod.rs

@@ -4,54 +4,64 @@ use duckdb::vtab::arrow::{
 };
 use vortex_array::arrays::StructArray;
 use vortex_array::arrow::{FromArrowArray, IntoArrowArray};
-use vortex_array::builders::ArrayBuilder;
 use vortex_array::validity::Validity;
-use vortex_array::{Array, ArrayRef, Canonical};
-use vortex_error::VortexResult;
+use vortex_array::{Array, ArrayRef};
+use vortex_dtype::FieldNames;
+use vortex_error::{VortexResult, vortex_err};
 
 pub trait ToDuckDB {
     fn to_duckdb(&self, chunk: &mut dyn WritableVector) -> VortexResult<()>;
 }
 
 impl ToDuckDB for ArrayRef {
     fn to_duckdb(&self, chunk: &mut dyn WritableVector) -> VortexResult<()> {
-        write_arrow_array_to_vector(&self.clone().into_arrow_preferred()?, chunk)?;
+        write_arrow_array_to_vector(&self.clone().into_arrow_preferred()?, chunk)
+            .map_err(|e| vortex_err!("Failed to convert vrotex duckdb array: {}", e.to_string()))
     }
 }
 
-struct SizedFlatVector<'a> {
-    pub vector: &'a FlatVector,
+struct NamedDataChunk<'a> {
+    pub chunk: &'a DataChunkHandle,
+    pub names: FieldNames,
+}
+
+struct SizedFlatVector {
+    pub vector: FlatVector,
     pub len: usize,
 }
 
 pub trait FromDuckDB<V> {
-    fn from_duckdb(vector: &V) -> VortexResult<ArrayRef>;
+    fn from_duckdb(vector: V) -> VortexResult<ArrayRef>;
 }
 
-impl FromDuckDB<DataChunkHandle> for ArrayRef {
-    fn from_duckdb(chunk: &DataChunkHandle) -> VortexResult<ArrayRef> {
+impl<'a> FromDuckDB<&'a NamedDataChunk<'a>> for ArrayRef {
+    fn from_duckdb(named_chunk: &'a NamedDataChunk<'a>) -> VortexResult<ArrayRef> {
+        let chunk = &named_chunk.chunk;
+        let names = &named_chunk.names;
         let len = chunk.len();
 
         let columns = (0..chunk.num_columns())
             .map(|i| {
                 let vector = chunk.flat_vector(i);
-                ArrayRef::from_duckdb(&SizedFlatVector {
-                    vector: &vector,
-                    len,
-                })
+                let array = ArrayRef::from_duckdb(SizedFlatVector { vector, len })?;
+                // Figure out the column names
+                Ok((names[i].clone(), array))
             })
             .collect::<VortexResult<Vec<_>>>()?;
 
-        let (names, arrays) = columns.iter().unzip();
+        let (names, arrays): (Vec<_>, Vec<_>) = columns.into_iter().unzip();
 
         // TODO(joe): extract validity
-        StructArray::try_new(names, arrays, len, Validity::AllValid).map(StructArray::to_array)
+        StructArray::try_new(names.into(), arrays, len, Validity::AllValid)
+            .map(StructArray::into_array)
     }
 }
 
 impl FromDuckDB<SizedFlatVector> for ArrayRef {
-    fn from_duckdb(vector: &SizedFlatVector) -> VortexResult<ArrayRef> {
-        let arrow_arr = flat_vector_to_arrow_array(&mut vector.vector.clone(), vector.len)?;
+    fn from_duckdb(mut sized_vector: SizedFlatVector) -> VortexResult<ArrayRef> {
+        let len = sized_vector.len;
+        let arrow_arr = flat_vector_to_arrow_array(&mut sized_vector.vector, len)
+            .map_err(|e| vortex_err!("Failed to convert duckdb array to vortex: {}", e))?;
         Ok(ArrayRef::from_arrow(arrow_arr, true))
     }
 }

vortex-duckdb/src/convert/mod.rs

@@ -1,2 +1,5 @@
 mod array;
 mod types;
+
+pub use array::{FromDuckDB, ToDuckDB};
+pub use types::{FromDuckDBType, ToDuckDBType};

vortex-duckdb/src/convert/types/from.rs

@@ -55,14 +55,12 @@ fn from_duckdb_list(list: LogicalTypeHandle) -> DType {
 }
 
 fn from_duckdb_struct(struct_: LogicalTypeHandle) -> StructDType {
-    struct_
-        .num_children()
-        .into_iter()
+    (0..struct_.num_children())
         .map(|i| {
             // TODO: is there struct field nullability
             let child_nullability = Nullable;
             let child_name = struct_.child_name(i);
-            let child_type = FromDuckDBType::from_duckdb(struct_.child(i), child_nullability);
+            let child_type = DType::from_duckdb(struct_.child(i), child_nullability);
             (child_name, child_type)
         })
         .collect()

vortex-duckdb/src/convert/types/mod.rs

@@ -1,2 +1,5 @@
 mod from;
 mod to;
+
+pub use from::FromDuckDBType;
+pub use to::ToDuckDBType;

vortex-duckdb/src/convert/types/to.rs

@@ -1,6 +1,6 @@
 use duckdb::core::{LogicalTypeHandle, LogicalTypeId};
 use vortex_dtype::{DType, PType};
-use vortex_error::{VortexResult, vortex_bail, vortex_err};
+use vortex_error::{VortexResult, vortex_bail};
 
 pub trait ToDuckDBType {
     fn to_duckdb_type(&self) -> VortexResult<LogicalTypeHandle>;

vortex-duckdb/src/lib.rs

@@ -1 +1,3 @@
 mod convert;
+
+pub use convert::{FromDuckDB, FromDuckDBType, ToDuckDB, ToDuckDBType};