@@ -6,9 +6,9 @@ use std::ptr::null_mut;
6
6
7
7
use crate :: vtab:: vector:: Inserter ;
8
8
use arrow:: array:: {
9
- as_boolean_array, as_large_list_array, as_list_array, as_primitive_array, as_string_array, as_struct_array , Array ,
10
- ArrayData , AsArray , BooleanArray , Decimal128Array , FixedSizeListArray , GenericListArray , OffsetSizeTrait ,
11
- PrimitiveArray , StringArray , StructArray ,
9
+ as_boolean_array, as_generic_binary_array , as_large_list_array, as_list_array, as_primitive_array, as_string_array,
10
+ as_struct_array , Array , ArrayData , AsArray , BinaryArray , BooleanArray , Decimal128Array , FixedSizeListArray ,
11
+ GenericListArray , OffsetSizeTrait , PrimitiveArray , StringArray , StructArray ,
12
12
} ;
13
13
14
14
use arrow:: {
@@ -230,6 +230,9 @@ pub fn record_batch_to_duckdb_data_chunk(
230
230
DataType :: Utf8 => {
231
231
string_array_to_vector ( as_string_array ( col. as_ref ( ) ) , & mut chunk. flat_vector ( i) ) ;
232
232
}
233
+ DataType :: Binary => {
234
+ binary_array_to_vector ( as_generic_binary_array ( col. as_ref ( ) ) , & mut chunk. flat_vector ( i) ) ;
235
+ }
233
236
DataType :: List ( _) => {
234
237
list_array_to_vector ( as_list_array ( col. as_ref ( ) ) , & mut chunk. list_vector ( i) ) ?;
235
238
}
@@ -430,6 +433,15 @@ fn string_array_to_vector(array: &StringArray, out: &mut FlatVector) {
430
433
}
431
434
}
432
435
436
+ fn binary_array_to_vector ( array : & BinaryArray , out : & mut FlatVector ) {
437
+ assert ! ( array. len( ) <= out. capacity( ) ) ;
438
+
439
+ for i in 0 ..array. len ( ) {
440
+ let s = array. value ( i) ;
441
+ out. insert ( i, s) ;
442
+ }
443
+ }
444
+
433
445
fn list_array_to_vector < O : OffsetSizeTrait + AsPrimitive < usize > > (
434
446
array : & GenericListArray < O > ,
435
447
out : & mut ListVector ,
@@ -443,6 +455,9 @@ fn list_array_to_vector<O: OffsetSizeTrait + AsPrimitive<usize>>(
443
455
DataType :: Utf8 => {
444
456
string_array_to_vector ( as_string_array ( value_array. as_ref ( ) ) , & mut child) ;
445
457
}
458
+ DataType :: Binary => {
459
+ binary_array_to_vector ( as_generic_binary_array ( value_array. as_ref ( ) ) , & mut child) ;
460
+ }
446
461
_ => {
447
462
return Err ( "Nested list is not supported yet." . into ( ) ) ;
448
463
}
@@ -469,6 +484,9 @@ fn fixed_size_list_array_to_vector(
469
484
DataType :: Utf8 => {
470
485
string_array_to_vector ( as_string_array ( value_array. as_ref ( ) ) , & mut child) ;
471
486
}
487
+ DataType :: Binary => {
488
+ binary_array_to_vector ( as_generic_binary_array ( value_array. as_ref ( ) ) , & mut child) ;
489
+ }
472
490
_ => {
473
491
return Err ( "Nested array is not supported yet." . into ( ) ) ;
474
492
}
@@ -493,6 +511,9 @@ fn struct_array_to_vector(array: &StructArray, out: &mut StructVector) -> Result
493
511
DataType :: Utf8 => {
494
512
string_array_to_vector ( as_string_array ( column. as_ref ( ) ) , & mut out. child ( i) ) ;
495
513
}
514
+ DataType :: Binary => {
515
+ binary_array_to_vector ( as_generic_binary_array ( column. as_ref ( ) ) , & mut out. child ( i) ) ;
516
+ }
496
517
DataType :: List ( _) => {
497
518
list_array_to_vector ( as_list_array ( column. as_ref ( ) ) , & mut out. list_vector_child ( i) ) ?;
498
519
}
@@ -560,10 +581,10 @@ mod test {
560
581
use crate :: { Connection , Result } ;
561
582
use arrow:: {
562
583
array:: {
563
- Array , ArrayRef , AsArray , Date32Array , Date64Array , Decimal256Array , FixedSizeListArray , Float64Array ,
564
- GenericListArray , Int32Array , ListArray , OffsetSizeTrait , PrimitiveArray , StringArray , StructArray ,
565
- Time32SecondArray , Time64MicrosecondArray , TimestampMicrosecondArray , TimestampMillisecondArray ,
566
- TimestampNanosecondArray , TimestampSecondArray ,
584
+ Array , ArrayRef , AsArray , BinaryArray , Date32Array , Date64Array , Decimal256Array , FixedSizeListArray ,
585
+ Float64Array , GenericListArray , Int32Array , ListArray , OffsetSizeTrait , PrimitiveArray , StringArray ,
586
+ StructArray , Time32SecondArray , Time64MicrosecondArray , TimestampMicrosecondArray ,
587
+ TimestampMillisecondArray , TimestampNanosecondArray , TimestampSecondArray ,
567
588
} ,
568
589
buffer:: { OffsetBuffer , ScalarBuffer } ,
569
590
datatypes:: { i256, ArrowPrimitiveType , DataType , Field , Fields , Schema } ,
@@ -924,4 +945,23 @@ mod test {
924
945
)
925
946
) ;
926
947
}
948
+
949
+ #[ test]
950
+ fn test_arrow_binary ( ) {
951
+ let byte_array = BinaryArray :: from_iter_values ( [ b"test" ] . iter ( ) ) ;
952
+ let arc: ArrayRef = Arc :: new ( byte_array) ;
953
+ let batch = RecordBatch :: try_from_iter ( vec ! [ ( "x" , arc) ] ) . unwrap ( ) ;
954
+
955
+ let db = Connection :: open_in_memory ( ) . unwrap ( ) ;
956
+ db. register_table_function :: < ArrowVTab > ( "arrow" ) . unwrap ( ) ;
957
+
958
+ let mut stmt = db. prepare ( "SELECT * FROM arrow(?, ?)" ) . unwrap ( ) ;
959
+
960
+ let mut arr = stmt. query_arrow ( arrow_recordbatch_to_query_params ( batch) ) . unwrap ( ) ;
961
+ let rb = arr. next ( ) . expect ( "no record batch" ) ;
962
+
963
+ let column = rb. column ( 0 ) . as_any ( ) . downcast_ref :: < BinaryArray > ( ) . unwrap ( ) ;
964
+ assert_eq ! ( column. len( ) , 1 ) ;
965
+ assert_eq ! ( column. value( 0 ) , b"test" ) ;
966
+ }
927
967
}
0 commit comments