Merge pull request #405 from splitgraph/pg-json-to-arrow-transport

Add PG JSON to Arrow transport rules

Author: wangxiaoying

Cargo.lock

@@ -44,9 +44,9 @@ r2d2 = {version = "0.8", optional = true}
 r2d2-oracle = {version = "0.5.0", features = ["chrono"], optional = true}
 r2d2_mysql = {version = "21.0", optional = true}
 r2d2_postgres = {version = "0.18.1", optional = true}
-r2d2_sqlite = {version = "0.18", optional = true}
+r2d2_sqlite = {version = "0.20.0", optional = true}
 regex = {version = "1", optional = true}
-rusqlite = {version = "0.25", features = ["column_decltype", "chrono", "bundled"], optional = true}
+rusqlite = {version = "0.27.0", features = ["column_decltype", "chrono", "bundled"], optional = true}
 rust_decimal = {version = "1", features = ["db-postgres"], optional = true}
 rust_decimal_macros = {version = "1", optional = true}
 tiberius = {version = "0.5", features = ["rust_decimal", "chrono"], optional = true}

connectorx/Cargo.toml

@@ -3,7 +3,7 @@ use crate::constants::SECONDS_IN_DAY;
 use arrow::array::{
     ArrayBuilder, BooleanBuilder, Date32Builder, Date64Builder, Float32Builder, Float64Builder,
     Int32Builder, Int64Builder, LargeBinaryBuilder, StringBuilder, Time64NanosecondBuilder,
-    UInt32Builder, UInt64Builder,
+    TimestampNanosecondBuilder, UInt32Builder, UInt64Builder,
 };
 use arrow::datatypes::Field;
 use arrow::datatypes::{DataType as ArrowDataType, TimeUnit};
@@ -140,34 +140,44 @@ impl ArrowAssoc for Option<String> {
 }
 
 impl ArrowAssoc for DateTime<Utc> {
-    type Builder = Float64Builder;
+    type Builder = TimestampNanosecondBuilder;
 
-    fn builder(_nrows: usize) -> Float64Builder {
-        unimplemented!()
+    fn builder(nrows: usize) -> Self::Builder {
+        TimestampNanosecondBuilder::with_capacity(nrows)
     }
 
-    fn append(_builder: &mut Self::Builder, _value: DateTime<Utc>) -> Result<()> {
-        unimplemented!()
+    #[throws(ArrowDestinationError)]
+    fn append(builder: &mut Self::Builder, value: DateTime<Utc>) {
+        builder.append_value(value.timestamp_nanos())
     }
 
-    fn field(_header: &str) -> Field {
-        unimplemented!()
+    fn field(header: &str) -> Field {
+        Field::new(
+            header,
+            ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
+            true,
+        )
     }
 }
 
 impl ArrowAssoc for Option<DateTime<Utc>> {
-    type Builder = Float64Builder;
+    type Builder = TimestampNanosecondBuilder;
 
-    fn builder(_nrows: usize) -> Float64Builder {
-        unimplemented!()
+    fn builder(nrows: usize) -> Self::Builder {
+        TimestampNanosecondBuilder::with_capacity(nrows)
     }
 
-    fn append(_builder: &mut Self::Builder, _value: Option<DateTime<Utc>>) -> Result<()> {
-        unimplemented!()
+    #[throws(ArrowDestinationError)]
+    fn append(builder: &mut Self::Builder, value: Option<DateTime<Utc>>) {
+        builder.append_option(value.map(|x| x.timestamp_nanos()))
     }
 
-    fn field(_header: &str) -> Field {
-        unimplemented!()
+    fn field(header: &str) -> Field {
+        Field::new(
+            header,
+            ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
+            false,
+        )
     }
 }
 

connectorx/src/destinations/arrow/arrow_assoc.rs

@@ -90,7 +90,7 @@ where
             let l1query = limit1_query(query, &SQLiteDialect {})?;
 
             let is_sucess = conn.query_row(l1query.as_str(), [], |row| {
-                for (j, col) in row.columns().iter().enumerate() {
+                for (j, col) in row.as_ref().columns().iter().enumerate() {
                     if j >= names.len() {
                         names.push(col.name().to_string());
                     }
@@ -142,14 +142,15 @@ where
         }
 
         // tried all queries but all get empty result set
-        let mut stmt = conn.prepare(self.queries[0].as_str())?;
-        let rows = stmt.query([])?;
-
-        if let Some(cnames) = rows.column_names() {
-            self.names = cnames.into_iter().map(|s| s.to_string()).collect();
-            // set all columns as string (align with pandas)
-            self.schema = vec![SQLiteTypeSystem::Text(false); self.names.len()];
-        }
+        let stmt = conn.prepare(self.queries[0].as_str())?;
+
+        self.names = stmt
+            .column_names()
+            .into_iter()
+            .map(|s| s.to_string())
+            .collect();
+        // set all columns as string (align with pandas)
+        self.schema = vec![SQLiteTypeSystem::Text(false); self.names.len()];
     }
 
     #[throws(SQLiteSourceError)]

connectorx/src/sources/sqlite/mod.rs

@@ -13,6 +13,7 @@ use num_traits::ToPrimitive;
 use postgres::NoTls;
 use postgres_openssl::MakeTlsConnector;
 use rust_decimal::Decimal;
+use serde_json::Value;
 use std::marker::PhantomData;
 use thiserror::Error;
 use uuid::Uuid;
@@ -57,6 +58,8 @@ macro_rules! impl_postgres_transport {
                 { UUID[Uuid]                 => LargeUtf8[String]         | conversion option }
                 { Char[&'r str]              => LargeUtf8[String]         | conversion none }
                 { ByteA[Vec<u8>]             => LargeBinary[Vec<u8>]      | conversion auto }
+                { JSON[Value]                => LargeUtf8[String]         | conversion option }
+                { JSONB[Value]               => LargeUtf8[String]         | conversion none }
             }
         );
     }
@@ -83,3 +86,9 @@ impl<P, C> TypeConversion<Decimal, f64> for PostgresArrowTransport<P, C> {
             .unwrap_or_else(|| panic!("cannot convert decimal {:?} to float64", val))
     }
 }
+
+impl<P, C> TypeConversion<Value, String> for PostgresArrowTransport<P, C> {
+    fn convert(val: Value) -> String {
+        val.to_string()
+    }
+}