feat: support oid & name type for postgres

datafusion-contrib · Dec 23, 2024 · 08fdab4 · 08fdab4
1 parent 00a2af6
commit 08fdab4
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/src/sql/arrow_sql_gen/postgres.rs b/src/sql/arrow_sql_gen/postgres.rs
@@ -255,6 +255,9 @@ pub fn rows_to_arrow(rows: &[Row], projected_schema: &Option<SchemaRef>) -> Resu
                 Type::INT8 => {
                     handle_primitive_type!(builder, Type::INT8, Int64Builder, i64, row, i);
                 }
+                Type::OID => {
+                    handle_primitive_type!(builder, Type::OID, UInt32Builder, u32, row, i);
+                }
                 Type::FLOAT4 => {
                     handle_primitive_type!(builder, Type::FLOAT4, Float32Builder, f32, row, i);
                 }
@@ -267,6 +270,9 @@ pub fn rows_to_arrow(rows: &[Row], projected_schema: &Option<SchemaRef>) -> Resu
                 Type::VARCHAR => {
                     handle_primitive_type!(builder, Type::VARCHAR, StringBuilder, &str, row, i);
                 }
+                Type::NAME => {
+                    handle_primitive_type!(builder, Type::NAME, StringBuilder, &str, row, i);
+                }
                 Type::BYTEA => {
                     handle_primitive_type!(builder, Type::BYTEA, BinaryBuilder, Vec<u8>, row, i);
                 }
@@ -835,9 +841,10 @@ fn map_column_type_to_data_type(column_type: &Type) -> Option<DataType> {
         Type::INT2 => Some(DataType::Int16),
         Type::INT4 => Some(DataType::Int32),
         Type::INT8 | Type::MONEY => Some(DataType::Int64),
+        Type::OID => Some(DataType::UInt32),
         Type::FLOAT4 => Some(DataType::Float32),
         Type::FLOAT8 => Some(DataType::Float64),
-        Type::TEXT | Type::VARCHAR | Type::BPCHAR | Type::UUID => Some(DataType::Utf8),
+        Type::TEXT | Type::VARCHAR | Type::BPCHAR | Type::UUID | Type::NAME => Some(DataType::Utf8),
         Type::BYTEA => Some(DataType::Binary),
         Type::BOOL => Some(DataType::Boolean),
         Type::JSON => Some(DataType::LargeUtf8),

diff --git a/src/sql/arrow_sql_gen/postgres/schema.rs b/src/sql/arrow_sql_gen/postgres/schema.rs
@@ -14,15 +14,15 @@ pub(crate) fn pg_data_type_to_arrow_type(
         "smallint" => Ok(DataType::Int16),
         "integer" | "int" | "int4" => Ok(DataType::Int32),
         "bigint" | "int8" | "money" => Ok(DataType::Int64),
+        "oid" => Ok(DataType::UInt32),
         "numeric" | "decimal" => {
             let (precision, scale) = parse_numeric_type(pg_type)?;
             Ok(DataType::Decimal128(precision, scale))
         }
         "real" | "float4" => Ok(DataType::Float32),
         "double precision" | "float8" => Ok(DataType::Float64),
-        "character" | "char" | "character varying" | "varchar" | "text" | "bpchar" | "uuid" => {
-            Ok(DataType::Utf8)
-        }
+        "character" | "char" | "character varying" | "varchar" | "text" | "bpchar" | "uuid"
+        | "name" => Ok(DataType::Utf8),
         "bytea" => Ok(DataType::Binary),
         "date" => Ok(DataType::Date32),
         "time" | "time without time zone" => Ok(DataType::Time64(TimeUnit::Nanosecond)),
@@ -196,6 +196,10 @@ mod tests {
             pg_data_type_to_arrow_type("bigint", None).expect("Failed to convert bigint"),
             DataType::Int64
         );
+        assert_eq!(
+            pg_data_type_to_arrow_type("oid", None).expect("Failed to convert oid"),
+            DataType::UInt32
+        );
         assert_eq!(
             pg_data_type_to_arrow_type("real", None).expect("Failed to convert real"),
             DataType::Float32
@@ -224,6 +228,10 @@ mod tests {
             pg_data_type_to_arrow_type("text", None).expect("Failed to convert text"),
             DataType::Utf8
         );
+        assert_eq!(
+            pg_data_type_to_arrow_type("name", None).expect("Failed to convert name"),
+            DataType::Utf8
+        );
 
         // Test date/time types
         assert_eq!(