Expose lit_with_metadata and add unit test

Author: timsaucer

python/datafusion/__init__.py

@@ -21,6 +21,10 @@
 See https://datafusion.apache.org/python for more information.
 """
 
+from __future__ import annotations
+
+from typing import Any
+
 try:
     import importlib.metadata as importlib_metadata
 except ImportError:
@@ -130,3 +134,18 @@ def str_lit(value):
 def lit(value) -> Expr:
     """Create a literal expression."""
     return Expr.literal(value)
+
+
+def literal_with_metadata(value: Any, metadata: dict[str, str]) -> Expr:
+    """Creates a new expression representing a scalar value with metadata.
+
+    Args:
+        value: A valid PyArrow scalar value or easily castable to one.
+        metadata: Metadata to attach to the expression.
+    """
+    return Expr.literal_with_metadata(value, metadata)
+
+
+def lit_with_metadata(value: Any, metadata: dict[str, str]) -> Expr:
+    """Alias for literal_with_metadata."""
+    return literal_with_metadata(value, metadata)

python/datafusion/expr.py

@@ -435,6 +435,20 @@ def literal(value: Any) -> Expr:
             value = pa.scalar(value)
         return Expr(expr_internal.RawExpr.literal(value))
 
+    @staticmethod
+    def literal_with_metadata(value: Any, metadata: dict[str, str]) -> Expr:
+        """Creates a new expression representing a scalar value with metadata.
+
+        Args:
+            value: A valid PyArrow scalar value or easily castable to one.
+            metadata: Metadata to attach to the expression.
+        """
+        if isinstance(value, str):
+            value = pa.scalar(value, type=pa.string_view())
+        value = value if isinstance(value, pa.Scalar) else pa.scalar(value)
+
+        return Expr(expr_internal.RawExpr.literal_with_metadata(value, metadata))
+
     @staticmethod
     def string_literal(value: str) -> Expr:
         """Creates a new expression representing a UTF8 literal value.

python/tests/test_expr.py

@@ -19,7 +19,14 @@
 
 import pyarrow as pa
 import pytest
-from datafusion import SessionContext, col, functions, lit
+from datafusion import (
+    SessionContext,
+    col,
+    functions,
+    lit,
+    lit_with_metadata,
+    literal_with_metadata,
+)
 from datafusion.expr import (
     Aggregate,
     AggregateFunction,
@@ -824,3 +831,52 @@ def test_expr_functions(ctx, function, expected_result):
 
     assert len(result) == 1
     assert result[0].column(0).equals(expected_result)
+
+
+def test_literal_metadata(ctx):
+    result = (
+        ctx.from_pydict({"a": [1]})
+        .select(
+            lit(1).alias("no_metadata"),
+            lit_with_metadata(2, {"key1": "value1"}).alias("lit_with_metadata_fn"),
+            literal_with_metadata(3, {"key2": "value2"}).alias(
+                "literal_with_metadata_fn"
+            ),
+        )
+        .collect()
+    )
+
+    expected_schema = pa.schema(
+        [
+            pa.field("no_metadata", pa.int64(), nullable=False),
+            pa.field(
+                "lit_with_metadata_fn",
+                pa.int64(),
+                nullable=False,
+                metadata={"key1": "value1"},
+            ),
+            pa.field(
+                "literal_with_metadata_fn",
+                pa.int64(),
+                nullable=False,
+                metadata={"key2": "value2"},
+            ),
+        ]
+    )
+
+    expected = pa.RecordBatch.from_pydict(
+        {
+            "no_metadata": pa.array([1]),
+            "lit_with_metadata_fn": pa.array([2]),
+            "literal_with_metadata_fn": pa.array([3]),
+        },
+        schema=expected_schema,
+    )
+
+    assert result[0] == expected
+
+    # Testing result[0].schema == expected_schema does not check each key/value pair
+    # so we want to explicitly test these
+    for expected_field in expected_schema:
+        actual_field = result[0].schema.field(expected_field.name)
+        assert expected_field.metadata == actual_field.metadata

src/expr.rs

@@ -18,7 +18,7 @@
 use datafusion::logical_expr::expr::AggregateFunctionParams;
 use datafusion::logical_expr::utils::exprlist_to_fields;
 use datafusion::logical_expr::{
-    ExprFuncBuilder, ExprFunctionExt, LogicalPlan, WindowFunctionDefinition,
+    lit_with_metadata, ExprFuncBuilder, ExprFunctionExt, LogicalPlan, WindowFunctionDefinition,
 };
 use pyo3::IntoPyObjectExt;
 use pyo3::{basic::CompareOp, prelude::*};
@@ -282,6 +282,14 @@ impl PyExpr {
         lit(value.0).into()
     }
 
+    #[staticmethod]
+    pub fn literal_with_metadata(
+        value: PyScalarValue,
+        metadata: HashMap<String, String>,
+    ) -> PyExpr {
+        lit_with_metadata(value.0, metadata).into()
+    }
+
     #[staticmethod]
     pub fn column(value: &str) -> PyExpr {
         col(value).into()