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

datafusion_remote_table/
schema.rs
