arrow_pg/
datatypes.rs

1use std::sync::Arc;
2
3#[cfg(not(feature = "datafusion"))]
4use arrow::{datatypes::*, record_batch::RecordBatch};
5#[cfg(feature = "datafusion")]
6use datafusion::arrow::{datatypes::*, record_batch::RecordBatch};
7
8use pgwire::api::portal::Format;
9use pgwire::api::results::FieldInfo;
10use pgwire::api::Type;
11use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
12use pgwire::messages::data::DataRow;
13use postgres_types::Kind;
14
15use crate::row_encoder::RowEncoder;
16
17#[cfg(feature = "datafusion")]
18pub mod df;
19
20pub fn into_pg_type(arrow_type: &DataType) -> PgWireResult<Type> {
21    Ok(match arrow_type {
22        DataType::Null => Type::UNKNOWN,
23        DataType::Boolean => Type::BOOL,
24        DataType::Int8 | DataType::UInt8 => Type::CHAR,
25        DataType::Int16 | DataType::UInt16 => Type::INT2,
26        DataType::Int32 | DataType::UInt32 => Type::INT4,
27        DataType::Int64 | DataType::UInt64 => Type::INT8,
28        DataType::Timestamp(_, tz) => {
29            if tz.is_some() {
30                Type::TIMESTAMPTZ
31            } else {
32                Type::TIMESTAMP
33            }
34        }
35        DataType::Time32(_) | DataType::Time64(_) => Type::TIME,
36        DataType::Date32 | DataType::Date64 => Type::DATE,
37        DataType::Interval(_) => Type::INTERVAL,
38        DataType::Binary
39        | DataType::FixedSizeBinary(_)
40        | DataType::LargeBinary
41        | DataType::BinaryView => Type::BYTEA,
42        DataType::Float16 | DataType::Float32 => Type::FLOAT4,
43        DataType::Float64 => Type::FLOAT8,
44        DataType::Decimal128(_, _) => Type::NUMERIC,
45        DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Type::TEXT,
46        DataType::List(field) | DataType::FixedSizeList(field, _) | DataType::LargeList(field) => {
47            match field.data_type() {
48                DataType::Boolean => Type::BOOL_ARRAY,
49                DataType::Int8 | DataType::UInt8 => Type::CHAR_ARRAY,
50                DataType::Int16 | DataType::UInt16 => Type::INT2_ARRAY,
51                DataType::Int32 | DataType::UInt32 => Type::INT4_ARRAY,
52                DataType::Int64 | DataType::UInt64 => Type::INT8_ARRAY,
53                DataType::Timestamp(_, tz) => {
54                    if tz.is_some() {
55                        Type::TIMESTAMPTZ_ARRAY
56                    } else {
57                        Type::TIMESTAMP_ARRAY
58                    }
59                }
60                DataType::Time32(_) | DataType::Time64(_) => Type::TIME_ARRAY,
61                DataType::Date32 | DataType::Date64 => Type::DATE_ARRAY,
62                DataType::Interval(_) => Type::INTERVAL_ARRAY,
63                DataType::FixedSizeBinary(_)
64                | DataType::Binary
65                | DataType::LargeBinary
66                | DataType::BinaryView => Type::BYTEA_ARRAY,
67                DataType::Float16 | DataType::Float32 => Type::FLOAT4_ARRAY,
68                DataType::Float64 => Type::FLOAT8_ARRAY,
69                DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Type::TEXT_ARRAY,
70                struct_type @ DataType::Struct(_) => Type::new(
71                    Type::RECORD_ARRAY.name().into(),
72                    Type::RECORD_ARRAY.oid(),
73                    Kind::Array(into_pg_type(struct_type)?),
74                    Type::RECORD_ARRAY.schema().into(),
75                ),
76                list_type => {
77                    return Err(PgWireError::UserError(Box::new(ErrorInfo::new(
78                        "ERROR".to_owned(),
79                        "XX000".to_owned(),
80                        format!("Unsupported List Datatype {list_type}"),
81                    ))));
82                }
83            }
84        }
85        DataType::Dictionary(_, value_type) => into_pg_type(value_type)?,
86        DataType::Struct(fields) => {
87            let name: String = fields
88                .iter()
89                .map(|x| x.name().clone())
90                .reduce(|a, b| a + ", " + &b)
91                .map(|x| format!("({x})"))
92                .unwrap_or("()".to_string());
93            let kind = Kind::Composite(
94                fields
95                    .iter()
96                    .map(|x| {
97                        into_pg_type(x.data_type())
98                            .map(|_type| postgres_types::Field::new(x.name().clone(), _type))
99                    })
100                    .collect::<Result<Vec<_>, PgWireError>>()?,
101            );
102            Type::new(name, Type::RECORD.oid(), kind, Type::RECORD.schema().into())
103        }
104        _ => {
105            return Err(PgWireError::UserError(Box::new(ErrorInfo::new(
106                "ERROR".to_owned(),
107                "XX000".to_owned(),
108                format!("Unsupported Datatype {arrow_type}"),
109            ))));
110        }
111    })
112}
113
114pub fn arrow_schema_to_pg_fields(schema: &Schema, format: &Format) -> PgWireResult<Vec<FieldInfo>> {
115    schema
116        .fields()
117        .iter()
118        .enumerate()
119        .map(|(idx, f)| {
120            let pg_type = into_pg_type(f.data_type())?;
121            Ok(FieldInfo::new(
122                f.name().into(),
123                None,
124                None,
125                pg_type,
126                format.format_for(idx),
127            ))
128        })
129        .collect::<PgWireResult<Vec<FieldInfo>>>()
130}
131
132pub fn encode_recordbatch(
133    fields: Arc<Vec<FieldInfo>>,
134    record_batch: RecordBatch,
135) -> Box<impl Iterator<Item = PgWireResult<DataRow>>> {
136    let mut row_stream = RowEncoder::new(record_batch, fields);
137    Box::new(std::iter::from_fn(move || row_stream.next_row()))
138}
arrow_pg/datatypes.rs

arrow_pg/
datatypes.rs
