From 65fe0fb58c6a0a26d2e1d98c3bc7c187305c2bb4 Mon Sep 17 00:00:00 2001 From: jeadie Date: Wed, 24 Apr 2024 11:56:36 +1000 Subject: [PATCH 01/10] support UTF8[] --- src/vtab/arrow.rs | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/src/vtab/arrow.rs b/src/vtab/arrow.rs index fa378c9c..8a17b76e 100644 --- a/src/vtab/arrow.rs +++ b/src/vtab/arrow.rs @@ -422,17 +422,20 @@ fn list_array_to_vector>( match value_array.data_type() { dt if dt.is_primitive() => { primitive_array_to_vector(value_array.as_ref(), &mut child)?; - for i in 0..array.len() { - let offset = array.value_offsets()[i]; - let length = array.value_length(i); - out.set_entry(i, offset.as_(), length.as_()); - } + } + DataType::Utf8 => { + string_array_to_vector(as_string_array(value_array.as_ref()), &mut child); } _ => { return Err("Nested list is not supported yet.".into()); } } + for i in 0..array.len() { + let offset = array.value_offsets()[i]; + let length = array.value_length(i); + out.set_entry(i, offset.as_(), length.as_()); + } Ok(()) } @@ -452,10 +455,19 @@ fn fixed_size_list_array_to_vector( } out.set_len(value_array.len()); } + DataType::Utf8 => { + string_array_to_vector(as_string_array(value_array.as_ref()), &mut child); + } _ => { return Err("Nested list is not supported yet.".into()); } } + for i in 0..array.len() { + let offset = array.value_offset(i); + let length = array.value_length(); + out.set_entry(i, offset as usize, length as usize); + } + out.set_len(value_array.len()); Ok(()) } From f3d9f77717def8ce458b495751396c9111ca2f88 Mon Sep 17 00:00:00 2001 From: jeadie Date: Wed, 24 Apr 2024 15:45:11 +1000 Subject: [PATCH 02/10] add tests --- src/vtab/arrow.rs | 67 ++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 61 insertions(+), 6 deletions(-) diff --git a/src/vtab/arrow.rs b/src/vtab/arrow.rs index 8a17b76e..90ff59b2 100644 --- a/src/vtab/arrow.rs +++ b/src/vtab/arrow.rs @@ -555,12 +555,8 @@ mod test { use crate::{Connection, Result}; use arrow::{ array::{ - Array, ArrayRef, AsArray, Date32Array, Date64Array, Decimal256Array, Float64Array, Int32Array, - PrimitiveArray, StringArray, StructArray, Time32SecondArray, Time64MicrosecondArray, - TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, - }, - datatypes::{i256, ArrowPrimitiveType, DataType, Field, Fields, Schema}, - record_batch::RecordBatch, + Array, ArrayRef, AsArray, Date32Array, Date64Array, Decimal256Array, Float64Array, GenericListArray, Int32Array, ListArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray, Time32SecondArray, Time64MicrosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray + }, buffer::{OffsetBuffer, ScalarBuffer}, datatypes::{i256, ArrowNativeType, ArrowPrimitiveType, DataType, Field, Fields, Schema}, record_batch::RecordBatch }; use std::{error::Error, sync::Arc}; @@ -688,6 +684,65 @@ mod test { Ok(()) } + fn check_generic_array_roundtrip( + arry: GenericListArray, + ) -> Result<(), Box> where T: OffsetSizeTrait{ + + let expected_output_array = arry.clone(); + + let db = Connection::open_in_memory()?; + db.register_table_function::("arrow")?; + + // Roundtrip a record batch from Rust to DuckDB and back to Rust + let schema = Schema::new(vec![Field::new("a", arry.data_type().clone(), false)]); + + let rb = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(arry.clone())])?; + let param = arrow_recordbatch_to_query_params(rb); + let mut stmt = db.prepare("select a from arrow(?, ?)")?; + let rb = stmt.query_arrow(param)?.next().expect("no record batch"); + + let output_any_array = rb.column(0); + match (output_any_array.data_type(), expected_output_array.data_type()) { + // TODO: DuckDB doesnt return timestamp_tz properly yet, so we just check that the units are the same + (DataType::Timestamp(unit_a, _), DataType::Timestamp(unit_b, _)) => assert_eq!(unit_a, unit_b), + (a, b) => assert_eq!(a, b), + } + + let maybe_output_array = output_any_array.as_primitive_opt::(); + + match maybe_output_array { + Some(output_array) => { + // Check that the output array is the same as the input array + assert_eq!(output_array.len(), expected_output_array.len()); + for i in 0..output_array.len() { + assert_eq!(output_array.is_valid(i), expected_output_array.is_valid(i)); + if output_array.is_valid(i) { + assert_eq!(output_array.value(i), expected_output_array.value(i)); + } + } + } + None => { + panic!("Output array is not a PrimitiveArray {:?}", rb.column(0).data_type()); + } + } + + Ok(()) + } + + + #[test] + fn test_array_roundtrip() -> Result<(), Box> { + + let list_array = ListArray::new( + Arc::new(Field::new("item", DataType::Utf8, true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 4, 5])), + Arc::new(StringArray::from(vec![Some("foo"), Some("baz"), Some("bar"), Some("foo"), Some("baz")])), None + ); + check_generic_array_roundtrip(list_array)?; + + Ok(()) + } + #[test] fn test_timestamp_roundtrip() -> Result<(), Box> { check_rust_primitive_array_roundtrip(Int32Array::from(vec![1, 2, 3]), Int32Array::from(vec![1, 2, 3]))?; From 2652a54f39f2905dbe7022c1a060d07decee0e87 Mon Sep 17 00:00:00 2001 From: jeadie Date: Wed, 24 Apr 2024 16:05:42 +1000 Subject: [PATCH 03/10] fix test --- src/vtab/arrow.rs | 38 +++++++++++++------------------------- 1 file changed, 13 insertions(+), 25 deletions(-) diff --git a/src/vtab/arrow.rs b/src/vtab/arrow.rs index 90ff59b2..3e40cff3 100644 --- a/src/vtab/arrow.rs +++ b/src/vtab/arrow.rs @@ -687,58 +687,46 @@ mod test { fn check_generic_array_roundtrip( arry: GenericListArray, ) -> Result<(), Box> where T: OffsetSizeTrait{ - + let expected_output_array = arry.clone(); - + let db = Connection::open_in_memory()?; db.register_table_function::("arrow")?; - + // Roundtrip a record batch from Rust to DuckDB and back to Rust let schema = Schema::new(vec![Field::new("a", arry.data_type().clone(), false)]); - + let rb = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(arry.clone())])?; let param = arrow_recordbatch_to_query_params(rb); let mut stmt = db.prepare("select a from arrow(?, ?)")?; let rb = stmt.query_arrow(param)?.next().expect("no record batch"); - + let output_any_array = rb.column(0); - match (output_any_array.data_type(), expected_output_array.data_type()) { - // TODO: DuckDB doesnt return timestamp_tz properly yet, so we just check that the units are the same - (DataType::Timestamp(unit_a, _), DataType::Timestamp(unit_b, _)) => assert_eq!(unit_a, unit_b), - (a, b) => assert_eq!(a, b), - } - - let maybe_output_array = output_any_array.as_primitive_opt::(); - - match maybe_output_array { + assert!(output_any_array.data_type().equals_datatype(expected_output_array.data_type())); + + match output_any_array.as_list_opt::() { Some(output_array) => { - // Check that the output array is the same as the input array assert_eq!(output_array.len(), expected_output_array.len()); for i in 0..output_array.len() { assert_eq!(output_array.is_valid(i), expected_output_array.is_valid(i)); if output_array.is_valid(i) { - assert_eq!(output_array.value(i), expected_output_array.value(i)); + assert!(expected_output_array.value(i).eq(&output_array.value(i))); } } } - None => { - panic!("Output array is not a PrimitiveArray {:?}", rb.column(0).data_type()); - } + None => panic!("Expected GenericListArray"), } - + Ok(()) } - #[test] fn test_array_roundtrip() -> Result<(), Box> { - - let list_array = ListArray::new( + check_generic_array_roundtrip(ListArray::new( Arc::new(Field::new("item", DataType::Utf8, true)), OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 4, 5])), Arc::new(StringArray::from(vec![Some("foo"), Some("baz"), Some("bar"), Some("foo"), Some("baz")])), None - ); - check_generic_array_roundtrip(list_array)?; + ))?; Ok(()) } From 092a2a9cd7b15dda55a2dd33d4b842432b63884f Mon Sep 17 00:00:00 2001 From: Max Gabrielsson Date: Wed, 24 Apr 2024 12:42:58 +0200 Subject: [PATCH 04/10] format --- src/vtab/arrow.rs | 47 +++++++++++++++++++++++++++++++---------------- 1 file changed, 31 insertions(+), 16 deletions(-) diff --git a/src/vtab/arrow.rs b/src/vtab/arrow.rs index 3e40cff3..3a2f3827 100644 --- a/src/vtab/arrow.rs +++ b/src/vtab/arrow.rs @@ -555,8 +555,14 @@ mod test { use crate::{Connection, Result}; use arrow::{ array::{ - Array, ArrayRef, AsArray, Date32Array, Date64Array, Decimal256Array, Float64Array, GenericListArray, Int32Array, ListArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray, Time32SecondArray, Time64MicrosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray - }, buffer::{OffsetBuffer, ScalarBuffer}, datatypes::{i256, ArrowNativeType, ArrowPrimitiveType, DataType, Field, Fields, Schema}, record_batch::RecordBatch + Array, ArrayRef, AsArray, Date32Array, Date64Array, Decimal256Array, Float64Array, GenericListArray, + Int32Array, ListArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray, Time32SecondArray, + Time64MicrosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, + TimestampSecondArray, + }, + buffer::{OffsetBuffer, ScalarBuffer}, + datatypes::{i256, ArrowNativeType, ArrowPrimitiveType, DataType, Field, Fields, Schema}, + record_batch::RecordBatch, }; use std::{error::Error, sync::Arc}; @@ -684,26 +690,28 @@ mod test { Ok(()) } - fn check_generic_array_roundtrip( - arry: GenericListArray, - ) -> Result<(), Box> where T: OffsetSizeTrait{ - - let expected_output_array = arry.clone(); - + fn check_generic_array_roundtrip(arry: GenericListArray) -> Result<(), Box> + where + T: OffsetSizeTrait, + { + let expected_output_array = arry.clone(); + let db = Connection::open_in_memory()?; db.register_table_function::("arrow")?; - + // Roundtrip a record batch from Rust to DuckDB and back to Rust let schema = Schema::new(vec![Field::new("a", arry.data_type().clone(), false)]); - + let rb = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(arry.clone())])?; let param = arrow_recordbatch_to_query_params(rb); let mut stmt = db.prepare("select a from arrow(?, ?)")?; let rb = stmt.query_arrow(param)?.next().expect("no record batch"); - + let output_any_array = rb.column(0); - assert!(output_any_array.data_type().equals_datatype(expected_output_array.data_type())); - + assert!(output_any_array + .data_type() + .equals_datatype(expected_output_array.data_type())); + match output_any_array.as_list_opt::() { Some(output_array) => { assert_eq!(output_array.len(), expected_output_array.len()); @@ -716,7 +724,7 @@ mod test { } None => panic!("Expected GenericListArray"), } - + Ok(()) } @@ -724,8 +732,15 @@ mod test { fn test_array_roundtrip() -> Result<(), Box> { check_generic_array_roundtrip(ListArray::new( Arc::new(Field::new("item", DataType::Utf8, true)), - OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 4, 5])), - Arc::new(StringArray::from(vec![Some("foo"), Some("baz"), Some("bar"), Some("foo"), Some("baz")])), None + OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 4, 5])), + Arc::new(StringArray::from(vec![ + Some("foo"), + Some("baz"), + Some("bar"), + Some("foo"), + Some("baz"), + ])), + None, ))?; Ok(()) From e106d79807920b977fcb75e948253c59352d8f72 Mon Sep 17 00:00:00 2001 From: Max Gabrielsson Date: Wed, 24 Apr 2024 14:08:34 +0200 Subject: [PATCH 05/10] clippy --- src/vtab/arrow.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/vtab/arrow.rs b/src/vtab/arrow.rs index 3a2f3827..3d6c23c2 100644 --- a/src/vtab/arrow.rs +++ b/src/vtab/arrow.rs @@ -561,7 +561,7 @@ mod test { TimestampSecondArray, }, buffer::{OffsetBuffer, ScalarBuffer}, - datatypes::{i256, ArrowNativeType, ArrowPrimitiveType, DataType, Field, Fields, Schema}, + datatypes::{i256, ArrowPrimitiveType, DataType, Field, Fields, Schema}, record_batch::RecordBatch, }; use std::{error::Error, sync::Arc}; From 741adae4d9ef395d04b2dc83008867ced724c2af Mon Sep 17 00:00:00 2001 From: Max Gabrielsson Date: Wed, 24 Apr 2024 22:43:23 +0200 Subject: [PATCH 06/10] bump cause github is broken --- Cargo.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Cargo.toml b/Cargo.toml index c5674e9b..d4223645 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -84,6 +84,7 @@ pretty_assertions = "1.4.0" path = "libduckdb-sys" version = "0.10.1" + [package.metadata.docs.rs] features = ['vtab', 'chrono'] all-features = false From c5c104ce72238e25ac950a41a46974ff8651316a Mon Sep 17 00:00:00 2001 From: jeadie Date: Mon, 3 Jun 2024 20:57:13 +1000 Subject: [PATCH 07/10] add support for DuckDB arrays when using Arrow's FixedSizeList --- src/vtab/arrow.rs | 4 ++-- src/vtab/logical_type.rs | 9 +++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/vtab/arrow.rs b/src/vtab/arrow.rs index 3d6c23c2..bde98c5d 100644 --- a/src/vtab/arrow.rs +++ b/src/vtab/arrow.rs @@ -196,8 +196,8 @@ pub fn to_duckdb_logical_type(data_type: &DataType) -> Result Self { + unsafe { + Self { + ptr: duckdb_create_array_type(child_type.ptr, array_size), + } + } + } + /// Creates a decimal type from its `width` and `scale`. pub fn decimal(width: u8, scale: u8) -> Self { unsafe { From 666a6499eea0f97ffde472d1a630f8f6564be1ca Mon Sep 17 00:00:00 2001 From: jeadie Date: Mon, 3 Jun 2024 21:00:10 +1000 Subject: [PATCH 08/10] fmt --- src/vtab/arrow.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/vtab/arrow.rs b/src/vtab/arrow.rs index bde98c5d..740667ed 100644 --- a/src/vtab/arrow.rs +++ b/src/vtab/arrow.rs @@ -197,7 +197,10 @@ pub fn to_duckdb_logical_type(data_type: &DataType) -> Result Date: Tue, 4 Jun 2024 20:04:27 +0200 Subject: [PATCH 09/10] add ArrayVector --- crates/duckdb/src/vtab/arrow.rs | 74 ++++++++++++++++++++-------- crates/duckdb/src/vtab/data_chunk.rs | 7 ++- crates/duckdb/src/vtab/vector.rs | 43 ++++++++++++++++ 3 files changed, 102 insertions(+), 22 deletions(-) diff --git a/crates/duckdb/src/vtab/arrow.rs b/crates/duckdb/src/vtab/arrow.rs index e613a86b..f1b8e9fe 100644 --- a/crates/duckdb/src/vtab/arrow.rs +++ b/crates/duckdb/src/vtab/arrow.rs @@ -1,5 +1,5 @@ use super::{ - vector::{FlatVector, ListVector, Vector}, + vector::{ArrayVector, FlatVector, ListVector, Vector}, BindInfo, DataChunk, Free, FunctionInfo, InitInfo, LogicalType, LogicalTypeId, StructVector, VTab, }; use std::ptr::null_mut; @@ -237,7 +237,7 @@ pub fn record_batch_to_duckdb_data_chunk( list_array_to_vector(as_large_list_array(col.as_ref()), &mut chunk.list_vector(i))?; } DataType::FixedSizeList(_, _) => { - fixed_size_list_array_to_vector(as_fixed_size_list_array(col.as_ref()), &mut chunk.list_vector(i))?; + fixed_size_list_array_to_vector(as_fixed_size_list_array(col.as_ref()), &mut chunk.array_vector(i))?; } DataType::Struct(_) => { let struct_array = as_struct_array(col.as_ref()); @@ -458,33 +458,21 @@ fn list_array_to_vector>( fn fixed_size_list_array_to_vector( array: &FixedSizeListArray, - out: &mut ListVector, + out: &mut ArrayVector, ) -> Result<(), Box> { let value_array = array.values(); let mut child = out.child(value_array.len()); match value_array.data_type() { dt if dt.is_primitive() => { primitive_array_to_vector(value_array.as_ref(), &mut child)?; - for i in 0..array.len() { - let offset = array.value_offset(i); - let length = array.value_length(); - out.set_entry(i, offset as usize, length as usize); - } - out.set_len(value_array.len()); } DataType::Utf8 => { string_array_to_vector(as_string_array(value_array.as_ref()), &mut child); } _ => { - return Err("Nested list is not supported yet.".into()); + return Err("Nested array is not supported yet.".into()); } } - for i in 0..array.len() { - let offset = array.value_offset(i); - let length = array.value_length(); - out.set_entry(i, offset as usize, length as usize); - } - out.set_len(value_array.len()); Ok(()) } @@ -514,7 +502,7 @@ fn struct_array_to_vector(array: &StructArray, out: &mut StructVector) -> Result DataType::FixedSizeList(_, _) => { fixed_size_list_array_to_vector( as_fixed_size_list_array(column.as_ref()), - &mut out.list_vector_child(i), + &mut out.array_vector_child(i), )?; } DataType::Struct(_) => { @@ -572,10 +560,10 @@ mod test { use crate::{Connection, Result}; use arrow::{ array::{ - Array, ArrayRef, AsArray, Date32Array, Date64Array, Decimal256Array, Float64Array, GenericListArray, - Int32Array, ListArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray, Time32SecondArray, - Time64MicrosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, - TimestampSecondArray, + Array, ArrayRef, AsArray, Date32Array, Date64Array, Decimal256Array, FixedSizeListArray, Float64Array, + GenericListArray, Int32Array, ListArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray, + Time32SecondArray, Time64MicrosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, }, buffer::{OffsetBuffer, ScalarBuffer}, datatypes::{i256, ArrowPrimitiveType, DataType, Field, Fields, Schema}, @@ -763,6 +751,50 @@ mod test { Ok(()) } + //field: FieldRef, size: i32, values: ArrayRef, nulls: Option + #[test] + fn test_fixed_array_roundtrip() -> Result<(), Box> { + let array = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + 2, + Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)])), + None, + ); + + let expected_output_array = array.clone(); + + let db = Connection::open_in_memory()?; + db.register_table_function::("arrow")?; + + // Roundtrip a record batch from Rust to DuckDB and back to Rust + let schema = Schema::new(vec![Field::new("a", array.data_type().clone(), false)]); + + let rb = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array.clone())])?; + let param = arrow_recordbatch_to_query_params(rb); + let mut stmt = db.prepare("select a from arrow(?, ?)")?; + let rb = stmt.query_arrow(param)?.next().expect("no record batch"); + + let output_any_array = rb.column(0); + assert!(output_any_array + .data_type() + .equals_datatype(expected_output_array.data_type())); + + match output_any_array.as_fixed_size_list_opt() { + Some(output_array) => { + assert_eq!(output_array.len(), expected_output_array.len()); + for i in 0..output_array.len() { + assert_eq!(output_array.is_valid(i), expected_output_array.is_valid(i)); + if output_array.is_valid(i) { + assert!(expected_output_array.value(i).eq(&output_array.value(i))); + } + } + } + None => panic!("Expected FixedSizeListArray"), + } + + Ok(()) + } + #[test] fn test_primitive_roundtrip_contains_nulls() -> Result<(), Box> { let mut builder = arrow::array::PrimitiveBuilder::::new(); diff --git a/crates/duckdb/src/vtab/data_chunk.rs b/crates/duckdb/src/vtab/data_chunk.rs index 6e472773..3bc6d874 100644 --- a/crates/duckdb/src/vtab/data_chunk.rs +++ b/crates/duckdb/src/vtab/data_chunk.rs @@ -1,6 +1,6 @@ use super::{ logical_type::LogicalType, - vector::{FlatVector, ListVector, StructVector}, + vector::{ArrayVector, FlatVector, ListVector, StructVector}, }; use crate::ffi::{ duckdb_create_data_chunk, duckdb_data_chunk, duckdb_data_chunk_get_column_count, duckdb_data_chunk_get_size, @@ -35,6 +35,11 @@ impl DataChunk { ListVector::from(unsafe { duckdb_data_chunk_get_vector(self.ptr, idx as u64) }) } + /// Get a array vector from the column index. + pub fn array_vector(&self, idx: usize) -> ArrayVector { + ArrayVector::from(unsafe { duckdb_data_chunk_get_vector(self.ptr, idx as u64) }) + } + /// Get struct vector at the column index: `idx`. pub fn struct_vector(&self, idx: usize) -> StructVector { StructVector::from(unsafe { duckdb_data_chunk_get_vector(self.ptr, idx as u64) }) diff --git a/crates/duckdb/src/vtab/vector.rs b/crates/duckdb/src/vtab/vector.rs index bf61cff4..030cf6ee 100644 --- a/crates/duckdb/src/vtab/vector.rs +++ b/crates/duckdb/src/vtab/vector.rs @@ -1,5 +1,7 @@ use std::{any::Any, ffi::CString, slice}; +use libduckdb_sys::{duckdb_array_type_array_size, duckdb_array_vector_get_child}; + use super::LogicalType; use crate::ffi::{ duckdb_list_entry, duckdb_list_vector_get_child, duckdb_list_vector_get_size, duckdb_list_vector_reserve, @@ -170,6 +172,42 @@ impl ListVector { } } +/// A array vector. (fixed-size list) +pub struct ArrayVector { + /// ArrayVector does not own the vector pointer. + ptr: duckdb_vector, +} + +impl From for ArrayVector { + fn from(ptr: duckdb_vector) -> Self { + Self { ptr } + } +} + +impl ArrayVector { + /// Get the logical type of this ArrayVector. + pub fn logical_type(&self) -> LogicalType { + LogicalType::from(unsafe { duckdb_vector_get_column_type(self.ptr) }) + } + + pub fn get_array_size(&self) -> u64 { + let ty = self.logical_type(); + unsafe { duckdb_array_type_array_size(ty.ptr) as u64 } + } + + /// Returns the child vector. + /// capacity should be a multiple of the array size. + // TODO: not ideal interface. Where should we keep count. + pub fn child(&self, capacity: usize) -> FlatVector { + FlatVector::with_capacity(unsafe { duckdb_array_vector_get_child(self.ptr) }, capacity) + } + + /// Set primitive data to the child node. + pub fn set_child(&self, data: &[T]) { + self.child(data.len()).copy(data); + } +} + /// A struct vector. pub struct StructVector { /// ListVector does not own the vector pointer. @@ -198,6 +236,11 @@ impl StructVector { ListVector::from(unsafe { duckdb_struct_vector_get_child(self.ptr, idx as u64) }) } + /// Take the child as [ArrayVector]. + pub fn array_vector_child(&self, idx: usize) -> ArrayVector { + ArrayVector::from(unsafe { duckdb_struct_vector_get_child(self.ptr, idx as u64) }) + } + /// Get the logical type of this struct vector. pub fn logical_type(&self) -> LogicalType { LogicalType::from(unsafe { duckdb_vector_get_column_type(self.ptr) }) From 22bcc02bfe5f2732647b90ba06bea39733abb228 Mon Sep 17 00:00:00 2001 From: Max Gabrielsson Date: Tue, 4 Jun 2024 20:06:42 +0200 Subject: [PATCH 10/10] update path in remote test --- crates/duckdb/src/extension.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/duckdb/src/extension.rs b/crates/duckdb/src/extension.rs index 1fa54c03..1a6a9690 100644 --- a/crates/duckdb/src/extension.rs +++ b/crates/duckdb/src/extension.rs @@ -38,7 +38,7 @@ mod test { let db = Connection::open_in_memory()?; assert_eq!( 300f32, - db.query_row::(r#"SELECT SUM(value) FROM read_parquet('https://github.com/wangfenjin/duckdb-rs/raw/main/examples/int32_decimal.parquet');"#, [], |r| r.get(0))? + db.query_row::(r#"SELECT SUM(value) FROM read_parquet('https://github.com/duckdb/duckdb-rs/raw/main/crates/duckdb/examples/int32_decimal.parquet');"#, [], |r| r.get(0))? ); Ok(()) }