chore: Update to latest main on upstream DataFusion

Cargo.toml

@@ -32,24 +32,24 @@ resolver = "3"
 
 [workspace.dependencies]
 tokio = { version = "1.52" }
-pyo3 = { version = "0.28" }
-pyo3-async-runtimes = { version = "0.28" }
+pyo3 = { version = "0.29" }
+pyo3-async-runtimes = { version = "0.29" }
 pyo3-log = "0.13.3"
 chrono = { version = "0.4", default-features = false }
-arrow = { version = "58" }
-arrow-array = { version = "58" }
-arrow-schema = { version = "58" }
-arrow-select = { version = "58" }
-datafusion = { version = "54" }
-datafusion-substrait = { version = "54" }
-datafusion-proto = { version = "54" }
-datafusion-ffi = { version = "54" }
-datafusion-catalog = { version = "54", default-features = false }
-datafusion-common = { version = "54", default-features = false }
-datafusion-functions-aggregate = { version = "54" }
-datafusion-functions-window = { version = "54" }
-datafusion-spark = { version = "54" }
-datafusion-expr = { version = "54" }
+arrow = { version = "59" }
+arrow-array = { version = "59" }
+arrow-schema = { version = "59" }
+arrow-select = { version = "59" }
+datafusion = { version = "54.1.0" }
+datafusion-substrait = { version = "54.1.0" }
+datafusion-proto = { version = "54.1.0" }
+datafusion-ffi = { version = "54.1.0" }
+datafusion-catalog = { version = "54.1.0", default-features = false }
+datafusion-common = { version = "54.1.0", default-features = false }
+datafusion-functions-aggregate = { version = "54.1.0" }
+datafusion-functions-window = { version = "54.1.0" }
+datafusion-spark = { version = "54.1.0" }
+datafusion-expr = { version = "54.1.0" }
 prost = "0.14.3"
 serde_json = "1"
 uuid = { version = "1.23" }
@@ -62,7 +62,7 @@ url = "2"
 log = "0.4.29"
 parking_lot = "0.12"
 prost-types = "0.14.3"                                                # keep in line with `datafusion-substrait`
-pyo3-build-config = "0.28"
+pyo3-build-config = "0.29"
 datafusion-python-util = { path = "crates/util", version = "54.0.0" }
 
 [profile.release]
@@ -72,3 +72,13 @@ codegen-units = 2
 # We cannot publish to crates.io with any patches in the below section. Developers
 # must remove any entries in this section before creating a release candidate.
 [patch.crates-io]
+datafusion = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
+datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
+datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
+datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
+datafusion-catalog = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
+datafusion-common = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
+datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
+datafusion-functions-window = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
+datafusion-spark = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }
+datafusion-expr = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" }

crates/core/Cargo.toml

@@ -53,7 +53,7 @@ datafusion = { workspace = true, features = ["avro", "unicode_expressions"] }
 datafusion-substrait = { workspace = true, optional = true }
 datafusion-proto = { workspace = true }
 datafusion-ffi = { workspace = true }
-datafusion-spark = { workspace = true }
+datafusion-spark = { workspace = true, features = ["core"] }
 prost = { workspace = true } # keep in line with `datafusion-substrait`
 serde_json = { workspace = true }
 uuid = { workspace = true, features = ["v4"] }

crates/core/src/array.rs

@@ -63,10 +63,10 @@ impl PyArrowArrayExportable {
         };
 
         let ffi_schema = FFI_ArrowSchema::try_from(&field)?;
-        let schema_capsule = PyCapsule::new(py, ffi_schema, Some(cr"arrow_schema".into()))?;
+        let schema_capsule = PyCapsule::new_with_value(py, ffi_schema, cr"arrow_schema")?;
 
         let ffi_array = FFI_ArrowArray::new(&self.array.to_data());
-        let array_capsule = PyCapsule::new(py, ffi_array, Some(cr"arrow_array".into()))?;
+        let array_capsule = PyCapsule::new_with_value(py, ffi_array, cr"arrow_array")?;
 
         Ok((schema_capsule, array_capsule))
     }

crates/core/src/codec.rs

@@ -100,9 +100,13 @@ use datafusion::logical_expr::{
     TypeSignature, Volatility, WindowUDF, WindowUDFImpl,
 };
 use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
+use datafusion::physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec};
-use datafusion_proto::physical_plan::{DefaultPhysicalExtensionCodec, PhysicalExtensionCodec};
+use datafusion_proto::physical_plan::{
+    DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension,
+};
 use pyo3::prelude::*;
 use pyo3::sync::PyOnceLock;
 use pyo3::types::{PyBytes, PyTuple};
@@ -483,12 +487,18 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
         buf: &[u8],
         inputs: &[Arc<dyn ExecutionPlan>],
         ctx: &TaskContext,
+        proto_converter: &dyn PhysicalProtoConverterExtension,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        self.inner.try_decode(buf, inputs, ctx)
+        self.inner.try_decode(buf, inputs, ctx, proto_converter)
     }
 
-    fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
-        self.inner.try_encode(node, buf)
+    fn try_encode(
+        &self,
+        node: Arc<dyn ExecutionPlan>,
+        buf: &mut Vec<u8>,
+        proto_converter: &dyn PhysicalProtoConverterExtension,
+    ) -> Result<()> {
+        self.inner.try_encode(node, buf, proto_converter)
     }
 
     fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
@@ -509,16 +519,22 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
         self.inner.try_decode_udf(name, buf)
     }
 
-    fn try_encode_expr(&self, node: &Arc<dyn PhysicalExpr>, buf: &mut Vec<u8>) -> Result<()> {
-        self.inner.try_encode_expr(node, buf)
+    fn try_encode_expr(
+        &self,
+        node: &Arc<dyn PhysicalExpr>,
+        buf: &mut Vec<u8>,
+        ctx: &PhysicalExprEncodeCtx<'_>,
+    ) -> Result<()> {
+        self.inner.try_encode_expr(node, buf, ctx)
     }
 
     fn try_decode_expr(
         &self,
         buf: &[u8],
         inputs: &[Arc<dyn PhysicalExpr>],
+        ctx: &PhysicalExprDecodeCtx<'_>,
     ) -> Result<Arc<dyn PhysicalExpr>> {
-        self.inner.try_decode_expr(buf, inputs)
+        self.inner.try_decode_expr(buf, inputs, ctx)
     }
 
     fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> {

crates/core/src/context.rs

@@ -1371,12 +1371,10 @@ impl PySessionContext {
         &self,
         py: Python<'py>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
-        let name = cr"datafusion_task_context_provider".into();
-
         let ctx_provider = Arc::clone(&self.ctx) as Arc<dyn TaskContextProvider>;
         let ffi_ctx_provider = FFI_TaskContextProvider::from(&ctx_provider);
 
-        PyCapsule::new(py, ffi_ctx_provider, Some(name))
+        PyCapsule::new_with_value(py, ffi_ctx_provider, cr"datafusion_task_context_provider")
     }
 
     pub fn __datafusion_logical_extension_codec__<'py>(

crates/core/src/dataframe.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use std::collections::HashMap;
-use std::ffi::{CStr, CString};
+use std::ffi::CStr;
 use std::ptr::NonNull;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -1237,8 +1237,7 @@ impl PyDataFrame {
         // destructor provided by PyO3 will drop the stream unless ownership is
         // transferred to PyArrow during import.
         let stream = FFI_ArrowArrayStream::new(reader);
-        let name = CString::new(ARROW_ARRAY_STREAM_NAME.to_bytes()).unwrap();
-        let capsule = PyCapsule::new(py, stream, Some(name))?;
+        let capsule = PyCapsule::new_with_value(py, stream, ARROW_ARRAY_STREAM_NAME)?;
         Ok(capsule)
     }
 
@@ -1317,7 +1316,8 @@ impl PyDataFrame {
             None => Vec::new(), // Empty vector means fill null for all columns
         };
 
-        let df = self.df.as_ref().clone().fill_null(scalar_value.0, cols)?;
+        let cols = cols.iter().map(String::as_str).collect::<Vec<_>>();
+        let df = self.df.as_ref().fill_null(&scalar_value.0, &cols)?;
         Ok(Self::new(df))
     }
 }

crates/core/src/expr/create_external_table.rs

@@ -88,7 +88,7 @@ impl PyCreateExternalTable {
         let create = CreateExternalTable {
             schema: Arc::new(schema.into()),
             name: name.into(),
-            location,
+            locations: vec![location],
             file_type,
             table_partition_cols,
             if_not_exists,
@@ -118,8 +118,8 @@ impl PyCreateExternalTable {
         Ok(self.create.name.to_string())
     }
 
-    pub fn location(&self) -> String {
-        self.create.location.clone()
+    pub fn locations(&self) -> Vec<String> {
+        self.create.locations.clone()
     }
 
     pub fn file_type(&self) -> String {

crates/core/src/expr/dml.rs

@@ -18,6 +18,7 @@
 use datafusion::logical_expr::dml::InsertOp;
 use datafusion::logical_expr::{DmlStatement, WriteOp};
 use pyo3::IntoPyObjectExt;
+use pyo3::exceptions::PyNotImplementedError;
 use pyo3::prelude::*;
 
 use super::logical_node::LogicalNode;
@@ -71,8 +72,8 @@ impl PyDmlStatement {
         })
     }
 
-    pub fn op(&self) -> PyWriteOp {
-        self.dml.op.clone().into()
+    pub fn op(&self) -> PyResult<PyWriteOp> {
+        self.dml.op.clone().try_into()
     }
 
     pub fn input(&self) -> PyLogicalPlan {
@@ -112,16 +113,21 @@ pub enum PyWriteOp {
     Truncate,
 }
 
-impl From<WriteOp> for PyWriteOp {
-    fn from(write_op: WriteOp) -> Self {
+impl TryFrom<WriteOp> for PyWriteOp {
+    type Error = PyErr;
+
+    fn try_from(write_op: WriteOp) -> Result<Self, Self::Error> {
         match write_op {
-            WriteOp::Insert(InsertOp::Append) => PyWriteOp::Append,
-            WriteOp::Insert(InsertOp::Overwrite) => PyWriteOp::Overwrite,
-            WriteOp::Insert(InsertOp::Replace) => PyWriteOp::Replace,
-            WriteOp::Update => PyWriteOp::Update,
-            WriteOp::Delete => PyWriteOp::Delete,
-            WriteOp::Ctas => PyWriteOp::Ctas,
-            WriteOp::Truncate => PyWriteOp::Truncate,
+            WriteOp::Insert(InsertOp::Append) => Ok(PyWriteOp::Append),
+            WriteOp::Insert(InsertOp::Overwrite) => Ok(PyWriteOp::Overwrite),
+            WriteOp::Insert(InsertOp::Replace) => Ok(PyWriteOp::Replace),
+            WriteOp::Update => Ok(PyWriteOp::Update),
+            WriteOp::Delete => Ok(PyWriteOp::Delete),
+            WriteOp::Ctas => Ok(PyWriteOp::Ctas),
+            WriteOp::Truncate => Ok(PyWriteOp::Truncate),
+            unsupported => Err(PyNotImplementedError::new_err(format!(
+                "DataFusion write operation {unsupported:?} is not supported"
+            ))),
         }
     }
 }

crates/core/src/expr/drop_catalog_schema.rs

@@ -18,9 +18,8 @@
 use std::fmt::{self, Display, Formatter};
 use std::sync::Arc;
 
-use datafusion::common::SchemaReference;
+use datafusion::common::{SchemaReference, TableReference};
 use datafusion::logical_expr::DropCatalogSchema;
-use datafusion::sql::TableReference;
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;

crates/core/src/expr/recursive_query.rs

@@ -22,6 +22,7 @@ use pyo3::IntoPyObjectExt;
 use pyo3::prelude::*;
 
 use super::logical_node::LogicalNode;
+use crate::errors::PyDataFusionResult;
 use crate::sql::logical::PyLogicalPlan;
 
 #[pyclass(
@@ -67,15 +68,15 @@ impl PyRecursiveQuery {
         static_term: PyLogicalPlan,
         recursive_term: PyLogicalPlan,
         is_distinct: bool,
-    ) -> Self {
-        Self {
-            query: RecursiveQuery {
+    ) -> PyDataFusionResult<Self> {
+        Ok(Self {
+            query: RecursiveQuery::try_new(
                 name,
-                static_term: static_term.plan(),
-                recursive_term: recursive_term.plan(),
+                static_term.plan(),
+                recursive_term.plan(),
                 is_distinct,
-            },
-        }
+            )?,
+        })
     }
 
     fn name(&self) -> PyResult<String> {

crates/core/src/sql/logical.rs

@@ -135,7 +135,7 @@ impl PyLogicalPlan {
             LogicalPlan::Dml(plan) => PyDmlStatement::from(plan.clone()).to_variant(py),
             LogicalPlan::Ddl(plan) => match plan {
                 DdlStatement::CreateExternalTable(plan) => {
-                    PyCreateExternalTable::from(plan.clone()).to_variant(py)
+                    PyCreateExternalTable::from(plan.as_ref().clone()).to_variant(py)
                 }
                 DdlStatement::CreateMemoryTable(plan) => {
                     PyCreateMemoryTable::from(plan.clone()).to_variant(py)
@@ -154,7 +154,7 @@ impl PyLogicalPlan {
                     PyDropCatalogSchema::from(plan.clone()).to_variant(py)
                 }
                 DdlStatement::CreateFunction(plan) => {
-                    PyCreateFunction::from(plan.clone()).to_variant(py)
+                    PyCreateFunction::from(plan.as_ref().clone()).to_variant(py)
                 }
                 DdlStatement::DropFunction(plan) => {
                     PyDropFunction::from(plan.clone()).to_variant(py)

crates/core/src/substrait.rs

@@ -131,8 +131,8 @@ impl PySubstraitSerializer {
     }
 
     #[staticmethod]
-    pub fn deserialize_bytes(proto_bytes: Vec<u8>, py: Python) -> PyDataFusionResult<PyPlan> {
-        let plan = wait_for_future(py, serializer::deserialize_bytes(proto_bytes))??;
+    pub fn deserialize_bytes(proto_bytes: Vec<u8>, _py: Python) -> PyDataFusionResult<PyPlan> {
+        let plan = serializer::deserialize_bytes(&proto_bytes)?;
         Ok(PyPlan { plan: *plan })
     }
 }

crates/util/src/lib.rs

@@ -164,7 +164,8 @@ pub fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &str) -> PyResult<()
         )));
     }
 
-    let capsule_name = unsafe { capsule_name.unwrap().as_cstr().to_str()? };
+    let capsule_name = unsafe { capsule_name.unwrap().as_cstr().to_str() }
+        .map_err(|err| PyValueError::new_err(err.to_string()))?;
     if capsule_name != name {
         return Err(PyValueError::new_err(format!(
             "Expected name '{name}' in PyCapsule, instead got '{capsule_name}'"
@@ -208,10 +209,9 @@ pub fn create_logical_extension_capsule<'py>(
     py: Python<'py>,
     codec: &FFI_LogicalExtensionCodec,
 ) -> PyResult<Bound<'py, PyCapsule>> {
-    let name = cr"datafusion_logical_extension_codec".into();
     let codec = codec.clone();
 
-    PyCapsule::new(py, codec, Some(name))
+    PyCapsule::new_with_value(py, codec, cr"datafusion_logical_extension_codec")
 }
 
 pub fn ffi_logical_codec_from_pycapsule(obj: Bound<PyAny>) -> PyResult<FFI_LogicalExtensionCodec> {
@@ -235,10 +235,9 @@ pub fn create_physical_extension_capsule<'py>(
     py: Python<'py>,
     codec: &FFI_PhysicalExtensionCodec,
 ) -> PyResult<Bound<'py, PyCapsule>> {
-    let name = cr"datafusion_physical_extension_codec".into();
     let codec = codec.clone();
 
-    PyCapsule::new(py, codec, Some(name))
+    PyCapsule::new_with_value(py, codec, cr"datafusion_physical_extension_codec")
 }
 
 /// Define a `<fn_name>(obj) -> PyResult<Arc<$output_type>>` extractor that

examples/datafusion-ffi-example/src/aggregate_udf.rs

@@ -50,12 +50,10 @@ impl MySumUDF {
         &self,
         py: Python<'py>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
-        let name = cr"datafusion_aggregate_udf".into();
-
         let func = Arc::new(AggregateUDF::from(self.clone()));
         let provider = FFI_AggregateUDF::from(func);
 
-        PyCapsule::new(py, provider, Some(name))
+        PyCapsule::new_with_value(py, provider, cr"datafusion_aggregate_udf")
     }
 }