Skip to content

Commit ba5c31e

Browse files
Amar1729Amar1729
committed
postgres add boolean array
1 parent 1d2ad11 commit ba5c31e

File tree

5 files changed

+42
-8
lines changed

5 files changed

+42
-8
lines changed

‎connectorx-python/connectorx/tests/test_postgres.py‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -450,7 +450,7 @@ def test_postgres_with_index_col(postgres_url: str) -> None:
450450

451451

452452
def test_postgres_types_binary(postgres_url: str) -> None:
453-
query = "SELECT test_date, test_timestamp, test_timestamptz, test_int16, test_int64, test_float32, test_numeric, test_bpchar, test_char, test_varchar, test_uuid, test_time, test_json, test_jsonb, test_bytea, test_enum, test_f4array, test_f8array, test_narray, test_i2array, test_i4array, test_i8array, test_citext, test_ltree, test_lquery, test_ltxtquery FROM test_types"
453+
query = "SELECT test_date, test_timestamp, test_timestamptz, test_int16, test_int64, test_float32, test_numeric, test_bpchar, test_char, test_varchar, test_uuid, test_time, test_json, test_jsonb, test_bytea, test_enum, test_f4array, test_f8array, test_narray, test_boolarray, test_i2array, test_i4array, test_i8array, test_citext, test_ltree, test_lquery, test_ltxtquery FROM test_types"
454454
df = read_sql(postgres_url, query)
455455
expected = pd.DataFrame(
456456
index=range(4),
@@ -538,6 +538,9 @@ def test_postgres_types_binary(postgres_url: str) -> None:
538538
"test_narray": pd.Series(
539539
[[], None, [521.34], [0.12, 333.33, 22.22]], dtype="object"
540540
),
541+
"test_boolarray": pd.Series(
542+
[[True, True], [], [False, False], None], dtype="object",
543+
),
541544
"test_i2array": pd.Series(
542545
[[-1, 0, 1], [], [-32768, 32767], None], dtype="object"
543546
),
@@ -560,7 +563,7 @@ def test_postgres_types_binary(postgres_url: str) -> None:
560563

561564

562565
def test_postgres_types_csv(postgres_url: str) -> None:
563-
query = "SELECT test_date, test_timestamp, test_timestamptz, test_int16, test_int64, test_float32, test_numeric, test_bpchar, test_char, test_varchar, test_uuid, test_time, test_json, test_jsonb, test_bytea, test_enum::text, test_f4array, test_f8array, test_narray, test_i2array, test_i4array, test_i8array, test_citext, test_ltree FROM test_types"
566+
query = "SELECT test_date, test_timestamp, test_timestamptz, test_int16, test_int64, test_float32, test_numeric, test_bpchar, test_char, test_varchar, test_uuid, test_time, test_json, test_jsonb, test_bytea, test_enum::text, test_f4array, test_f8array, test_narray, test_boolarray, test_i2array, test_i4array, test_i8array, test_citext, test_ltree FROM test_types"
564567
df = read_sql(postgres_url, query, protocol="csv")
565568
expected = pd.DataFrame(
566569
index=range(4),
@@ -648,6 +651,9 @@ def test_postgres_types_csv(postgres_url: str) -> None:
648651
"test_narray": pd.Series(
649652
[[], None, [521.34], [0.12, 333.33, 22.22]], dtype="object"
650653
),
654+
"test_boolarray": pd.Series(
655+
[[True, True], [], [False, False], None], dtype="object",
656+
),
651657
"test_i2array": pd.Series(
652658
[[-1, 0, 1], [], [-32768, 32767], None], dtype="object"
653659
),
@@ -666,7 +672,7 @@ def test_postgres_types_csv(postgres_url: str) -> None:
666672

667673

668674
def test_postgres_types_cursor(postgres_url: str) -> None:
669-
query = "SELECT test_date, test_timestamp, test_timestamptz, test_int16, test_int64, test_float32, test_numeric, test_bpchar, test_char, test_varchar, test_uuid, test_time, test_json, test_jsonb, test_bytea, test_enum::text, test_f4array, test_f8array, test_narray, test_i2array, test_i4array, test_i8array, test_citext, test_ltree FROM test_types"
675+
query = "SELECT test_date, test_timestamp, test_timestamptz, test_int16, test_int64, test_float32, test_numeric, test_bpchar, test_char, test_varchar, test_uuid, test_time, test_json, test_jsonb, test_bytea, test_enum::text, test_f4array, test_f8array, test_narray, test_boolarray, test_i2array, test_i4array, test_i8array, test_citext, test_ltree FROM test_types"
670676
df = read_sql(postgres_url, query, protocol="cursor")
671677
expected = pd.DataFrame(
672678
index=range(4),
@@ -754,6 +760,9 @@ def test_postgres_types_cursor(postgres_url: str) -> None:
754760
"test_narray": pd.Series(
755761
[[], None, [521.34], [0.12, 333.33, 22.22]], dtype="object"
756762
),
763+
"test_boolarray": pd.Series(
764+
[[True, True], [], [False, False], None], dtype="object",
765+
),
757766
"test_i2array": pd.Series(
758767
[[-1, 0, 1], [], [-32768, 32767], None], dtype="object"
759768
),
@@ -772,7 +781,7 @@ def test_postgres_types_cursor(postgres_url: str) -> None:
772781

773782

774783
def test_postgres_types_simple(postgres_url: str) -> None:
775-
query = "SELECT test_date, test_timestamp, test_timestamptz, test_int16, test_int64, test_float32, test_numeric, test_bpchar, test_char, test_varchar, test_uuid, test_time, test_bytea, test_enum, test_f4array, test_f8array, test_narray, test_i2array, test_i4array, test_i8array FROM test_types"
784+
query = "SELECT test_date, test_timestamp, test_timestamptz, test_int16, test_int64, test_float32, test_numeric, test_bpchar, test_char, test_varchar, test_uuid, test_time, test_bytea, test_enum, test_f4array, test_f8array, test_narray, test_boolarray, test_i2array, test_i4array, test_i8array FROM test_types"
776785
df = read_sql(postgres_url, query, protocol="simple")
777786
expected = pd.DataFrame(
778787
index=range(4),
@@ -842,6 +851,9 @@ def test_postgres_types_simple(postgres_url: str) -> None:
842851
"test_narray": pd.Series(
843852
[[], None, [521.34], [0.12, 333.33, 22.22]], dtype="object"
844853
),
854+
"test_boolarray": pd.Series(
855+
[[True, True], [], [False, False], None], dtype="object",
856+
),
845857
"test_i2array": pd.Series(
846858
[[-1, 0, 1], [], [-32768, 32767], None], dtype="object"
847859
),

‎connectorx-python/src/pandas/destination.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,9 @@ impl<'a> Destination for PandasDestination<'a> {
175175
PandasBlockType::Int64Array => {
176176
self.allocate_array::<super::pandas_columns::PyList>(dt, placement)?;
177177
}
178+
PandasBlockType::BooleanArray => {
179+
self.allocate_array::<super::pandas_columns::PyList>(dt, placement)?;
180+
}
178181
PandasBlockType::String => {
179182
self.allocate_array::<PyString>(dt, placement)?;
180183
}
@@ -208,6 +211,17 @@ impl<'a> Destination for PandasDestination<'a> {
208211
.collect()
209212
}
210213
}
214+
PandasBlockType::BooleanArray => {
215+
let bblock = ArrayBlock::<bool>::extract(buf)?;
216+
let bcols = bblock.split()?;
217+
for (&cid, bcol) in block.cids.iter().zip_eq(bcols) {
218+
partitioned_columns[cid] = bcol
219+
.partition(counts)
220+
.into_iter()
221+
.map(|c| Box::new(c) as _)
222+
.collect()
223+
}
224+
}
211225
PandasBlockType::Float64 => {
212226
let fblock = Float64Block::extract(buf)?;
213227
let fcols = fblock.split()?;

‎connectorx-python/src/pandas/typesystem.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ pub enum PandasTypeSystem {
99
F64Array(bool),
1010
I64Array(bool),
1111
Bool(bool),
12+
BoolArray(bool),
1213
Char(bool),
1314
Str(bool),
1415
BoxStr(bool),
@@ -23,6 +24,7 @@ pub enum PandasBlockType {
2324
Boolean(bool), // bool indicates nullablity
2425
Int64(bool),
2526
Float64,
27+
BooleanArray,
2628
Int64Array,
2729
Float64Array,
2830
String,
@@ -54,6 +56,7 @@ impl From<PandasTypeSystem> for PandasBlockType {
5456
PandasTypeSystem::Bool(nullable) => PandasBlockType::Boolean(nullable),
5557
PandasTypeSystem::I64(nullable) => PandasBlockType::Int64(nullable),
5658
PandasTypeSystem::F64(_) => PandasBlockType::Float64,
59+
PandasTypeSystem::BoolArray(_) => PandasBlockType::BooleanArray,
5760
PandasTypeSystem::F64Array(_) => PandasBlockType::Float64Array,
5861
PandasTypeSystem::I64Array(_) => PandasBlockType::Int64Array,
5962
PandasTypeSystem::String(_)
@@ -74,6 +77,7 @@ impl_typesystem! {
7477
{ F64Array => Vec<f64> }
7578
{ I64Array => Vec<i64> }
7679
{ Bool => bool }
80+
{ BoolArray => Vec<bool> }
7781
{ Char => char }
7882
{ Str => &'r str }
7983
{ BoxStr => Box<str> }

‎connectorx/src/sources/postgres/typesystem.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ pub enum PostgresTypeSystem {
1414
Int2(bool),
1515
Int4(bool),
1616
Int8(bool),
17+
BoolArray(bool),
1718
Float4Array(bool),
1819
Float8Array(bool),
1920
NumericArray(bool),
@@ -45,6 +46,7 @@ impl_typesystem! {
4546
{ Float4 => f32 }
4647
{ Float8 => f64 }
4748
{ Numeric => Decimal }
49+
{ BoolArray => Vec<bool> }
4850
{ Int2Array => Vec<i16> }
4951
{ Int4Array => Vec<i32> }
5052
{ Int8Array => Vec<i64> }
@@ -75,6 +77,7 @@ impl<'a> From<&'a Type> for PostgresTypeSystem {
7577
"float4" => Float4(true),
7678
"float8" => Float8(true),
7779
"numeric" => Numeric(true),
80+
"_bool" => BoolArray(true),
7881
"_int2" => Int2Array(true),
7982
"_int4" => Int4Array(true),
8083
"_int8" => Int8Array(true),

‎scripts/postgres.sql‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ CREATE TABLE IF NOT EXISTS test_types(
6060
test_f4array REAL[],
6161
test_f8array DOUBLE PRECISION[],
6262
test_narray NUMERIC(5,2)[],
63+
test_boolarray BOOLEAN[],
6364
test_i2array SMALLINT[],
6465
test_i4array Integer[],
6566
test_i8array BIGINT[],
@@ -69,10 +70,10 @@ CREATE TABLE IF NOT EXISTS test_types(
6970
test_ltxtquery ltxtquery
7071
);
7172

72-
INSERT INTO test_types VALUES ('1970-01-01', '1970-01-01 00:00:01', '1970-01-01 00:00:01-00', 0, -9223372036854775808, NULL, NULL, 'a', 'a', NULL, '86b494cc-96b2-11eb-9298-3e22fbb9fe9d', '08:12:40', '1 year 2 months 3 days', '{"customer": "John Doe", "items": {"product": "Beer","qty": 6}}', '{"product": "Beer","qty": 6}', NULL, 'happy','{}', '{}', '{}', '{-1, 0, 1}', '{-1, 0, 1123}', '{-9223372036854775808, 9223372036854775807}', 'str_citext', 'A.B.C.D', '*.B.*', 'A & B*');
73-
INSERT INTO test_types VALUES ('2000-02-28', '2000-02-28 12:00:10', '2000-02-28 12:00:10-04', 1, 0, 3.1415926535, 521.34, 'bb', 'b', 'bb', '86b49b84-96b2-11eb-9298-3e22fbb9fe9d', NULL, '2 weeks ago', '{"customer": "Lily Bush", "items": {"product": "Diaper","qty": 24}}', '{"product": "Diaper","qty": 24}', 'Здра́вствуйте', 'very happy', NULL, NULL, NULL, '{}', '{}', '{}', '', 'A.B.E', 'A.*', 'A | B');
74-
INSERT INTO test_types VALUES ('2038-01-18', '2038-01-18 23:59:59', '2038-01-18 23:59:59+08', 2, 9223372036854775807, 2.71, '1e-130', 'ccc', NULL, 'c', '86b49c42-96b2-11eb-9298-3e22fbb9fe9d', '23:00:10', '3 months 2 days ago', '{"customer": "Josh William", "items": {"product": "Toy Car","qty": 1}}', '{"product": "Toy Car","qty": 1}', '', 'ecstatic', '{123.123}', '{-1e-307, 1e308}', '{521.34}', '{-32768, 32767}', '{-2147483648, 2147483647}', '{0}', 's', 'A', '*', 'A@');
75-
INSERT INTO test_types VALUES (NULL, NULL, NULL, 3, NULL, 0.00, -1e-37, NULL, 'd', 'defghijklm', NULL, '18:30:00', '3 year', NULL, NULL, '😜', NULL, '{-1e-37, 1e37}', '{0.000234, -12.987654321}', '{0.12, 333.33, 22.22}', NULL, NULL, NULL, NULL, NULL, NULL, NULL);
73+
INSERT INTO test_types VALUES ('1970-01-01', '1970-01-01 00:00:01', '1970-01-01 00:00:01-00', 0, -9223372036854775808, NULL, NULL, 'a', 'a', NULL, '86b494cc-96b2-11eb-9298-3e22fbb9fe9d', '08:12:40', '1 year 2 months 3 days', '{"customer": "John Doe", "items": {"product": "Beer","qty": 6}}', '{"product": "Beer","qty": 6}', NULL, 'happy','{}', '{}', '{}', '{TRUE,TRUE}', '{-1, 0, 1}', '{-1, 0, 1123}', '{-9223372036854775808, 9223372036854775807}', 'str_citext', 'A.B.C.D', '*.B.*', 'A & B*');
74+
INSERT INTO test_types VALUES ('2000-02-28', '2000-02-28 12:00:10', '2000-02-28 12:00:10-04', 1, 0, 3.1415926535, 521.34, 'bb', 'b', 'bb', '86b49b84-96b2-11eb-9298-3e22fbb9fe9d', NULL, '2 weeks ago', '{"customer": "Lily Bush", "items": {"product": "Diaper","qty": 24}}', '{"product": "Diaper","qty": 24}', 'Здра́вствуйте', 'very happy', NULL, NULL, NULL, '{}', '{}', '{}', '{}', '', 'A.B.E', 'A.*', 'A | B');
75+
INSERT INTO test_types VALUES ('2038-01-18', '2038-01-18 23:59:59', '2038-01-18 23:59:59+08', 2, 9223372036854775807, 2.71, '1e-130', 'ccc', NULL, 'c', '86b49c42-96b2-11eb-9298-3e22fbb9fe9d', '23:00:10', '3 months 2 days ago', '{"customer": "Josh William", "items": {"product": "Toy Car","qty": 1}}', '{"product": "Toy Car","qty": 1}', '', 'ecstatic', '{123.123}', '{-1e-307, 1e308}', '{521.34}', '{FALSE, FALSE}', '{-32768, 32767}', '{-2147483648, 2147483647}', '{0}', 's', 'A', '*', 'A@');
76+
INSERT INTO test_types VALUES (NULL, NULL, NULL, 3, NULL, 0.00, -1e-37, NULL, 'd', 'defghijklm', NULL, '18:30:00', '3 year', NULL, NULL, '😜', NULL, '{-1e-37, 1e37}', '{0.000234, -12.987654321}', '{0.12, 333.33, 22.22}', NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL);
7677

7778
CREATE OR REPLACE FUNCTION increment(i integer) RETURNS integer AS $$
7879
BEGIN

0 commit comments

Comments
 (0)