update datafusion to 52 & update deps (#525)

* update datafusion to 52 & update deps

* Update byte-unit to 5.2

* python: update FFI_TableProvider::new function

* Downgrade adbc to 0.21

* upgrade datafusion-python

* python: Fix an issue with RawTableProvider

* python/pyproject: use compatible version of maturin

* Use datafusion-python 52

* Fix segfault in RawTableProvider by keeping SessionContext alive

The FFI_TableProvider stores a Weak reference to the TaskContextProvider,
but the Arc was created locally and dropped immediately. Use a static
OnceLock to keep it alive. Also update the __datafusion_table_provider__
signature for the v52 protocol and bump datafusion dependency to >=52.0.0.

* core/tests/flight: add a workaround for df.count()

Author: hozan23

Cargo.lock

@@ -28,24 +28,25 @@ license = "Apache-2.0"
 description = "Extend the capabilities of DataFusion to support additional data sources via implementations of the `TableProvider` trait."
 
 [workspace.dependencies]
-arrow = "57.0.0"
-arrow-array = { version = "57.0.0" }
-arrow-flight = { version = "57.0.0", features = [
-  "flight-sql-experimental",
-  "tls-ring",
+arrow = "57.2"
+arrow-array = { version = "57.2" }
+arrow-flight = { version = "57.2", features = [
+    "flight-sql-experimental",
+    "tls-ring",
 ] }
-arrow-ipc = { version = "57.0.0" }
-arrow-schema = { version = "57.2.0", features = ["serde"] }
-arrow-json = "57.0.0"
-arrow-odbc = { version = "21.0.0" }
-datafusion = { version = "51", default-features = false }
-datafusion-expr = { version = "51" }
-datafusion-federation = { version = "0.4.11" }
-datafusion-ffi = { version = "51" }
-datafusion-proto = { version = "51" }
-datafusion-physical-expr = { version = "51" }
-datafusion-physical-plan = { version = "51" }
+arrow-ipc = { version = "57.2" }
+arrow-schema = { version = "57.2", features = ["serde"] }
+arrow-json = "57.2"
+arrow-odbc = { version = "21.0" }
+datafusion = { version = "52.0", default-features = false }
+datafusion-expr = { version = "52.0" }
+datafusion-federation = { version = "0.4" }
+datafusion-ffi = { version = "52.0" }
+datafusion-proto = { version = "52.0" }
+datafusion-physical-expr = { version = "52.0" }
+datafusion-physical-plan = { version = "52.0" }
+datafusion-python = { version = "52.0" }
 datafusion-table-providers = { path = "core" }
-duckdb = { version = "=1.3.2", package = "spiceai_duckdb_fork" } # Forked to add support for duckdb_scan_arrow, pending: https://github.com/duckdb/duckdb-rs/pull/488
-adbc_core = { version = "0.21.0" }
-adbc_driver_manager = { version = "0.21.0" }
+duckdb = { version = "=1.3", package = "spiceai_duckdb_fork" } # Forked to add support for duckdb_scan_arrow, pending: https://github.com/duckdb/duckdb-rs/pull/488
+adbc_core = { version = "0.21" }
+adbc_driver_manager = { version = "0.21" }

Cargo.toml

@@ -18,18 +18,7 @@ arrow-flight = { workspace = true, optional = true, features = [
 arrow-schema = { workspace = true, optional = true, features = ["serde"] }
 arrow-json = { workspace = true }
 arrow-odbc = { workspace = true, optional = true }
-async-stream = { version = "0.3", optional = true }
-async-trait = "0.1"
-base64 = { version = "0.22.1", optional = true }
-bb8 = { version = "0.9", optional = true }
-bb8-postgres = { version = "0.9", optional = true }
-bigdecimal = "0.4"
-byteorder = "1.5.0"
-bytes = { version = "1.11.0", optional = true }
-byte-unit = { version = "5.1.4", optional = true }
-chrono = "0.4"
-clickhouse = { version = "0.13.3", optional = true }
-dashmap = "6.1.0"
+
 datafusion = { workspace = true, default-features = false, features = ["sql"]}
 datafusion-expr = { workspace = true, optional = true }
 datafusion-federation = { workspace = true, features = [
@@ -45,13 +34,28 @@ duckdb = { workspace = true, features = [
   "vtab-arrow",
   "appender-arrow",
 ], optional = true }
-libduckdb-sys = { version = "=1.3.0", optional = true }
+adbc_driver_manager = { workspace = true, optional = true }
+adbc_core = { workspace = true, optional = true }
+
+async-stream = { version = "0.3", optional = true }
+async-trait = "0.1"
+base64 = { version = "0.22", optional = true }
+bb8 = { version = "0.9", optional = true }
+bb8-postgres = { version = "0.9", optional = true }
+bigdecimal = "0.4"
+byteorder = "1.5"
+bytes = { version = "1.11", optional = true }
+byte-unit = { version = "5.2", optional = true }
+chrono = "0.4"
+clickhouse = { version = "0.13", optional = true }
+dashmap = "6.1"
+libduckdb-sys = { version = "=1.3", optional = true }
 dyn-clone = { version = "1.0", optional = true }
-fallible-iterator = "0.3.0"
-fundu = "2.0.1"
+fallible-iterator = "0.3"
+fundu = "2.0"
 futures = "0.3"
 geo-types = "0.7"
-itertools = "0.14.0"
+itertools = "0.14"
 mysql_async = { version = "0.36", features = [
   "native-tls-tls",
   "chrono",
@@ -61,13 +65,13 @@ mysql_async = { version = "0.36", features = [
 native-tls = { version = "0.2", optional = true }
 num-bigint = "0.4"
 odbc-api = { version = "19.0", optional = true }
-pem = { version = "3.0.4", optional = true }
+pem = { version = "3.0", optional = true }
 postgres-native-tls = { version = "0.5.0", optional = true }
-prost = { version = "=0.14.1", optional = true }
+prost = { version = "=0.14", optional = true }
 rand = { version = "0.9" }
 regex = { version = "1" }
 r2d2 = { version = "0.8", optional = true }
-rusqlite = { version = "0.37.0", optional = true }
+rusqlite = { version = "0.37", optional = true }
 sea-query = { version = "0.32", features = [
   "backend-sqlite",
   "backend-postgres",
@@ -77,7 +81,7 @@ sea-query = { version = "0.32", features = [
   "with-time",
   "with-chrono",
 ] }
-secrecy = "0.10.3"
+secrecy = "0.10"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 sha2 = "0.10"
@@ -90,18 +94,17 @@ tokio-postgres = { version = "0.7", features = [
   "with-serde_json-1",
   "with-geo-types-0_7",
 ], optional = true }
-tokio-rusqlite = { version = "0.7.0", optional = true }
+tokio-rusqlite = { version = "0.7", optional = true }
 tonic = { version = "0.14", optional = true, features = [
   "tls-native-roots",
   "tls-webpki-roots",
 ] }
 tracing = "0.1"
-trust-dns-resolver = "0.23.2"
-url = "2.5.4"
+trust-dns-resolver = "0.23"
+url = "2.5"
 uuid = { version = "1.18", optional = true }
-rust_decimal = { version = "1.38.0", features = ["db-postgres"] }
-adbc_driver_manager = { workspace = true, optional = true }
-adbc_core = { workspace = true, optional = true }
+rust_decimal = { version = "1.38", features = ["db-postgres"] }
+
 r2d2_adbc = { version = "0.2.0", optional = true }
 
 [dev-dependencies]
@@ -203,4 +206,4 @@ required-features = ["sqlite"]
 [[example]]
 name = "adbc"
 path = "examples/adbc_duckdb.rs"
-required-features = ["adbc"]
+required-features = ["adbc"]

core/Cargo.toml

@@ -213,10 +213,15 @@ async fn test_flight_sql_data_source() -> datafusion::common::Result<()> {
         .await
         .unwrap();
     let df = ctx.sql("select col1 from fsql").await.unwrap();
-    assert_eq!(
-        df.count().await.unwrap(),
-        rows_per_partition * num_partitions
-    );
+
+    let row_count: usize = df
+        .collect()
+        .await
+        .unwrap()
+        .iter()
+        .map(|b| b.num_rows())
+        .sum();
+    assert_eq!(row_count, rows_per_partition * num_partitions);
     let df = ctx.sql("select sum(col2) from fsql").await?;
     let rb = df
         .collect()

core/tests/flight/mod.rs

@@ -16,17 +16,19 @@ doc = false
 [dependencies]
 arrow = { workspace = true }
 arrow-flight = { workspace = true, optional = true }
-datafusion = { workspace = true, features = ["pyarrow"] }
+datafusion = { workspace = true }
 datafusion-ffi = { workspace = true }
+datafusion-python = { workspace = true } 
 datafusion-table-providers = { workspace = true }
-pyo3 = { version = "0.26.0" }
+duckdb = { workspace = true, optional = true }
+
+pyo3 = { version = "0.26" }
 tokio = { version = "1.46", features = [
     "macros",
     "rt",
     "rt-multi-thread",
     "sync",
 ] }
-duckdb = { workspace = true, optional = true }
 
 [features]
 default = [

python/Cargo.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["maturin>=1.5.1,<1.6.0"]
+requires = ["maturin>=1.5.1,<1.9.0"]
 build-backend = "maturin"
 
 [project]
@@ -8,7 +8,7 @@ version = "0.1.0"
 description = "Build and run queries against data"
 readme = "../README.md"
 license = { file = "../LICENSE" }
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 keywords = ["datafusion", "dataframe", "rust", "query-engine"]
 classifier = [
     "Development Status :: 2 - Pre-Alpha",
@@ -24,7 +24,7 @@ classifier = [
     "Programming Language :: Python",
     "Programming Language :: Rust",
 ]
-dependencies = ["datafusion>=45.0.0"]
+dependencies = ["datafusion>=52.0.0"]
 
 [project.urls]
 repository = "https://github.com/datafusion-contrib/datafusion-table-providers"

python/pyproject.toml

@@ -3,7 +3,7 @@ use std::{
     sync::{Arc, OnceLock},
 };
 
-use datafusion::catalog::TableProvider;
+use datafusion::{catalog::TableProvider, execution::TaskContextProvider, prelude::SessionContext};
 use datafusion_ffi::table_provider::FFI_TableProvider;
 use pyo3::{prelude::*, types::PyCapsule};
 
@@ -19,11 +19,20 @@ pub(crate) fn get_tokio_runtime() -> &'static tokio::runtime::Runtime {
     RUNTIME.get_or_init(|| tokio::runtime::Runtime::new().expect("Failed to create Tokio runtime"))
 }
 
+/// Returns a static reference to a default SessionContext.
+/// This ensures the Arc backing the Weak reference inside
+/// FFI_TaskContextProvider is never dropped.
+fn get_default_session_context() -> &'static Arc<SessionContext> {
+    static CTX: OnceLock<Arc<SessionContext>> = OnceLock::new();
+    CTX.get_or_init(|| Arc::new(SessionContext::default()))
+}
+
 #[pymethods]
 impl RawTableProvider {
     fn __datafusion_table_provider__<'py>(
         &self,
         py: Python<'py>,
+        _session: Bound<'py, PyAny>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
         let name = CString::new("datafusion_table_provider").unwrap();
 
@@ -33,10 +42,13 @@ impl RawTableProvider {
             None
         };
 
+        let ctx = Arc::clone(get_default_session_context()) as Arc<dyn TaskContextProvider>;
         let provider = FFI_TableProvider::new(
             Arc::clone(&self.table),
             self.supports_pushdown_filters,
             runtime,
+            &ctx,
+            None,
         );
 
         PyCapsule::new(py, provider, Some(name.clone()))