Skip to content

Commit c5288e0

Browse files
committed
_bool (postgres): impl arrow_assoc for postgres
Modify `impl_arrow_assoc_vec` to take three arguments. Previously, its macro assumed a `MutablePrimitiveArray<$T>`; however, the primitive for booleans is not `MutablePrimitiveArray<bool>` but instead `MutableBooleanArray`. Make calling conventions easier by defining an extra macro to use in the general primitive type case.
1 parent 1695450 commit c5288e0

File tree

3 files changed

+21
-11
lines changed

3 files changed

+21
-11
lines changed

connectorx/src/destinations/arrow2/arrow_assoc.rs

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -62,12 +62,12 @@ impl_arrow_assoc!(f64, ArrowDataType::Float64, MutablePrimitiveArray<f64>);
6262
impl_arrow_assoc!(bool, ArrowDataType::Boolean, MutableBooleanArray);
6363

6464
macro_rules! impl_arrow_assoc_vec {
65-
($T:ty, $AT:expr) => {
65+
($T:ty, $PT:ty, $AT:expr) => {
6666
impl ArrowAssoc for Vec<$T> {
67-
type Builder = MutableListArray<i64, MutablePrimitiveArray<$T>>;
67+
type Builder = MutableListArray<i64, $PT>;
6868

6969
fn builder(nrows: usize) -> Self::Builder {
70-
MutableListArray::<i64, MutablePrimitiveArray<$T>>::with_capacity(nrows)
70+
MutableListArray::<i64, $PT>::with_capacity(nrows)
7171
}
7272

7373
#[inline]
@@ -86,10 +86,10 @@ macro_rules! impl_arrow_assoc_vec {
8686
}
8787

8888
impl ArrowAssoc for Option<Vec<$T>> {
89-
type Builder = MutableListArray<i64, MutablePrimitiveArray<$T>>;
89+
type Builder = MutableListArray<i64, $PT>;
9090

9191
fn builder(nrows: usize) -> Self::Builder {
92-
MutableListArray::<i64, MutablePrimitiveArray<$T>>::with_capacity(nrows)
92+
MutableListArray::<i64, $PT>::with_capacity(nrows)
9393
}
9494

9595
#[inline]
@@ -114,12 +114,19 @@ macro_rules! impl_arrow_assoc_vec {
114114
};
115115
}
116116

117-
impl_arrow_assoc_vec!(i32, ArrowDataType::Int32);
118-
impl_arrow_assoc_vec!(i64, ArrowDataType::Int64);
119-
impl_arrow_assoc_vec!(u32, ArrowDataType::UInt32);
120-
impl_arrow_assoc_vec!(u64, ArrowDataType::UInt64);
121-
impl_arrow_assoc_vec!(f32, ArrowDataType::Float32);
122-
impl_arrow_assoc_vec!(f64, ArrowDataType::Float64);
117+
macro_rules! impl_arrow_assoc_primitive_vec {
118+
($T:ty, $AT:expr) => {
119+
impl_arrow_assoc_vec!($T, MutablePrimitiveArray<$T>, $AT);
120+
};
121+
}
122+
123+
impl_arrow_assoc_vec!(bool, MutableBooleanArray, ArrowDataType::Boolean);
124+
impl_arrow_assoc_primitive_vec!(i32, ArrowDataType::Int32);
125+
impl_arrow_assoc_primitive_vec!(i64, ArrowDataType::Int64);
126+
impl_arrow_assoc_primitive_vec!(u32, ArrowDataType::UInt32);
127+
impl_arrow_assoc_primitive_vec!(u64, ArrowDataType::UInt64);
128+
impl_arrow_assoc_primitive_vec!(f32, ArrowDataType::Float32);
129+
impl_arrow_assoc_primitive_vec!(f64, ArrowDataType::Float64);
123130

124131
impl ArrowAssoc for &str {
125132
type Builder = MutableUtf8Array<i64>;

connectorx/src/destinations/arrow2/typesystem.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ pub enum Arrow2TypeSystem {
1616
Date64(bool),
1717
Time64(bool),
1818
DateTimeTz(bool),
19+
BoolArray(bool),
1920
Int32Array(bool),
2021
Int64Array(bool),
2122
UInt32Array(bool),
@@ -41,6 +42,7 @@ impl_typesystem! {
4142
{ Date64 => NaiveDateTime }
4243
{ Time64 => NaiveTime }
4344
{ DateTimeTz => DateTime<Utc> }
45+
{ BoolArray => Vec<bool> }
4446
{ Int32Array => Vec<i32> }
4547
{ Int64Array => Vec<i64> }
4648
{ UInt32Array => Vec<u32> }

connectorx/src/transports/postgres_arrow2.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ macro_rules! impl_postgres_transport {
6161
{ ByteA[Vec<u8>] => LargeBinary[Vec<u8>] | conversion auto }
6262
{ JSON[Value] => LargeUtf8[String] | conversion option }
6363
{ JSONB[Value] => LargeUtf8[String] | conversion none }
64+
{ BoolArray[Vec<bool>] => BoolArray[Vec<bool>] | conversion auto_vec }
6465
{ Int2Array[Vec<i16>] => Int64Array[Vec<i64>] | conversion auto_vec }
6566
{ Int4Array[Vec<i32>] => Int64Array[Vec<i64>] | conversion auto_vec }
6667
{ Int8Array[Vec<i64>] => Int64Array[Vec<i64>] | conversion auto }

0 commit comments

Comments
 (0)