refactor: introducing datafusion feature in arrow-pg (#93)

* feat: add datafusion feature to arrow-pg

* refactor: move datafusion encoding logic into arrow-pg

* ci: add test for cargo-pg default feature

Signed-off-by: Ning Sun <[email protected]>

---------

Signed-off-by: Ning Sun <[email protected]>

Author: sunng87

.github/workflows/ci.yml

@@ -53,6 +53,9 @@ jobs:
         override: true
     - name: Build and run tests
       run: cargo test --all-features
+    - name: Test arrow-pg default features
+      working-directory: arrow-pg
+      run: cargo test
 
   integration:
     name: Integration tests

.gitignore

@@ -1,4 +1,5 @@
 /target
 .direnv
 .envrc
-.vscode
+.vscode
+.aider*

Cargo.lock

@@ -12,10 +12,16 @@ documentation.workspace = true
 readme = "../README.md"
 rust-version.workspace = true
 
+[features]
+default = ["arrow"]
+arrow = ["dep:arrow"]
+datafusion = ["dep:datafusion"]
+
 [dependencies]
-arrow.workspace = true
+arrow = { workspace = true, optional = true }
 bytes.workspace = true
 chrono.workspace = true
+datafusion = { workspace = true, optional = true }
 futures.workspace = true
 pgwire = { workspace = true, features = ["server-api"] }
 postgres-types.workspace = true

arrow-pg/Cargo.toml

@@ -1,7 +1,10 @@
 use std::sync::Arc;
 
-use arrow::datatypes::*;
-use arrow::record_batch::RecordBatch;
+#[cfg(not(feature = "datafusion"))]
+use arrow::{datatypes::*, record_batch::RecordBatch};
+#[cfg(feature = "datafusion")]
+use datafusion::arrow::{datatypes::*, record_batch::RecordBatch};
+
 use pgwire::api::portal::Format;
 use pgwire::api::results::FieldInfo;
 use pgwire::api::Type;
@@ -11,6 +14,9 @@ use postgres_types::Kind;
 
 use crate::row_encoder::RowEncoder;
 
+#[cfg(feature = "datafusion")]
+pub mod df;
+
 pub fn into_pg_type(arrow_type: &DataType) -> PgWireResult<Type> {
     Ok(match arrow_type {
         DataType::Null => Type::UNKNOWN,

arrow-pg/src/datatypes.rs

@@ -17,9 +17,9 @@ use pgwire::messages::data::DataRow;
 use rust_decimal::prelude::ToPrimitive;
 use rust_decimal::Decimal;
 
-use arrow_pg::datatypes::{arrow_schema_to_pg_fields, encode_recordbatch, into_pg_type};
+use super::{arrow_schema_to_pg_fields, encode_recordbatch, into_pg_type};
 
-pub(crate) async fn encode_dataframe<'a>(
+pub async fn encode_dataframe<'a>(
     df: DataFrame,
     format: &Format,
 ) -> PgWireResult<QueryResponse<'a>> {
@@ -51,7 +51,7 @@ pub(crate) async fn encode_dataframe<'a>(
 /// If the type is empty or unknown, we fallback to datafusion inferenced type
 /// from `inferenced_types`.
 /// An error will be raised when neither sources can provide type information.
-pub(crate) fn deserialize_parameters<S>(
+pub fn deserialize_parameters<S>(
     portal: &Portal<S>,
     inferenced_types: &[Option<&DataType>],
 ) -> PgWireResult<ParamValues>

datafusion-postgres/src/datatypes.rs renamed to arrow-pg/src/datatypes/df.rs

@@ -3,11 +3,13 @@ use std::io::Write;
 use std::str::FromStr;
 use std::sync::Arc;
 
-use arrow::array::*;
-use arrow::datatypes::*;
+#[cfg(not(feature = "datafusion"))]
+use arrow::{array::*, datatypes::*};
 use bytes::BufMut;
 use bytes::BytesMut;
 use chrono::{NaiveDate, NaiveDateTime};
+#[cfg(feature = "datafusion")]
+use datafusion::arrow::{array::*, datatypes::*};
 use pgwire::api::results::DataRowEncoder;
 use pgwire::api::results::FieldFormat;
 use pgwire::error::PgWireError;

arrow-pg/src/encoder.rs

@@ -1,3 +1,8 @@
+//! Arrow data encoding and type mapping for Postgres(pgwire).
+
+// #[cfg(all(feature = "arrow", feature = "datafusion"))]
+// compile_error!("Feature arrow and datafusion cannot be enabled at same time. Use no-default-features when activating datafusion");
+
 pub mod datatypes;
 pub mod encoder;
 mod error;

arrow-pg/src/lib.rs

@@ -1,19 +1,36 @@
 use std::{str::FromStr, sync::Arc};
 
-use arrow::array::{
-    timezone::Tz, Array, BinaryArray, BooleanArray, Date32Array, Date64Array, Decimal128Array,
-    LargeBinaryArray, PrimitiveArray, StringArray, Time32MillisecondArray, Time32SecondArray,
-    Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray,
-    TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray,
-};
+#[cfg(not(feature = "datafusion"))]
 use arrow::{
+    array::{
+        timezone::Tz, Array, BinaryArray, BooleanArray, Date32Array, Date64Array, Decimal128Array,
+        LargeBinaryArray, PrimitiveArray, StringArray, Time32MillisecondArray, Time32SecondArray,
+        Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray,
+        TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray,
+    },
     datatypes::{
         DataType, Date32Type, Date64Type, Float32Type, Float64Type, Int16Type, Int32Type,
         Int64Type, Int8Type, Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
         Time64NanosecondType, TimeUnit, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
     },
     temporal_conversions::{as_date, as_time},
 };
+#[cfg(feature = "datafusion")]
+use datafusion::arrow::{
+    array::{
+        timezone::Tz, Array, BinaryArray, BooleanArray, Date32Array, Date64Array, Decimal128Array,
+        LargeBinaryArray, PrimitiveArray, StringArray, Time32MillisecondArray, Time32SecondArray,
+        Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray,
+        TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray,
+    },
+    datatypes::{
+        DataType, Date32Type, Date64Type, Float32Type, Float64Type, Int16Type, Int32Type,
+        Int64Type, Int8Type, Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
+        Time64NanosecondType, TimeUnit, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
+    },
+    temporal_conversions::{as_date, as_time},
+};
+
 use bytes::{BufMut, BytesMut};
 use chrono::{DateTime, TimeZone, Utc};
 use pgwire::api::results::FieldFormat;

arrow-pg/src/list_encoder.rs

@@ -1,6 +1,10 @@
 use std::sync::Arc;
 
+#[cfg(not(feature = "datafusion"))]
 use arrow::array::RecordBatch;
+#[cfg(feature = "datafusion")]
+use datafusion::arrow::array::RecordBatch;
+
 use pgwire::{
     api::results::{DataRowEncoder, FieldInfo},
     error::PgWireResult,