Merge branch 'main' into formatter-options

Author: kosiew

Cargo.lock

@@ -34,25 +34,25 @@ protoc = [ "datafusion-substrait/protoc" ]
 substrait = ["dep:datafusion-substrait"]
 
 [dependencies]
-tokio = { version = "1.43", features = ["macros", "rt", "rt-multi-thread", "sync"] }
-pyo3 = { version = "0.23", features = ["extension-module", "abi3", "abi3-py39"] }
-pyo3-async-runtimes = { version = "0.23", features = ["tokio-runtime"]}
-arrow = { version = "54.2.1", features = ["pyarrow"] }
-datafusion = { version = "46.0.1", features = ["avro", "unicode_expressions"] }
-datafusion-substrait = { version = "46.0.1", optional = true }
-datafusion-proto = { version = "46.0.1" }
-datafusion-ffi = { version = "46.0.1" }
+tokio = { version = "1.44", 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" }
 prost = "0.13.1" # keep in line with `datafusion-substrait`
-uuid = { version = "1.12", features = ["v4"] }
+uuid = { version = "1.16", features = ["v4"] }
 mimalloc = { version = "0.1", optional = true, default-features = false, features = ["local_dynamic_tls"] }
-async-trait = "0.1.73"
+async-trait = "0.1.88"
 futures = "0.3"
-object_store = { version = "0.11.0", features = ["aws", "gcp", "azure", "http"] }
+object_store = { version = "0.12.0", features = ["aws", "gcp", "azure", "http"] }
 url = "2"
 
 [build-dependencies]
 prost-types = "0.13.1" # keep in line with `datafusion-substrait`
-pyo3-build-config = "0.23"
+pyo3-build-config = "0.24"
 
 [lib]
 name = "datafusion_python"

Cargo.toml

@@ -406,9 +406,17 @@ def column(value: str) -> Expr:
         """Creates a new expression representing a column."""
         return Expr(expr_internal.RawExpr.column(value))
 
-    def alias(self, name: str) -> Expr:
-        """Assign a name to the expression."""
-        return Expr(self.expr.alias(name))
+    def alias(self, name: str, metadata: Optional[dict[str, str]] = None) -> Expr:
+        """Assign a name to the expression.
+
+        Args:
+            name: The name to assign to the expression.
+            metadata: Optional metadata to attach to the expression.
+
+        Returns:
+            A new expression with the assigned name.
+        """
+        return Expr(self.expr.alias(name, metadata))
 
     def sort(self, ascending: bool = True, nulls_first: bool = True) -> SortExpr:
         """Creates a sort :py:class:`Expr` from an existing :py:class:`Expr`.

python/datafusion/expr.py

@@ -372,9 +372,18 @@ def order_by(expr: Expr, ascending: bool = True, nulls_first: bool = True) -> So
     return SortExpr(expr, ascending=ascending, nulls_first=nulls_first)
 
 
-def alias(expr: Expr, name: str) -> Expr:
-    """Creates an alias expression."""
-    return Expr(f.alias(expr.expr, name))
+def alias(expr: Expr, name: str, metadata: Optional[dict[str, str]] = None) -> Expr:
+    """Creates an alias expression with an optional metadata dictionary.
+
+    Args:
+        expr: The expression to alias
+        name: The alias name
+        metadata: Optional metadata to attach to the column
+
+    Returns:
+        An expression with the given alias
+    """
+    return Expr(f.alias(expr.expr, name, metadata))
 
 
 def col(name: str) -> Expr:

python/datafusion/functions.py

@@ -338,7 +338,6 @@ def test_bit_and_bool_fns(df, name, expr, result):
             ),
             [7, 9],
         ),
-        ("last_value", f.last_value(column("a")), [3, 6]),
         (
             "last_value_ordered",
             f.last_value(column("a"), order_by=[column("a").sort(ascending=False)]),

python/tests/test_aggregation.py

@@ -247,3 +247,8 @@ def test_fill_null(df):
     assert result.column(0) == pa.array([1, 2, 100])
     assert result.column(1) == pa.array([4, 25, 6])
     assert result.column(2) == pa.array([1234, 1234, 8])
+
+
+def test_alias_with_metadata(df):
+    df = df.select(col("a").alias("b", {"key": "value"}))
+    assert df.schema().field("b").metadata == {b"key": b"value"}

python/tests/test_expr.py

@@ -1231,3 +1231,8 @@ def test_between_default(df):
 
     actual = df.collect()[0].to_pydict()
     assert actual == expected
+
+
+def test_alias_with_metadata(df):
+    df = df.select(f.alias(f.col("a"), "b", {"key": "value"}))
+    assert df.schema().field("b").metadata == {b"key": b"value"}

python/tests/test_functions.py

@@ -307,7 +307,7 @@ impl PyDataFrame {
 
     #[pyo3(signature = (*args))]
     fn select(&self, args: Vec<PyExpr>) -> PyDataFusionResult<Self> {
-        let expr = args.into_iter().map(|e| e.into()).collect();
+        let expr: Vec<Expr> = args.into_iter().map(|e| e.into()).collect();
         let df = self.df.as_ref().clone().select(expr)?;
         Ok(Self::new(df))
     }

src/dataframe.rs

@@ -275,7 +275,9 @@ impl DisplayAs for DatasetExec {
         Python::with_gil(|py| {
             let number_of_fragments = self.fragments.bind(py).len();
             match t {
-                DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                DisplayFormatType::Default
+                | DisplayFormatType::Verbose
+                | DisplayFormatType::TreeRender => {
                     let projected_columns: Vec<String> = self
                         .schema
                         .fields()

src/dataset_exec.rs

@@ -22,6 +22,7 @@ use datafusion::logical_expr::{
 };
 use pyo3::IntoPyObjectExt;
 use pyo3::{basic::CompareOp, prelude::*};
+use std::collections::HashMap;
 use std::convert::{From, Into};
 use std::sync::Arc;
 use window::PyWindowFrame;
@@ -275,8 +276,9 @@ impl PyExpr {
     }
 
     /// assign a name to the PyExpr
-    pub fn alias(&self, name: &str) -> PyExpr {
-        self.expr.clone().alias(name).into()
+    #[pyo3(signature = (name, metadata=None))]
+    pub fn alias(&self, name: &str, metadata: Option<HashMap<String, String>>) -> PyExpr {
+        self.expr.clone().alias_with_metadata(name, metadata).into()
     }
 
     /// Create a sort PyExpr from an existing PyExpr.
@@ -714,9 +716,19 @@ impl PyExpr {
                 | Operator::BitwiseXor
                 | Operator::BitwiseAnd
                 | Operator::BitwiseOr => DataTypeMap::map_from_arrow_type(&DataType::Binary),
-                Operator::AtArrow | Operator::ArrowAt => {
-                    Err(py_type_err(format!("Unsupported expr: ${op}")))
-                }
+                Operator::AtArrow
+                | Operator::ArrowAt
+                | Operator::Arrow
+                | Operator::LongArrow
+                | Operator::HashArrow
+                | Operator::HashLongArrow
+                | Operator::AtAt
+                | Operator::IntegerDivide
+                | Operator::HashMinus
+                | Operator::AtQuestion
+                | Operator::Question
+                | Operator::QuestionAnd
+                | 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),