postg_arrow/
import.rs

1use crate::types::arrow_to_pg_type;
2use arrow::datatypes::{Schema, DataType};
3use arrow::record_batch::RecordBatch;
4use arrow::array::{
5    Int16Array, Int32Array, Int64Array, Float32Array, Float64Array, BooleanArray, LargeBinaryArray, LargeStringArray
6};
7use bytes::{BufMut, BytesMut};
8use futures::{Stream, StreamExt};
9use sqlx::{PgConnection, Executor};
10use anyhow::{anyhow, Result};
11
12/// Generates a PostgreSQL CREATE TABLE DDL statement from an Arrow Schema.
13pub fn generate_create_table_ddl(table: &str, schema: &Schema) -> Result<String> {
14    let mut column_defs = Vec::new();
15    for field in schema.fields() {
16        let pg_type = arrow_to_pg_type(field.data_type())?;
17        let not_null = if field.is_nullable() { "" } else { " NOT NULL" };
18        column_defs.push(format!("\"{}\" {}{}", field.name(), pg_type, not_null));
19    }
20    let ddl = format!("CREATE TABLE \"{}\" ({});", table, column_defs.join(", "));
21    Ok(ddl)
22}
23
24/// Import an Arrow RecordBatch stream into a PostgreSQL table.
25pub async fn arrow_to_table(
26    conn: &mut PgConnection,
27    table: &str,
28    create_table: bool,
29    schema: &Schema,
30    mut stream: impl Stream<Item = Result<RecordBatch>> + Unpin,
31) -> Result<()> {
32    if create_table {
33        let ddl = generate_create_table_ddl(table, schema)?;
34        conn.execute(ddl.as_str()).await?;
35    }
36
37    let copy_query = format!("COPY \"{}\" FROM STDIN WITH (FORMAT binary)", table);
38    let mut copy_in = conn.copy_in_raw(copy_query.as_str()).await?;
39
40    let mut buf = BytesMut::new();
41    buf.extend_from_slice(b"PGCOPY\n\xff\r\n\0");
42    buf.put_i32(0);
43    buf.put_i32(0);
44    copy_in.send(buf.split().freeze()).await?;
45
46    let num_cols = schema.fields().len();
47
48    while let Some(batch) = stream.next().await {
49        let batch = batch?;
50        let rows = batch.num_rows();
51
52        for r in 0..rows {
53            buf.put_i16(num_cols as i16);
54            for c in 0..num_cols {
55                let col = batch.column(c);
56                if col.is_null(r) {
57                    buf.put_i32(-1);
58                } else {
59                    match schema.field(c).data_type() {
60                        DataType::Int16 => {
61                            let arr = col.as_any().downcast_ref::<Int16Array>().unwrap();
62                            buf.put_i32(2);
63                            buf.put_i16(arr.value(r));
64                        }
65                        DataType::Int32 => {
66                            let arr = col.as_any().downcast_ref::<Int32Array>().unwrap();
67                            buf.put_i32(4);
68                            buf.put_i32(arr.value(r));
69                        }
70                        DataType::Int64 => {
71                            let arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
72                            buf.put_i32(8);
73                            buf.put_i64(arr.value(r));
74                        }
75                        DataType::Float32 => {
76                            let arr = col.as_any().downcast_ref::<Float32Array>().unwrap();
77                            buf.put_i32(4);
78                            buf.put_f32(arr.value(r));
79                        }
80                        DataType::Float64 => {
81                            let arr = col.as_any().downcast_ref::<Float64Array>().unwrap();
82                            buf.put_i32(8);
83                            buf.put_f64(arr.value(r));
84                        }
85                        DataType::Boolean => {
86                            let arr = col.as_any().downcast_ref::<BooleanArray>().unwrap();
87                            buf.put_i32(1);
88                            buf.put_u8(if arr.value(r) { 1 } else { 0 });
89                        }
90                        DataType::LargeBinary => {
91                            let arr = col.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
92                            let val = arr.value(r);
93                            buf.put_i32(val.len() as i32);
94                            buf.extend_from_slice(val);
95                        }
96                        DataType::LargeUtf8 => {
97                            let arr = col.as_any().downcast_ref::<LargeStringArray>().unwrap();
98                            let val = arr.value(r);
99                            buf.put_i32(val.len() as i32);
100                            buf.extend_from_slice(val.as_bytes());
101                        }
102                        DataType::Utf8 => {
103                            let arr = col.as_any().downcast_ref::<arrow::array::StringArray>().unwrap();
104                            let val = arr.value(r);
105                            buf.put_i32(val.len() as i32);
106                            buf.extend_from_slice(val.as_bytes());
107                        }
108                        _ => {
109                            return Err(anyhow!("Unsupported type in import: {:?}", schema.field(c).data_type()));
110                        }
111                    }
112                }
113            }
114            if buf.len() > 8192 {
115                copy_in.send(buf.split().freeze()).await?;
116            }
117        }
118    }
119
120    buf.put_i16(-1);
121    copy_in.send(buf.split().freeze()).await?;
122    copy_in.finish().await?;
123
124    Ok(())
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use arrow::datatypes::{DataType, Field, TimeUnit};
131
132    #[test]
133    fn test_generate_create_table_ddl_basic() {
134        let fields = vec![
135            Field::new("id", DataType::Int32, false),
136            Field::new("name", DataType::Utf8, true),
137            Field::new("score", DataType::Float64, true),
138            Field::new("active", DataType::Boolean, false),
139        ];
140        let schema = Schema::new(fields);
141        let ddl = generate_create_table_ddl("users", &schema).unwrap();
142        assert_eq!(
143            ddl,
144            "CREATE TABLE \"users\" (\"id\" INT4 NOT NULL, \"name\" TEXT, \"score\" FLOAT8, \"active\" BOOL NOT NULL);"
145        );
146    }
147
148    #[test]
149    fn test_generate_create_table_ddl_all_types() {
150        let fields = vec![
151            Field::new("col_i16", DataType::Int16, false),
152            Field::new("col_i32", DataType::Int32, false),
153            Field::new("col_i64", DataType::Int64, false),
154            Field::new("col_f32", DataType::Float32, true),
155            Field::new("col_f64", DataType::Float64, true),
156            Field::new("col_bool", DataType::Boolean, false),
157            Field::new("col_utf8", DataType::Utf8, true),
158            Field::new("col_large_utf8", DataType::LargeUtf8, true),
159            Field::new("col_bin", DataType::Binary, true),
160            Field::new("col_large_bin", DataType::LargeBinary, true),
161            Field::new("col_date", DataType::Date32, true),
162            Field::new("col_time", DataType::Time64(TimeUnit::Microsecond), true),
163            Field::new("col_ts", DataType::Timestamp(TimeUnit::Microsecond, None), true),
164            Field::new(
165                "col_tstz",
166                DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
167                true,
168            ),
169            Field::new("col_decimal", DataType::Decimal128(38, 9), true),
170        ];
171        let schema = Schema::new(fields);
172        let ddl = generate_create_table_ddl("all_types", &schema).unwrap();
173        assert_eq!(
174            ddl,
175            "CREATE TABLE \"all_types\" (\
176             \"col_i16\" INT2 NOT NULL, \
177             \"col_i32\" INT4 NOT NULL, \
178             \"col_i64\" INT8 NOT NULL, \
179             \"col_f32\" FLOAT4, \
180             \"col_f64\" FLOAT8, \
181             \"col_bool\" BOOL NOT NULL, \
182             \"col_utf8\" TEXT, \
183             \"col_large_utf8\" TEXT, \
184             \"col_bin\" BYTEA, \
185             \"col_large_bin\" BYTEA, \
186             \"col_date\" DATE, \
187             \"col_time\" TIME, \
188             \"col_ts\" TIMESTAMP, \
189             \"col_tstz\" TIMESTAMPTZ, \
190             \"col_decimal\" NUMERIC);"
191        );
192    }
193
194    #[test]
195    fn test_generate_create_table_ddl_unsupported_type() {
196        let fields = vec![Field::new("invalid_col", DataType::Null, true)];
197        let schema = Schema::new(fields);
198        let err = generate_create_table_ddl("invalid_table", &schema);
199        assert!(err.is_err());
200    }
201
202    #[test]
203    fn test_import_interface_signature() {
204        // Verify type signature and Stream Item conformance
205        fn _assert_signature<'a, S>(_stream: S)
206        where
207            S: Stream<Item = Result<RecordBatch>> + Unpin + 'a,
208        {
209            let _ = |conn: &'a mut PgConnection,
210                     table: &'a str,
211                     create_table: bool,
212                     schema: &'a Schema,
213                     stream: S| {
214                arrow_to_table(conn, table, create_table, schema, stream)
215            };
216        }
217
218        _assert_signature(futures::stream::empty::<Result<RecordBatch>>());
219    }
220}
postg_arrow/import.rs

postg_arrow/
import.rs
