arrow_pg/
datatypes.rs

1use std::sync::Arc;
2
3#[cfg(not(feature = "datafusion"))]
4use arrow::{datatypes::*, record_batch::RecordBatch};
5#[cfg(feature = "postgis")]
6use arrow_schema::extension::ExtensionType;
7#[cfg(feature = "datafusion")]
8use datafusion::arrow::{datatypes::*, record_batch::RecordBatch};
9
10use pgwire::api::Type;
11use pgwire::api::portal::Format;
12use pgwire::api::results::FieldInfo;
13use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
14use pgwire::messages::data::DataRow;
15use pgwire::types::format::FormatOptions;
16use postgres_types::Kind;
17
18use crate::row_encoder::RowEncoder;
19
20#[cfg(feature = "datafusion")]
21pub mod df;
22
23pub fn into_pg_type(arrow_type: &DataType) -> PgWireResult<Type> {
24    let datatype = match arrow_type {
25        DataType::Null => Type::UNKNOWN,
26        DataType::Boolean => Type::BOOL,
27        DataType::Int8 => Type::INT2,
28        DataType::Int16 | DataType::UInt8 => Type::INT2,
29        DataType::Int32 | DataType::UInt16 => Type::INT4,
30        DataType::Int64 | DataType::UInt32 => Type::INT8,
31        DataType::UInt64 => Type::NUMERIC,
32        DataType::Timestamp(_, tz) => {
33            if tz.is_some() {
34                Type::TIMESTAMPTZ
35            } else {
36                Type::TIMESTAMP
37            }
38        }
39        DataType::Time32(_) | DataType::Time64(_) => Type::TIME,
40        DataType::Date32 | DataType::Date64 => Type::DATE,
41        DataType::Interval(_) | DataType::Duration(_) => Type::INTERVAL,
42        DataType::Binary
43        | DataType::FixedSizeBinary(_)
44        | DataType::LargeBinary
45        | DataType::BinaryView => Type::BYTEA,
46        DataType::Float16 | DataType::Float32 => Type::FLOAT4,
47        DataType::Float64 => Type::FLOAT8,
48        DataType::Decimal128(_, _) => Type::NUMERIC,
49        DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Type::TEXT,
50        DataType::List(field)
51        | DataType::FixedSizeList(field, _)
52        | DataType::LargeList(field)
53        | DataType::ListView(field)
54        | DataType::LargeListView(field) => match field.data_type() {
55            // Align with PostgreSQL: an array literal without type
56            // information such as `ARRAY[NULL]` has a `Null` element type and
57            // is reported as `text[]` by postgres.
58            DataType::Null => Type::TEXT_ARRAY,
59            DataType::Boolean => Type::BOOL_ARRAY,
60            DataType::Int8 => Type::INT2_ARRAY,
61            DataType::Int16 | DataType::UInt8 => Type::INT2_ARRAY,
62            DataType::Int32 | DataType::UInt16 => Type::INT4_ARRAY,
63            DataType::Int64 | DataType::UInt32 => Type::INT8_ARRAY,
64            DataType::UInt64 | DataType::Decimal128(_, _) => Type::NUMERIC_ARRAY,
65            DataType::Timestamp(_, tz) => {
66                if tz.is_some() {
67                    Type::TIMESTAMPTZ_ARRAY
68                } else {
69                    Type::TIMESTAMP_ARRAY
70                }
71            }
72            DataType::Time32(_) | DataType::Time64(_) => Type::TIME_ARRAY,
73            DataType::Date32 | DataType::Date64 => Type::DATE_ARRAY,
74            DataType::Interval(_) | DataType::Duration(_) => Type::INTERVAL_ARRAY,
75            DataType::FixedSizeBinary(_)
76            | DataType::Binary
77            | DataType::LargeBinary
78            | DataType::BinaryView => Type::BYTEA_ARRAY,
79            DataType::Float16 | DataType::Float32 => Type::FLOAT4_ARRAY,
80            DataType::Float64 => Type::FLOAT8_ARRAY,
81            DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Type::TEXT_ARRAY,
82            DataType::Struct(_) => Type::new(
83                Type::RECORD_ARRAY.name().into(),
84                Type::RECORD_ARRAY.oid(),
85                Kind::Array(field_into_pg_type(field)?),
86                Type::RECORD_ARRAY.schema().into(),
87            ),
88            list_type => {
89                return Err(PgWireError::UserError(Box::new(ErrorInfo::new(
90                    "ERROR".to_owned(),
91                    "XX000".to_owned(),
92                    format!("Unsupported List Datatype {list_type}"),
93                ))));
94            }
95        },
96        DataType::Dictionary(_, value_type) => into_pg_type(value_type.as_ref())?,
97        DataType::Struct(fields) => {
98            let name: String = fields
99                .iter()
100                .map(|x| x.name().clone())
101                .reduce(|a, b| a + ", " + &b)
102                .map(|x| format!("({x})"))
103                .unwrap_or("()".to_string());
104            let kind = Kind::Composite(
105                fields
106                    .iter()
107                    .map(|x| {
108                        field_into_pg_type(x)
109                            .map(|_type| postgres_types::Field::new(x.name().clone(), _type))
110                    })
111                    .collect::<Result<Vec<_>, PgWireError>>()?,
112            );
113            Type::new(name, Type::RECORD.oid(), kind, Type::RECORD.schema().into())
114        }
115        _ => {
116            return Err(PgWireError::UserError(Box::new(ErrorInfo::new(
117                "ERROR".to_owned(),
118                "XX000".to_owned(),
119                format!("Unsupported Datatype {arrow_type}"),
120            ))));
121        }
122    };
123
124    Ok(datatype)
125}
126
127pub fn field_into_pg_type(field: &Arc<Field>) -> PgWireResult<Type> {
128    let arrow_type = field.data_type();
129
130    match field.extension_type_name() {
131        // As of arrow 56, there are additional extension logical type that is
132        // defined using field metadata, for instance, json or geo.
133        //
134        // TODO: there is no fixed Geometry/Geography type id, here we use text
135        // for placeholder.
136        #[cfg(feature = "postgis")]
137        Some(geoarrow_schema::PointType::NAME) => Ok(Type::TEXT),
138        #[cfg(feature = "postgis")]
139        Some(geoarrow_schema::LineStringType::NAME) => Ok(Type::TEXT),
140        #[cfg(feature = "postgis")]
141        Some(geoarrow_schema::PolygonType::NAME) => Ok(Type::TEXT),
142        #[cfg(feature = "postgis")]
143        Some(geoarrow_schema::MultiPointType::NAME) => Ok(Type::TEXT),
144        #[cfg(feature = "postgis")]
145        Some(geoarrow_schema::MultiLineStringType::NAME) => Ok(Type::TEXT),
146        #[cfg(feature = "postgis")]
147        Some(geoarrow_schema::MultiPolygonType::NAME) => Ok(Type::TEXT),
148        #[cfg(feature = "postgis")]
149        Some(geoarrow_schema::GeometryCollectionType::NAME) => Ok(Type::TEXT),
150        #[cfg(feature = "postgis")]
151        Some(geoarrow_schema::GeometryType::NAME) => Ok(Type::TEXT),
152        #[cfg(feature = "postgis")]
153        Some(geoarrow_schema::RectType::NAME) => Ok(Type::TEXT),
154        #[cfg(feature = "postgis")]
155        Some(geoarrow_schema::WktType::NAME) => Ok(Type::TEXT),
156        #[cfg(feature = "postgis")]
157        Some(geoarrow_schema::WkbType::NAME) => Ok(Type::TEXT),
158
159        _ => into_pg_type(arrow_type),
160    }
161}
162
163pub fn arrow_schema_to_pg_fields(
164    schema: &Schema,
165    format: &Format,
166    data_format_options: Option<Arc<FormatOptions>>,
167) -> PgWireResult<Vec<FieldInfo>> {
168    let _ = data_format_options;
169    schema
170        .fields()
171        .iter()
172        .enumerate()
173        .map(|(idx, f)| {
174            let pg_type = field_into_pg_type(f)?;
175            let mut field_info =
176                FieldInfo::new(f.name().into(), None, None, pg_type, format.format_for(idx));
177            if let Some(data_format_options) = &data_format_options {
178                field_info = field_info.with_format_options(data_format_options.clone());
179            }
180
181            Ok(field_info)
182        })
183        .collect::<PgWireResult<Vec<FieldInfo>>>()
184}
185
186pub fn encode_recordbatch(
187    fields: Arc<Vec<FieldInfo>>,
188    record_batch: RecordBatch,
189) -> Box<impl Iterator<Item = PgWireResult<DataRow>>> {
190    let mut row_stream = RowEncoder::new(record_batch, fields);
191    Box::new(std::iter::from_fn(move || row_stream.next_row()))
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197
198    #[test]
199    fn null_list_is_text_array() {
200        // Align with PostgreSQL: `ARRAY[NULL]` has no element type
201        // information and is reported as `text[]`.
202        let ty = DataType::List(Arc::new(Field::new_list_field(DataType::Null, true)));
203        assert_eq!(into_pg_type(&ty).unwrap(), Type::TEXT_ARRAY);
204    }
205}
arrow_pg/datatypes.rs

arrow_pg/
datatypes.rs
