feat: update Expr.literal to handle None and convert values to arrays

Author: kosiew

python/datafusion/expr.py

@@ -444,11 +444,14 @@ def literal(value: Any) -> Expr:
 
         ``value`` must be a valid PyArrow scalar value or easily castable to one.
         """
-        if isinstance(value, str):
-            value = pa.scalar(value, type=pa.string_view())
-        if not isinstance(value, pa.Scalar):
-            value = pa.scalar(value)
-        return Expr(expr_internal.RawExpr.literal(value))
+        if value is None:
+            array = pa.array([None])
+        elif isinstance(value, str):
+            array = pa.array([value], type=pa.string_view())
+        else:
+            array = pa.array([value])
+
+        return Expr(expr_internal.RawExpr.literal(array))
 
     @staticmethod
     def literal_with_metadata(value: Any, metadata: dict[str, str]) -> Expr:
@@ -458,11 +461,14 @@ def literal_with_metadata(value: Any, metadata: dict[str, str]) -> Expr:
             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)
+        if value is None:
+            array = pa.array([None])
+        elif isinstance(value, str):
+            array = pa.array([value], type=pa.string_view())
+        else:
+            array = pa.array([value])
 
-        return Expr(expr_internal.RawExpr.literal_with_metadata(value, metadata))
+        return Expr(expr_internal.RawExpr.literal_with_metadata(array, metadata))
 
     @staticmethod
     def string_literal(value: str) -> Expr:
@@ -476,8 +482,8 @@ def string_literal(value: str) -> Expr:
         https://github.com/apache/datafusion/blob/86740bfd3d9831d6b7c1d0e1bf4a21d91598a0ac/datafusion/functions/src/core/arrow_cast.rs#L179
         """
         if isinstance(value, str):
-            value = pa.scalar(value, type=pa.string())
-            return Expr(expr_internal.RawExpr.literal(value))
+            array = pa.array([value], type=pa.string())
+            return Expr(expr_internal.RawExpr.literal(array))
         return Expr.literal(value)
 
     @staticmethod

python/tests/test_dataframe.py

@@ -302,6 +302,13 @@ def test_limit_with_offset(df):
     assert len(result.column(1)) == 1
 
 
+def test_literal_numeric_and_string(df):
+    df = df.select(literal(10), literal("foo"))
+    result = df.collect()[0]
+    assert result.column(0) == pa.array([10, 10, 10])
+    assert result.column(1) == pa.array(["foo", "foo", "foo"], type=pa.string_view())
+
+
 def test_head(df):
     df = df.head(1)
 

python/tests/test_functions.py

@@ -22,6 +22,7 @@
 import pytest
 from datafusion import SessionContext, column, literal, string_literal
 from datafusion import functions as f
+from datafusion.expr import Expr
 
 np.seterr(invalid="ignore")
 
@@ -110,6 +111,14 @@ def test_lit_arith(df):
     )
 
 
+def test_expr_literal_numeric_and_string():
+    """Ensure Expr.literal handles numeric and string values."""
+    num_expr = Expr.literal(7)
+    str_expr = Expr.literal("bar")
+    assert num_expr.to_variant() == 7
+    assert str_expr.to_variant() == "bar"
+
+
 def test_math_functions():
     ctx = SessionContext()
     # create a RecordBatch and a new DataFrame from it