Upgrade to DF 48

Author: timsaucer

Cargo.lock

@@ -37,11 +37,11 @@ substrait = ["dep:datafusion-substrait"]
 tokio = { version = "1.45", features = ["macros", "rt", "rt-multi-thread", "sync"] }
 pyo3 = { version = "0.24", features = ["extension-module", "abi3", "abi3-py39"] }
 pyo3-async-runtimes = { version = "0.24", features = ["tokio-runtime"]}
-arrow = { version = "55.0.0", features = ["pyarrow"] }
-datafusion = { version = "47.0.0", features = ["avro", "unicode_expressions"] }
-datafusion-substrait = { version = "47.0.0", optional = true }
-datafusion-proto = { version = "47.0.0" }
-datafusion-ffi = { version = "47.0.0" }
+arrow = { version = "55.1.0", features = ["pyarrow"] }
+datafusion = { version = "48.0.0", features = ["avro", "unicode_expressions"] }
+datafusion-substrait = { version = "48.0.0", optional = true }
+datafusion-proto = { version = "48.0.0" }
+datafusion-ffi = { version = "48.0.0" }
 prost = "0.13.1" # keep in line with `datafusion-substrait`
 uuid = { version = "1.16", features = ["v4"] }
 mimalloc = { version = "0.1", optional = true, default-features = false, features = ["local_dynamic_tls"] }
@@ -61,3 +61,10 @@ crate-type = ["cdylib", "rlib"]
 [profile.release]
 lto = true
 codegen-units = 1
+
+# TODO remove before merging to main
+[patch.crates-io]
+datafusion = { git = "https://github.com/apache/datafusion", rev = "33a32d4382bee7e3c705d0f55d05c24a115a2f98" }
+datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "33a32d4382bee7e3c705d0f55d05c24a115a2f98" }
+datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "33a32d4382bee7e3c705d0f55d05c24a115a2f98" }
+datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "33a32d4382bee7e3c705d0f55d05c24a115a2f98" }

Cargo.toml

@@ -24,6 +24,7 @@ use pyo3::IntoPyObjectExt;
 use pyo3::{basic::CompareOp, prelude::*};
 use std::collections::HashMap;
 use std::convert::{From, Into};
+use std::env::args;
 use std::sync::Arc;
 use window::PyWindowFrame;
 
@@ -150,7 +151,7 @@ impl PyExpr {
                 Ok(PyScalarVariable::new(data_type, variables).into_bound_py_any(py)?)
             }
             Expr::Like(value) => Ok(PyLike::from(value.clone()).into_bound_py_any(py)?),
-            Expr::Literal(value) => Ok(PyLiteral::from(value.clone()).into_bound_py_any(py)?),
+            Expr::Literal(value, metadata) => Ok(PyLiteral::new_with_metadata(value.clone(), metadata.clone()).into_bound_py_any(py)?),
             Expr::BinaryExpr(expr) => Ok(PyBinaryExpr::from(expr.clone()).into_bound_py_any(py)?),
             Expr::Not(expr) => Ok(PyNot::new(*expr.clone()).into_bound_py_any(py)?),
             Expr::IsNotNull(expr) => Ok(PyIsNotNull::new(*expr.clone()).into_bound_py_any(py)?),
@@ -377,7 +378,7 @@ impl PyExpr {
     /// Extracts the Expr value into a PyObject that can be shared with Python
     pub fn python_value(&self, py: Python) -> PyResult<PyObject> {
         match &self.expr {
-            Expr::Literal(scalar_value) => scalar_to_pyarrow(scalar_value, py),
+            Expr::Literal(scalar_value, _) => scalar_to_pyarrow(scalar_value, py),
             _ => Err(py_type_err(format!(
                 "Non Expr::Literal encountered in types: {:?}",
                 &self.expr
@@ -417,11 +418,13 @@ impl PyExpr {
                 params: AggregateFunctionParams { args, .. },
                 ..
             })
-            | Expr::ScalarFunction(ScalarFunction { args, .. })
-            | Expr::WindowFunction(WindowFunction {
-                params: WindowFunctionParams { args, .. },
-                ..
-            }) => Ok(args.iter().map(|arg| PyExpr::from(arg.clone())).collect()),
+            | Expr::ScalarFunction(ScalarFunction { args, .. }) => {
+                Ok(args.iter().map(|arg| PyExpr::from(arg.clone())).collect())
+            }
+            Expr::WindowFunction(boxed_window_fn) => {
+                let args = &boxed_window_fn.params.args;
+                Ok(args.iter().map(|arg| PyExpr::from(arg.clone())).collect())
+            }
 
             // Expr(s) that require more specific processing
             Expr::Case(Case {
@@ -600,10 +603,10 @@ impl PyExpr {
     ) -> PyDataFusionResult<PyExpr> {
         match &self.expr {
             Expr::AggregateFunction(agg_fn) => {
-                let window_fn = Expr::WindowFunction(WindowFunction::new(
+                let window_fn = Expr::WindowFunction(Box::new(WindowFunction::new(
                     WindowFunctionDefinition::AggregateUDF(agg_fn.func.clone()),
                     agg_fn.params.args.clone(),
-                ));
+                )));
 
                 add_builder_fns_to_window(
                     window_fn,
@@ -743,7 +746,7 @@ impl PyExpr {
                 | Operator::QuestionPipe => Err(py_type_err(format!("Unsupported expr: ${op}"))),
             },
             Expr::Cast(Cast { expr: _, data_type }) => DataTypeMap::map_from_arrow_type(data_type),
-            Expr::Literal(scalar_value) => DataTypeMap::map_from_scalar_value(scalar_value),
+            Expr::Literal(scalar_value, _) => DataTypeMap::map_from_scalar_value(scalar_value),
             _ => Err(py_type_err(format!(
                 "Non Expr::Literal encountered in types: {:?}",
                 expr

src/expr.rs

@@ -18,11 +18,23 @@
 use crate::errors::PyDataFusionError;
 use datafusion::common::ScalarValue;
 use pyo3::{prelude::*, IntoPyObjectExt};
+use std::collections::{BTreeMap, HashMap};
+use std::sync::Arc;
 
 #[pyclass(name = "Literal", module = "datafusion.expr", subclass)]
 #[derive(Clone)]
 pub struct PyLiteral {
     pub value: ScalarValue,
+    pub metadata: Option<BTreeMap<String, String>>,
+}
+
+impl PyLiteral {
+    pub fn new_with_metadata(
+        value: ScalarValue,
+        metadata: Option<BTreeMap<String, String>>,
+    ) -> PyLiteral {
+        Self { value, metadata }
+    }
 }
 
 impl From<PyLiteral> for ScalarValue {
@@ -33,7 +45,10 @@ impl From<PyLiteral> for ScalarValue {
 
 impl From<ScalarValue> for PyLiteral {
     fn from(value: ScalarValue) -> PyLiteral {
-        PyLiteral { value }
+        PyLiteral {
+            value,
+            metadata: None,
+        }
     }
 }
 

src/expr/literal.rs

@@ -118,51 +118,45 @@ impl PyWindowExpr {
     /// Returns order by columns in a window function expression
     pub fn get_sort_exprs(&self, expr: PyExpr) -> PyResult<Vec<PySortExpr>> {
         match expr.expr.unalias() {
-            Expr::WindowFunction(WindowFunction {
-                params: WindowFunctionParams { order_by, .. },
-                ..
-            }) => py_sort_expr_list(&order_by),
+            Expr::WindowFunction(boxed_window_fn) => {
+                py_sort_expr_list(&boxed_window_fn.params.order_by)
+            }
             other => Err(not_window_function_err(other)),
         }
     }
 
     /// Return partition by columns in a window function expression
     pub fn get_partition_exprs(&self, expr: PyExpr) -> PyResult<Vec<PyExpr>> {
         match expr.expr.unalias() {
-            Expr::WindowFunction(WindowFunction {
-                params: WindowFunctionParams { partition_by, .. },
-                ..
-            }) => py_expr_list(&partition_by),
+            Expr::WindowFunction(boxed_window_fn) => {
+                py_expr_list(&boxed_window_fn.params.partition_by)
+            }
             other => Err(not_window_function_err(other)),
         }
     }
 
     /// Return input args for window function
     pub fn get_args(&self, expr: PyExpr) -> PyResult<Vec<PyExpr>> {
         match expr.expr.unalias() {
-            Expr::WindowFunction(WindowFunction {
-                params: WindowFunctionParams { args, .. },
-                ..
-            }) => py_expr_list(&args),
+            Expr::WindowFunction(boxed_window_fn) => py_expr_list(&boxed_window_fn.params.args),
             other => Err(not_window_function_err(other)),
         }
     }
 
     /// Return window function name
     pub fn window_func_name(&self, expr: PyExpr) -> PyResult<String> {
         match expr.expr.unalias() {
-            Expr::WindowFunction(WindowFunction { fun, .. }) => Ok(fun.to_string()),
+            Expr::WindowFunction(boxed_window_fn) => Ok(boxed_window_fn.fun.to_string()),
             other => Err(not_window_function_err(other)),
         }
     }
 
     /// Returns a Pywindow frame for a given window function expression
     pub fn get_frame(&self, expr: PyExpr) -> Option<PyWindowFrame> {
         match expr.expr.unalias() {
-            Expr::WindowFunction(WindowFunction {
-                params: WindowFunctionParams { window_frame, .. },
-                ..
-            }) => Some(window_frame.into()),
+            Expr::WindowFunction(boxed_window_fn) => {
+                Some(boxed_window_fn.params.window_frame.into())
+            }
             _ => None,
         }
     }

src/expr/window.rs

@@ -103,7 +103,7 @@ fn array_cat(exprs: Vec<PyExpr>) -> PyExpr {
 #[pyo3(signature = (array, element, index=None))]
 fn array_position(array: PyExpr, element: PyExpr, index: Option<i64>) -> PyExpr {
     let index = ScalarValue::Int64(index);
-    let index = Expr::Literal(index);
+    let index = Expr::Literal(index, None);
     datafusion::functions_nested::expr_fn::array_position(array.into(), element.into(), index)
         .into()
 }
@@ -334,7 +334,7 @@ fn window(
         .unwrap_or(WindowFrame::new(order_by.as_ref().map(|v| !v.is_empty())));
 
     Ok(PyExpr {
-        expr: datafusion::logical_expr::Expr::WindowFunction(WindowFunction {
+        expr: datafusion::logical_expr::Expr::WindowFunction(Box::new(WindowFunction {
             fun,
             params: WindowFunctionParams {
                 args: args.into_iter().map(|x| x.expr).collect::<Vec<_>>(),
@@ -351,7 +351,7 @@ fn window(
                 window_frame,
                 null_treatment: None,
             },
-        }),
+        })),
     })
 }
 

src/functions.rs

@@ -61,7 +61,7 @@ fn extract_scalar_list<'py>(
         .iter()
         .map(|expr| match expr {
             // TODO: should we also leverage `ScalarValue::to_pyarrow` here?
-            Expr::Literal(v) => match v {
+            Expr::Literal(v, _) => match v {
                 // The unwraps here are for infallible conversions
                 ScalarValue::Boolean(Some(b)) => Ok(b.into_bound_py_any(py)?),
                 ScalarValue::Int8(Some(i)) => Ok(i.into_bound_py_any(py)?),
@@ -106,7 +106,7 @@ impl TryFrom<&Expr> for PyArrowFilterExpression {
             let op_module = Python::import(py, "operator")?;
             let pc_expr: PyDataFusionResult<Bound<'_, PyAny>> = match expr {
                 Expr::Column(Column { name, .. }) => Ok(pc.getattr("field")?.call1((name,))?),
-                Expr::Literal(scalar) => Ok(scalar_to_pyarrow(scalar, py)?.into_bound(py)),
+                Expr::Literal(scalar, _) => Ok(scalar_to_pyarrow(scalar, py)?.into_bound(py)),
                 Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
                     let operator = operator_to_py(op, &op_module)?;
                     let left = PyArrowFilterExpression::try_from(left.as_ref())?.0;

src/pyarrow_filter_expression.rs

@@ -300,13 +300,9 @@ impl WindowUDFImpl for MultiColumnWindowUDF {
         &self.signature
     }
 
-    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<arrow::datatypes::Field> {
+    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<arrow::datatypes::FieldRef> {
         // TODO: Should nullable always be `true`?
-        Ok(arrow::datatypes::Field::new(
-            field_args.name(),
-            self.return_type.clone(),
-            true,
-        ))
+        Ok(arrow::datatypes::Field::new(field_args.name(), self.return_type.clone(), true).into())
     }
 
     // TODO: Enable passing partition_evaluator_args to python?