fix: array type conversion in format_type

datafusion-postgres/export_pg_catalog_arrow.sh

@@ -173,7 +173,7 @@ def pg_type_to_arrow_type(pg_type, is_nullable=True):
         'character varying': pa.string(),
         'character': pa.string(),
         'name': pa.string(),
-        'oid': pa.uint32(),  # OIDs are unsigned
+        'oid': pa.int32(),  # OIDs are unsigned
         'regproc': pa.string(),
         'regtype': pa.string(),
         'regclass': pa.string(),

datafusion-postgres/src/pg_catalog/format_type.rs

@@ -2,10 +2,13 @@ use std::sync::Arc;
 
 use datafusion::{
     arrow::{
-        array::{Array, StringBuilder},
-        datatypes::DataType,
+        array::{Array, PrimitiveArray, StringBuilder},
+        datatypes::{DataType, Int32Type},
+    },
+    common::{
+        cast::{as_int32_array, as_int64_array},
+        DataFusionError,
     },
-    common::{cast::as_int32_array, DataFusionError},
     logical_expr::{
         ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature,
         Volatility,
@@ -14,12 +17,32 @@ use datafusion::{
 
 pub(crate) fn format_type_impl(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
     let args = ColumnarValue::values_to_arrays(args)?;
-    let type_oids = as_int32_array(&args[0])?;
+    let type_oids = match &args[0].data_type() {
+        DataType::Int32 => as_int32_array(&args[0])?,
+        DataType::Int64 => {
+            let oid_i64 = as_int64_array(&args[0])?;
+            &PrimitiveArray::<Int32Type>::from_iter(oid_i64.iter().map(|v| v.map(|x| x as i32)))
+        }
+        _ => {
+            return Err(DataFusionError::Internal(format!(
+                "Unexpected oid type in format_type: {}",
+                args[0].data_type()
+            )))
+        }
+    };
 
-    let typemods = if args.len() > 1 {
-        Some(as_int32_array(&args[1])?)
-    } else {
-        None
+    let typemods = match &args[1].data_type() {
+        DataType::Int32 => as_int32_array(&args[1])?,
+        DataType::Int64 => {
+            let oid_i64 = as_int64_array(&args[1])?;
+            &PrimitiveArray::<Int32Type>::from_iter(oid_i64.iter().map(|v| v.map(|x| x as i32)))
+        }
+        _ => {
+            return Err(DataFusionError::Internal(format!(
+                "Unexpected oid type in format_type: {}",
+                args[1].data_type()
+            )))
+        }
     };
 
     let mut result = StringBuilder::new();
@@ -31,9 +54,11 @@ pub(crate) fn format_type_impl(args: &[ColumnarValue]) -> Result<ColumnarValue,
         }
 
         let type_oid = type_oids.value(i);
-        let typemod = typemods
-            .map(|tm| if tm.is_null(i) { -1 } else { tm.value(i) })
-            .unwrap_or(-1);
+        let typemod = if typemods.is_null(i) {
+            -1
+        } else {
+            typemods.value(i)
+        };
 
         let formatted_type = format_postgres_type(type_oid, typemod);
         result.append_value(formatted_type);
@@ -211,8 +236,10 @@ impl FormatTypeUDF {
         Self {
             signature: Signature::one_of(
                 vec![
-                    TypeSignature::Exact(vec![DataType::Int64, DataType::Int32]),
                     TypeSignature::Exact(vec![DataType::Int32, DataType::Int32]),
+                    TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]),
+                    TypeSignature::Exact(vec![DataType::Int64, DataType::Int32]),
+                    TypeSignature::Exact(vec![DataType::Int64, DataType::Int64]),
                 ],
                 Volatility::Stable,
             ),

datafusion-postgres/src/pg_catalog/pg_attribute.rs

@@ -213,29 +213,29 @@ impl<C: CatalogInfo> PgAttributeTable<C> {
     /// Map DataFusion data types to PostgreSQL type information
     fn datafusion_to_pg_type(data_type: &DataType) -> (i32, i16, bool, &'static str, &'static str) {
         match data_type {
-            DataType::Boolean => (16, 1, true, "c", "p"),    // bool
-            DataType::Int8 => (18, 1, true, "c", "p"),       // char
-            DataType::Int16 => (21, 2, true, "s", "p"),      // int2
-            DataType::Int32 => (23, 4, true, "i", "p"),      // int4
-            DataType::Int64 => (20, 8, true, "d", "p"),      // int8
-            DataType::UInt8 => (21, 2, true, "s", "p"),      // Treat as int2
-            DataType::UInt16 => (23, 4, true, "i", "p"),     // Treat as int4
-            DataType::UInt32 => (20, 8, true, "d", "p"),     // Treat as int8
-            DataType::UInt64 => (1700, -1, false, "i", "m"), // Treat as numeric
-            DataType::Float32 => (700, 4, true, "i", "p"),   // float4
-            DataType::Float64 => (701, 8, true, "d", "p"),   // float8
-            DataType::Utf8 => (25, -1, false, "i", "x"),     // text
+            DataType::Boolean => (16, 1, true, "c", "p"),  // bool
+            DataType::Int8 => (18, 1, true, "c", "p"),     // char
+            DataType::Int16 => (21, 2, true, "s", "p"),    // int2
+            DataType::Int32 => (23, 4, true, "i", "p"),    // int4
+            DataType::Int64 => (20, 8, true, "d", "p"),    // int8
+            DataType::UInt8 => (18, 2, true, "s", "p"),    // char
+            DataType::UInt16 => (21, 4, true, "i", "p"),   // int2
+            DataType::UInt32 => (23, 8, true, "d", "p"),   // int4
+            DataType::UInt64 => (20, -1, false, "i", "m"), // int8
+            DataType::Float32 => (700, 4, true, "i", "p"), // float4
+            DataType::Float64 => (701, 8, true, "d", "p"), // float8
+            DataType::Utf8 => (25, -1, false, "i", "x"),   // text
             DataType::LargeUtf8 => (25, -1, false, "i", "x"), // text
-            DataType::Binary => (17, -1, false, "i", "x"),   // bytea
+            DataType::Binary => (17, -1, false, "i", "x"), // bytea
             DataType::LargeBinary => (17, -1, false, "i", "x"), // bytea
-            DataType::Date32 => (1082, 4, true, "i", "p"),   // date
-            DataType::Date64 => (1082, 4, true, "i", "p"),   // date
+            DataType::Date32 => (1082, 4, true, "i", "p"), // date
+            DataType::Date64 => (1082, 4, true, "i", "p"), // date
             DataType::Time32(_) => (1083, 8, true, "d", "p"), // time
             DataType::Time64(_) => (1083, 8, true, "d", "p"), // time
             DataType::Timestamp(_, _) => (1114, 8, true, "d", "p"), // timestamp
             DataType::Decimal128(_, _) => (1700, -1, false, "i", "m"), // numeric
             DataType::Decimal256(_, _) => (1700, -1, false, "i", "m"), // numeric
-            _ => (25, -1, false, "i", "x"),                  // Default to text for unknown types
+            _ => (25, -1, false, "i", "x"),                // Default to text for unknown types
         }
     }
 }