diff --git a/connectorx-python/connectorx/tests/test_arrow.py b/connectorx-python/connectorx/tests/test_arrow.py index 393bcf9ba4..d784d4f501 100644 --- a/connectorx-python/connectorx/tests/test_arrow.py +++ b/connectorx-python/connectorx/tests/test_arrow.py @@ -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) @@ -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) diff --git a/connectorx-python/connectorx/tests/test_postgres.py b/connectorx-python/connectorx/tests/test_postgres.py index 9d534c0b89..d99a4b8335 100644 --- a/connectorx-python/connectorx/tests/test_postgres.py +++ b/connectorx-python/connectorx/tests/test_postgres.py @@ -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) \ No newline at end of file diff --git a/connectorx-python/src/pandas/transports/postgres.rs b/connectorx-python/src/pandas/transports/postgres.rs index f08a52cf4d..c94d159130 100644 --- a/connectorx-python/src/pandas/transports/postgres.rs +++ b/connectorx-python/src/pandas/transports/postgres.rs @@ -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] | conversion option } { TimestampTz[DateTime] => DateTime[DateTime] | conversion auto } { Date[NaiveDate] => DateTime[DateTime] | conversion option } diff --git a/connectorx/src/sources/postgres/typesystem.rs b/connectorx/src/sources/postgres/typesystem.rs index b20bf96f9a..ebd4e0d2ef 100644 --- a/connectorx/src/sources/postgres/typesystem.rs +++ b/connectorx/src/sources/postgres/typesystem.rs @@ -36,6 +36,7 @@ pub enum PostgresTypeSystem { JSONB(bool), Enum(bool), HSTORE(bool), + Name(bool), } impl_typesystem! { @@ -56,8 +57,7 @@ impl_typesystem! { { VarcharArray | TextArray => Vec} { Bool => bool } { Char => i8 } - { Text | BpChar | VarChar | Enum => &'r str } - { ByteA => Vec } + { Text | BpChar | VarChar | Enum | Name => &'r str } { ByteA => Vec } { Time => NaiveTime } { Timestamp => NaiveDateTime } { TimestampTz => DateTime } @@ -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), diff --git a/connectorx/src/transports/postgres_arrow.rs b/connectorx/src/transports/postgres_arrow.rs index cda6ca4fc5..73c076fe84 100644 --- a/connectorx/src/transports/postgres_arrow.rs +++ b/connectorx/src/transports/postgres_arrow.rs @@ -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 } diff --git a/connectorx/src/transports/postgres_arrow2.rs b/connectorx/src/transports/postgres_arrow2.rs index 4d50ce412c..73d3b24c74 100644 --- a/connectorx/src/transports/postgres_arrow2.rs +++ b/connectorx/src/transports/postgres_arrow2.rs @@ -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 } diff --git a/connectorx/tests/test_polars.rs b/connectorx/tests/test_polars.rs index f528fb0047..4b56515a7b 100644 --- a/connectorx/tests/test_polars.rs +++ b/connectorx/tests/test_polars.rs @@ -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::::new(config, NoTls, 2).unwrap(); + let mut destination = Arrow2Destination::new(); + let dispatcher = Dispatcher::<_, _, PostgresArrow2Transport>::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); +} diff --git a/scripts/postgres.sql b/scripts/postgres.sql index 685ff4614c..134259970f 100644 --- a/scripts/postgres.sql +++ b/scripts/postgres.sql @@ -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