Merge pull request #1 from EricFecteau/arrow_rs_polars

Update arrow-rs with .polars

Author: EricFecteau

Cargo.toml

@@ -10,3 +10,5 @@ lto = true
 [workspace.dependencies]
 arrow = {version = "46", features = ["prettyprint", "ffi"]}
 arrow2 = {version = "0.17", default-features = false}
+polars = {version = "0.45", features=["dtype-u8", "dtype-u16", "lazy"]}
+polars-arrow = {version = "0.45"}

connectorx/Cargo.toml

@@ -24,6 +24,8 @@ chrono = "0.4"
 
 arrow = {workspace = true, optional = true}
 arrow2 = {workspace = true, default-features = false, optional = true}
+polars = {workspace = true, optional = true, features=["dtype-u8", "dtype-u16", "lazy"]}
+polars-arrow = {workspace = true, optional = true}
 bb8 = {version = "0.7", optional = true}
 bb8-tiberius = {version = "0.8", optional = true}
 csv = {version = "1", optional = true}
@@ -36,8 +38,6 @@ ndarray = {version = "0.15", optional = true}
 num-traits = {version = "0.2", optional = true}
 openssl = {version = "0.10", optional = true, features = ["vendored"]}
 oracle = {version = "0.5", optional = true}
-polars = {version = "0.45", optional = true, features=["dtype-u8", "dtype-u16", "lazy"]}
-polars-arrow = {version = "0.45", optional = true}
 postgres = {version = "0.19", features = ["with-chrono-0_4", "with-uuid-0_8", "with-serde_json-1"], optional = true}
 postgres-native-tls = {version = "0.5", optional = true}
 postgres-openssl = {version = "0.5", optional = true}
@@ -72,11 +72,12 @@ iai = "0.1"
 pprof = {version = "0.5", features = ["flamegraph"]}
 
 [features]
-all = ["src_sqlite", "src_postgres", "src_mysql", "src_mssql", "src_oracle", "src_bigquery", "src_csv", "src_dummy", "src_trino", "dst_arrow", "dst_arrow2", "federation", "fed_exec"]
+all = ["src_sqlite", "src_postgres", "src_mysql", "src_mssql", "src_oracle", "src_bigquery", "src_csv", "src_dummy", "src_trino", "dst_arrow", "dst_arrow2", "dst_arrow", "federation", "fed_exec", "dst_polars"]
 branch = []
 default = ["fptr"]
 dst_arrow = ["arrow"]
 dst_arrow2 = ["polars", "polars-arrow", "arrow2"]
+dst_polars = ["dst_arrow", "polars", "polars-arrow"]
 fptr = []
 src_bigquery = ["gcp-bigquery-client", "tokio"]
 src_csv = ["csv", "regex"]

connectorx/src/destinations/arrow/errors.rs

@@ -10,6 +10,10 @@ pub enum ArrowDestinationError {
     #[error(transparent)]
     ConnectorXError(#[from] crate::errors::ConnectorXError),
 
+    #[cfg(feature = "dst_polars")]
+    #[error(transparent)]
+    PolarsError(#[from] polars::error::PolarsError),
+
     /// Any other errors that are too trivial to be put here explicitly.
     #[error(transparent)]
     Other(#[from] anyhow::Error),

connectorx/src/destinations/arrow/mod.rs

@@ -22,6 +22,15 @@ use std::{
     sync::{Arc, Mutex},
 };
 
+#[cfg(feature = "dst_polars")]
+use {
+    arrow::ffi::to_ffi,
+    polars::prelude::{concat, DataFrame, IntoLazy, PlSmallStr, Series, UnionArgs},
+    polars_arrow::ffi::{import_array_from_c, import_field_from_c},
+    std::iter::FromIterator,
+    std::mem::transmute,
+};
+
 type Builder = Box<dyn Any + Send>;
 type Builders = Vec<Builder>;
 
@@ -125,6 +134,63 @@ impl ArrowDestination {
             .map_err(|e| anyhow!("mutex poisoned {}", e))?
     }
 
+    #[cfg(feature = "dst_polars")]
+    #[throws(ArrowDestinationError)]
+    pub fn polars(self) -> DataFrame {
+        // Convert to arrow first
+        let rbs = self.arrow()?;
+
+        // Ready LazyFrame vector for the chunks
+        let mut lf_vec = vec![];
+
+        for chunk in rbs.into_iter() {
+            // Column vector
+            let mut columns = Vec::with_capacity(chunk.num_columns());
+
+            // Arrow stores data by columns, therefore need to be Zero-copied by column
+            for (i, col) in chunk.columns().iter().enumerate() {
+                // Convert to ArrayData (arrow-rs)
+                let array = col.to_data();
+
+                // Convert to ffi with arrow-rs
+                let (out_array, out_schema) = to_ffi(&array).unwrap();
+
+                // Import field from ffi with polars
+                let field = unsafe {
+                    import_field_from_c(transmute::<
+                        &arrow::ffi::FFI_ArrowSchema,
+                        &polars_arrow::ffi::ArrowSchema,
+                    >(&out_schema))
+                }
+                .unwrap();
+
+                // Import data from ffi with polars
+                let data = unsafe {
+                    import_array_from_c(
+                        transmute::<arrow::ffi::FFI_ArrowArray, polars_arrow::ffi::ArrowArray>(
+                            out_array,
+                        ),
+                        field.dtype().clone(),
+                    )
+                }
+                .unwrap();
+
+                // Create Polars series from arrow column
+                columns.push(Series::from_arrow(
+                    PlSmallStr::from(chunk.schema().field(i).name()),
+                    data,
+                )?);
+            }
+
+            // Create DataFrame from the columns
+            lf_vec.push(DataFrame::from_iter(columns).lazy());
+        }
+
+        // Concat the chunks
+        let union_args = UnionArgs::default();
+        concat(lf_vec, union_args)?.collect()?
+    }
+
     #[throws(ArrowDestinationError)]
     pub fn record_batch(&mut self) -> Option<RecordBatch> {
         let mut guard = self