Skip to content

Commit

Permalink
Merge pull request #504 from Jordan-M-Young/feat/pg-name-support
Browse files Browse the repository at this point in the history 
Feat/pg name support
  • Loading branch information
@wangxiaoying
wangxiaoying authored Jul 3, 2023
2 parents e735e73 + 6172e71 commit f1d030e
Show file tree
Hide file tree
Showing 8 changed files with 67 additions and 10 deletions.
10 changes: 8 additions & 2 deletions connectorx-python/connectorx/tests/test_arrow.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ def test_arrow2(postgres_url: str) -> None:


def test_arrow2_type(postgres_url: str) -> None:
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_json, test_jsonb, test_f4array, test_f8array, test_narray, test_i2array, test_i4array, test_i8array, test_enum, test_ltree FROM test_types"
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_json, test_jsonb, test_f4array, test_f8array, test_narray, test_i2array, test_i4array, test_i8array, test_enum, test_ltree, test_name FROM test_types"
df = read_sql(postgres_url, query, return_type="arrow2")
df = df.to_pandas(date_as_object=False)
df.sort_values(by="test_int16", inplace=True, ignore_index=True)
Expand Down Expand Up @@ -181,7 +181,13 @@ def test_arrow2_type(postgres_url: str) -> None:
"test_enum": pd.Series(
["happy", "very happy", "ecstatic", None], dtype="object"
),
"test_ltree": pd.Series(["A.B.C.D", "A.B.E", "A", None], dtype="object"),
"test_ltree": pd.Series(
["A.B.C.D", "A.B.E", "A", None], dtype="object"
),
"test_name": pd.Series(
["0", "21", "someName", "101203203-1212323-22131235"]
)

},
)
assert_frame_equal(df, expected, check_names=True)
11 changes: 11 additions & 0 deletions connectorx-python/connectorx/tests/test_postgres.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1116,3 +1116,14 @@ def test_postgres_tls_fail(postgres_url_tls: str) -> None:
partition_range=(0, 2000),
partition_num=3,
)

def test_postgres_name_type(postgres_url: str) -> None:
# partition column can not have None
query = "SELECT test_name FROM test_types"
df = read_sql(postgres_url, query)
expected = pd.DataFrame(
data={
"test_name": pd.Series(["0", "21", "someName", "101203203-1212323-22131235"]),
},
)
assert_frame_equal(df, expected, check_names=True)
1 change: 1 addition & 0 deletions connectorx-python/src/pandas/transports/postgres.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ macro_rules! impl_postgres_transport {
{ Text[&'r str] => Str[&'r str] | conversion auto }
{ BpChar[&'r str] => Str[&'r str] | conversion none }
{ VarChar[&'r str] => Str[&'r str] | conversion none }
{ Name[&'r str] => Str[&'r str] | conversion none }
{ Timestamp[NaiveDateTime] => DateTime[DateTime<Utc>] | conversion option }
{ TimestampTz[DateTime<Utc>] => DateTime[DateTime<Utc>] | conversion auto }
{ Date[NaiveDate] => DateTime[DateTime<Utc>] | conversion option }
Expand Down
6 changes: 3 additions & 3 deletions connectorx/src/sources/postgres/typesystem.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ pub enum PostgresTypeSystem {
JSONB(bool),
Enum(bool),
HSTORE(bool),
Name(bool),
}

impl_typesystem! {
Expand All @@ -56,8 +57,7 @@ impl_typesystem! {
{ VarcharArray | TextArray => Vec<String>}
{ Bool => bool }
{ Char => i8 }
{ Text | BpChar | VarChar | Enum => &'r str }
{ ByteA => Vec<u8> }
{ Text | BpChar | VarChar | Enum | Name => &'r str } { ByteA => Vec<u8> }
{ Time => NaiveTime }
{ Timestamp => NaiveDateTime }
{ TimestampTz => DateTime<Utc> }
Expand Down Expand Up @@ -88,7 +88,7 @@ impl<'a> From<&'a Type> for PostgresTypeSystem {
"_text" => TextArray(true),
"bool" => Bool(true),
"char" => Char(true),
"text" | "citext" | "ltree" | "lquery" | "ltxtquery" => Text(true),
"text" | "citext" | "ltree" | "lquery" | "ltxtquery" | "name" => Text(true),
"bpchar" => BpChar(true),
"varchar" => VarChar(true),
"bytea" => ByteA(true),
Expand Down
1 change: 1 addition & 0 deletions connectorx/src/transports/postgres_arrow.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ macro_rules! impl_postgres_transport {
{ Text[&'r str] => LargeUtf8[String] | conversion owned }
{ BpChar[&'r str] => LargeUtf8[String] | conversion none }
{ VarChar[&'r str] => LargeUtf8[String] | conversion none }
{ Name[&'r str] => LargeUtf8[String] | conversion none }
{ Timestamp[NaiveDateTime] => Date64[NaiveDateTime] | conversion auto }
{ Date[NaiveDate] => Date32[NaiveDate] | conversion auto }
{ Time[NaiveTime] => Time64[NaiveTime] | conversion auto }
Expand Down
1 change: 1 addition & 0 deletions connectorx/src/transports/postgres_arrow2.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ macro_rules! impl_postgres_transport {
{ BpChar[&'r str] => LargeUtf8[String] | conversion none }
{ VarChar[&'r str] => LargeUtf8[String] | conversion none }
{ Enum[&'r str] => LargeUtf8[String] | conversion none }
{ Name[&'r str] => LargeUtf8[String] | conversion none }
{ Timestamp[NaiveDateTime] => Date64[NaiveDateTime] | conversion auto }
{ Date[NaiveDate] => Date32[NaiveDate] | conversion auto }
{ Time[NaiveTime] => Time64[NaiveTime] | conversion auto }
Expand Down
36 changes: 36 additions & 0 deletions connectorx/tests/test_polars.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -224,3 +224,39 @@ fn test_pg_pl_text_array() {
println!("{:?}", df);
assert_eq!(df, test_df);
}

#[test]

fn test_pg_pl_name() {
let _ = env_logger::builder().is_test(true).try_init();

let dburl = env::var("POSTGRES_URL").unwrap();

let queries = [CXQuery::naked("select test_name from test_types")];
let url = Url::parse(dburl.as_str()).unwrap();
let (config, _tls) = rewrite_tls_args(&url).unwrap();
let builder = PostgresSource::<BinaryProtocol, NoTls>::new(config, NoTls, 2).unwrap();
let mut destination = Arrow2Destination::new();
let dispatcher = Dispatcher::<_, _, PostgresArrow2Transport<BinaryProtocol, NoTls>>::new(
builder,
&mut destination,
&queries,
Some(format!("select * from test_types")),
);

dispatcher.run().expect("run dispatcher");

let s1 = "0";
let s2 = "21";
let s3 = "someName";
let s4 = "101203203-1212323-22131235";

let df: DataFrame = destination.polars().unwrap();
let test_df: DataFrame = df!(
"test_name" => &[s1,s2,s3,s4]
)
.unwrap();

println!("{:?}", df);
assert_eq!(df, test_df);
}
11 changes: 6 additions & 5 deletions scripts/postgres.sql
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,13 +68,14 @@ CREATE TABLE IF NOT EXISTS test_types(
test_lquery lquery,
test_ltxtquery ltxtquery,
test_varchararray VARCHAR[],
test_textarray TEXT[]
test_textarray TEXT[],
test_name NAME
);

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*',ARRAY['str1','str2'],ARRAY['text1','text2']);
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','{"0123456789","abcdefghijklmnopqrstuvwxyz","!@#$%^&*()_-+=~`:;<>?/"}','{"0123456789","abcdefghijklmnopqrstuvwxyz","!@#$%^&*()_-+=~`:;<>?/"}');
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@',ARRAY['',' '],ARRAY['',' ']);
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,'{}','{}');
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*',ARRAY['str1','str2'],ARRAY['text1','text2'],'0');
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','{"0123456789","abcdefghijklmnopqrstuvwxyz","!@#$%^&*()_-+=~`:;<>?/"}','{"0123456789","abcdefghijklmnopqrstuvwxyz","!@#$%^&*()_-+=~`:;<>?/"}','21');
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@',ARRAY['',' '],ARRAY['',' '],'someName');
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,'{}','{}','101203203-1212323-22131235');

CREATE OR REPLACE FUNCTION increment(i integer) RETURNS integer AS $$
BEGIN
Expand Down

0 comments on commit f1d030e

Please sign in to comment.