postg_arrow/
types.rs

1// Type mapping and conversions between Arrow and PostgreSQL types
2
3use arrow::datatypes::DataType;
4use postgres_types::Type;
5use anyhow::{anyhow, Result};
6
7pub fn pg_type_enum_to_arrow(pg_type: &Type) -> Result<DataType> {
8    match *pg_type {
9        Type::INT2 => Ok(DataType::Int16),
10        Type::INT4 => Ok(DataType::Int32),
11        Type::INT8 => Ok(DataType::Int64),
12        Type::FLOAT4 => Ok(DataType::Float32),
13        Type::FLOAT8 => Ok(DataType::Float64),
14        Type::BOOL => Ok(DataType::Boolean),
15        Type::TEXT | Type::VARCHAR => Ok(DataType::LargeUtf8),
16        Type::BYTEA => Ok(DataType::LargeBinary),
17        Type::DATE => Ok(DataType::LargeUtf8), // Downgraded
18        Type::TIME => Ok(DataType::LargeUtf8), // Downgraded
19        Type::TIMESTAMP => Ok(DataType::LargeUtf8), // Downgraded
20        Type::TIMESTAMPTZ => Ok(DataType::LargeUtf8), // Downgraded
21        Type::NUMERIC => Ok(DataType::LargeUtf8), // Downgraded
22        Type::UUID => Ok(DataType::LargeUtf8), // Mapped to string for broader compatibility
23        Type::JSON | Type::JSONB => Ok(DataType::LargeUtf8),
24        _ => Err(anyhow!("Unsupported PostgreSQL type OID: {}", pg_type.oid())),
25    }
26}
27
28pub fn pg_type_to_arrow(oid: u32) -> Result<DataType> {
29    let pg_type = Type::from_oid(oid).ok_or_else(|| anyhow!("Unsupported PostgreSQL type OID: {}", oid))?;
30    pg_type_enum_to_arrow(&pg_type)
31}
32
33pub fn pg_oid_to_arrow(oid: u32) -> Result<DataType> {
34    pg_type_to_arrow(oid)
35}
36
37pub fn arrow_to_pg_type(dt: &DataType) -> Result<String> {
38    match dt {
39        DataType::Int16 => Ok("INT2".to_string()),
40        DataType::Int32 => Ok("INT4".to_string()),
41        DataType::Int64 => Ok("INT8".to_string()),
42        DataType::Float32 => Ok("FLOAT4".to_string()),
43        DataType::Float64 => Ok("FLOAT8".to_string()),
44        DataType::Boolean => Ok("BOOL".to_string()),
45        DataType::Utf8 | DataType::LargeUtf8 => Ok("TEXT".to_string()),
46        DataType::Binary | DataType::LargeBinary => Ok("BYTEA".to_string()),
47        DataType::Date32 => Ok("DATE".to_string()),
48        DataType::Time64(_) => Ok("TIME".to_string()),
49        DataType::Timestamp(_, None) => Ok("TIMESTAMP".to_string()),
50        DataType::Timestamp(_, Some(_)) => Ok("TIMESTAMPTZ".to_string()),
51        DataType::Decimal128(_, _) => Ok("NUMERIC".to_string()),
52        _ => Err(anyhow!("Unsupported Arrow data type for PostgreSQL mapping: {:?}", dt)),
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use arrow::datatypes::TimeUnit;
59    use super::*;
60
61    #[test]
62    fn test_pg_type_enum_to_arrow_primitives() {
63        assert_eq!(pg_type_enum_to_arrow(&Type::INT2).unwrap(), DataType::Int16);
64        assert_eq!(pg_type_enum_to_arrow(&Type::INT4).unwrap(), DataType::Int32);
65        assert_eq!(pg_type_enum_to_arrow(&Type::INT8).unwrap(), DataType::Int64);
66        assert_eq!(pg_type_enum_to_arrow(&Type::FLOAT4).unwrap(), DataType::Float32);
67        assert_eq!(pg_type_enum_to_arrow(&Type::FLOAT8).unwrap(), DataType::Float64);
68        assert_eq!(pg_type_enum_to_arrow(&Type::BOOL).unwrap(), DataType::Boolean);
69    }
70
71    #[test]
72    fn test_pg_type_enum_to_arrow_strings_and_binary() {
73        assert_eq!(pg_type_enum_to_arrow(&Type::TEXT).unwrap(), DataType::LargeUtf8);
74        assert_eq!(pg_type_enum_to_arrow(&Type::VARCHAR).unwrap(), DataType::LargeUtf8);
75        assert_eq!(pg_type_enum_to_arrow(&Type::BYTEA).unwrap(), DataType::LargeBinary);
76    }
77
78    #[test]
79    fn test_pg_type_enum_to_arrow_temporal_and_complex() {
80        assert_eq!(pg_type_enum_to_arrow(&Type::DATE).unwrap(), DataType::LargeUtf8);
81        assert_eq!(
82            pg_type_enum_to_arrow(&Type::TIME).unwrap(),
83            DataType::LargeUtf8
84        );
85        assert_eq!(
86            pg_type_enum_to_arrow(&Type::TIMESTAMP).unwrap(),
87            DataType::LargeUtf8
88        );
89        assert_eq!(
90            pg_type_enum_to_arrow(&Type::TIMESTAMPTZ).unwrap(),
91            DataType::LargeUtf8
92        );
93        assert_eq!(
94            pg_type_enum_to_arrow(&Type::NUMERIC).unwrap(),
95            DataType::LargeUtf8
96        );
97        assert_eq!(pg_type_enum_to_arrow(&Type::UUID).unwrap(), DataType::LargeUtf8);
98        assert_eq!(pg_type_enum_to_arrow(&Type::JSON).unwrap(), DataType::LargeUtf8);
99        assert_eq!(pg_type_enum_to_arrow(&Type::JSONB).unwrap(), DataType::LargeUtf8);
100    }
101
102    #[test]
103    fn test_pg_type_enum_to_arrow_unsupported() {
104        assert!(pg_type_enum_to_arrow(&Type::POINT).is_err());
105        assert!(pg_type_enum_to_arrow(&Type::MONEY).is_err());
106    }
107
108    #[test]
109    fn test_pg_type_to_arrow_and_oid_alias() {
110        assert_eq!(pg_type_to_arrow(23).unwrap(), DataType::Int32); // INT4 OID is 23
111        assert_eq!(pg_oid_to_arrow(23).unwrap(), DataType::Int32);
112        assert!(pg_type_to_arrow(999999).is_err());
113        assert!(pg_oid_to_arrow(999999).is_err());
114    }
115
116    #[test]
117    fn test_arrow_to_pg_type_primitives() {
118        assert_eq!(arrow_to_pg_type(&DataType::Int16).unwrap(), "INT2");
119        assert_eq!(arrow_to_pg_type(&DataType::Int32).unwrap(), "INT4");
120        assert_eq!(arrow_to_pg_type(&DataType::Int64).unwrap(), "INT8");
121        assert_eq!(arrow_to_pg_type(&DataType::Float32).unwrap(), "FLOAT4");
122        assert_eq!(arrow_to_pg_type(&DataType::Float64).unwrap(), "FLOAT8");
123        assert_eq!(arrow_to_pg_type(&DataType::Boolean).unwrap(), "BOOL");
124    }
125
126    #[test]
127    fn test_arrow_to_pg_type_strings_and_binary() {
128        assert_eq!(arrow_to_pg_type(&DataType::Utf8).unwrap(), "TEXT");
129        assert_eq!(arrow_to_pg_type(&DataType::LargeUtf8).unwrap(), "TEXT");
130        assert_eq!(arrow_to_pg_type(&DataType::Binary).unwrap(), "BYTEA");
131        assert_eq!(arrow_to_pg_type(&DataType::LargeBinary).unwrap(), "BYTEA");
132    }
133
134    #[test]
135    fn test_arrow_to_pg_type_temporal_and_complex() {
136        assert_eq!(arrow_to_pg_type(&DataType::Date32).unwrap(), "DATE");
137        assert_eq!(
138            arrow_to_pg_type(&DataType::Time64(TimeUnit::Microsecond)).unwrap(),
139            "TIME"
140        );
141        assert_eq!(
142            arrow_to_pg_type(&DataType::Timestamp(TimeUnit::Microsecond, None)).unwrap(),
143            "TIMESTAMP"
144        );
145        assert_eq!(
146            arrow_to_pg_type(&DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into()))).unwrap(),
147            "TIMESTAMPTZ"
148        );
149        assert_eq!(
150            arrow_to_pg_type(&DataType::Decimal128(38, 9)).unwrap(),
151            "NUMERIC"
152        );
153    }
154
155    #[test]
156    fn test_arrow_to_pg_type_unsupported() {
157        assert!(arrow_to_pg_type(&DataType::Null).is_err());
158        assert!(arrow_to_pg_type(&DataType::Float16).is_err());
159        assert!(arrow_to_pg_type(&DataType::UInt32).is_err());
160    }
161}
postg_arrow/types.rs

postg_arrow/
types.rs
