refactor: introducing datafusion feature in arrow-pg

.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*

arrow-pg/Cargo.toml

@@ -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/src/datatypes.rs

@@ -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,

datafusion-postgres/src/datatypes.rs → arrow-pg/src/datatypes/df.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>

arrow-pg/src/encoder.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/lib.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/list_encoder.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/row_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,

arrow-pg/src/struct_encoder.rs

@@ -1,6 +1,10 @@
 use std::sync::Arc;
 
+#[cfg(not(feature = "datafusion"))]
 use arrow::array::{Array, StructArray};
+#[cfg(feature = "datafusion")]
+use datafusion::arrow::array::{Array, StructArray};
+
 use bytes::{BufMut, BytesMut};
 use pgwire::api::results::FieldFormat;
 use pgwire::error::PgWireResult;

datafusion-postgres/Cargo.toml

@@ -13,7 +13,7 @@ readme = "../README.md"
 rust-version.workspace = true
 
 [dependencies]
-arrow-pg = { path = "../arrow-pg", version = "0.1.1" }
+arrow-pg = { path = "../arrow-pg", version = "0.1.1", default-features = false, features = ["datafusion"] }
 bytes.workspace = true
 async-trait = "0.1"
 chrono.workspace = true

datafusion-postgres/src/handlers.rs

@@ -19,7 +19,7 @@ use pgwire::api::{ClientInfo, NoopErrorHandler, PgWireServerHandlers, Type};
 use pgwire::error::{PgWireError, PgWireResult};
 use tokio::sync::Mutex;
 
-use crate::datatypes;
+use arrow_pg::datatypes::df;
 use arrow_pg::datatypes::{arrow_schema_to_pg_fields, into_pg_type};
 
 pub struct HandlerFactory(pub Arc<DfSessionService>);
@@ -213,7 +213,7 @@ impl SimpleQueryHandler for DfSessionService {
             Ok(vec![Response::Execution(tag)])
         } else {
             // For non-INSERT queries, return a regular Query response
-            let resp = datatypes::encode_dataframe(df, &Format::UnifiedText).await?;
+            let resp = df::encode_dataframe(df, &Format::UnifiedText).await?;
             Ok(vec![Response::Query(resp)])
         }
     }
@@ -304,8 +304,7 @@ impl ExtendedQueryHandler for DfSessionService {
         let param_types = plan
             .get_parameter_types()
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-        let param_values =
-            datatypes::deserialize_parameters(portal, &ordered_param_types(&param_types))?; // Fixed: Use &param_types
+        let param_values = df::deserialize_parameters(portal, &ordered_param_types(&param_types))?; // Fixed: Use &param_types
         let plan = plan
             .clone()
             .replace_params_with_values(&param_values)
@@ -315,7 +314,7 @@ impl ExtendedQueryHandler for DfSessionService {
             .execute_logical_plan(plan)
             .await
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-        let resp = datatypes::encode_dataframe(dataframe, &portal.result_column_format).await?;
+        let resp = df::encode_dataframe(dataframe, &portal.result_column_format).await?;
         Ok(Response::Query(resp))
     }
 }

datafusion-postgres/src/lib.rs

@@ -1,4 +1,3 @@
-mod datatypes;
 mod handlers;
 pub mod pg_catalog;