datafusion_remote_table/
schema.rs

1use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
2use std::sync::Arc;
3
4#[derive(Debug, Clone)]
5pub enum RemoteType {
6    Postgres(PostgresType),
7}
8
9impl RemoteType {
10    pub fn to_arrow_type(&self) -> DataType {
11        match self {
12            RemoteType::Postgres(postgres_type) => postgres_type.to_arrow_type(),
13        }
14    }
15}
16
17#[derive(Debug, Clone)]
18pub enum PostgresType {
19    Bool,
20    Char,
21    Int2,
22    Int4,
23    Int8,
24    Float4,
25    Float8,
26    Text,
27    Varchar,
28    Bytea,
29    Date,
30    Timestamp,
31    TimestampTz,
32    Time,
33    Int2Array,
34    Int4Array,
35    Int8Array,
36    Float4Array,
37    Float8Array,
38    TextArray,
39    VarcharArray,
40    ByteaArray,
41    PostGisGeometry,
42}
43
44impl PostgresType {
45    pub fn to_arrow_type(&self) -> DataType {
46        match self {
47            PostgresType::Bool => DataType::Boolean,
48            PostgresType::Char => DataType::Utf8,
49            PostgresType::Int2 => DataType::Int16,
50            PostgresType::Int4 => DataType::Int32,
51            PostgresType::Int8 => DataType::Int64,
52            PostgresType::Float4 => DataType::Float32,
53            PostgresType::Float8 => DataType::Float64,
54            PostgresType::Text | PostgresType::Varchar => DataType::Utf8,
55            PostgresType::Bytea => DataType::Binary,
56            PostgresType::Date => DataType::Date32,
57            PostgresType::Timestamp => DataType::Timestamp(TimeUnit::Nanosecond, None),
58            PostgresType::TimestampTz => {
59                DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into()))
60            }
61            PostgresType::Time => DataType::Time64(TimeUnit::Nanosecond),
62            PostgresType::Int2Array => {
63                DataType::List(Arc::new(Field::new("", DataType::Int16, true)))
64            }
65            PostgresType::Int4Array => {
66                DataType::List(Arc::new(Field::new("", DataType::Int32, true)))
67            }
68            PostgresType::Int8Array => {
69                DataType::List(Arc::new(Field::new("", DataType::Int64, true)))
70            }
71            PostgresType::Float4Array => {
72                DataType::List(Arc::new(Field::new("", DataType::Float32, true)))
73            }
74            PostgresType::Float8Array => {
75                DataType::List(Arc::new(Field::new("", DataType::Float64, true)))
76            }
77            PostgresType::TextArray | PostgresType::VarcharArray => {
78                DataType::List(Arc::new(Field::new("", DataType::Utf8, true)))
79            }
80            PostgresType::ByteaArray => {
81                DataType::List(Arc::new(Field::new("", DataType::Binary, true)))
82            }
83            PostgresType::PostGisGeometry => DataType::Binary,
84        }
85    }
86}
87
88#[derive(Debug, Clone)]
89pub struct RemoteField {
90    pub name: String,
91    pub remote_type: RemoteType,
92    pub nullable: bool,
93}
94
95impl RemoteField {
96    pub fn new(name: impl Into<String>, remote_type: RemoteType, nullable: bool) -> Self {
97        RemoteField {
98            name: name.into(),
99            remote_type,
100            nullable,
101        }
102    }
103
104    pub fn to_arrow_field(&self) -> Field {
105        Field::new(
106            self.name.clone(),
107            self.remote_type.to_arrow_type(),
108            self.nullable,
109        )
110    }
111}
112
113#[derive(Debug, Clone)]
114pub struct RemoteSchema {
115    pub fields: Vec<RemoteField>,
116}
117
118impl RemoteSchema {
119    pub fn empty() -> Self {
120        RemoteSchema { fields: vec![] }
121    }
122    pub fn new(fields: Vec<RemoteField>) -> Self {
123        RemoteSchema { fields }
124    }
125
126    pub fn to_arrow_schema(&self) -> Schema {
127        let mut fields = vec![];
128        for remote_field in self.fields.iter() {
129            fields.push(remote_field.to_arrow_field());
130        }
131        Schema::new(fields)
132    }
133}
datafusion_remote_table/schema.rs

datafusion_remote_table/
schema.rs
