Change make_map to accept a Python dictionary

make_map now takes a dict for the common case and also supports
separate keys/values lists for column expressions. Non-Expr keys
and values are automatically converted to literals.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: timsaucer

python/datafusion/functions.py

@@ -18,7 +18,6 @@
 
 from __future__ import annotations
 
-import builtins
 from typing import TYPE_CHECKING, Any
 
 import pyarrow as pa
@@ -3348,29 +3347,47 @@ def empty(array: Expr) -> Expr:
 # map functions
 
 
-def make_map(*args: Expr) -> Expr:
-    """Returns a map created from key and value expressions.
+def make_map(
+    data: dict[Any, Any] | None = None,
+    keys: list[Any] | None = None,
+    values: list[Any] | None = None,
+) -> Expr:
+    """Returns a map expression.
+
+    Can be called with either a Python dictionary or separate ``keys``
+    and ``values`` lists.  Keys and values that are not already
+    :py:class:`~datafusion.expr.Expr` are automatically converted to
+    literal expressions.
 
-    Accepts an even number of arguments, alternating between keys and values.
-    For example, ``make_map(k1, v1, k2, v2)`` creates a map ``{k1: v1, k2: v2}``.
+    Args:
+        data: A Python dictionary of key-value pairs.
+        keys: A list of keys (use with ``values`` for column expressions).
+        values: A list of values (use with ``keys``).
 
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"a": [1]})
         >>> result = df.select(
-        ...     dfn.functions.make_map(
-        ...         dfn.lit("a"), dfn.lit(1),
-        ...         dfn.lit("b"), dfn.lit(2),
-        ...     ).alias("map"))
+        ...     dfn.functions.make_map({"a": 1, "b": 2}).alias("map"))
         >>> result.collect_column("map")[0].as_py()
         [('a', 1), ('b', 2)]
     """
-    if len(args) % 2 != 0:
-        msg = "make_map requires an even number of arguments"
+    if data is not None:
+        if keys is not None or values is not None:
+            msg = "Cannot specify both data and keys/values"
+            raise ValueError(msg)
+        key_list = list(data.keys())
+        value_list = list(data.values())
+    elif keys is not None and values is not None:
+        key_list = keys
+        value_list = values
+    else:
+        msg = "Must specify either data or both keys and values"
         raise ValueError(msg)
-    keys = [args[i].expr for i in builtins.range(0, len(args), 2)]
-    values = [args[i].expr for i in builtins.range(1, len(args), 2)]
-    return Expr(f.make_map(keys, values))
+
+    key_exprs = [k if isinstance(k, Expr) else Expr.literal(k) for k in key_list]
+    val_exprs = [v if isinstance(v, Expr) else Expr.literal(v) for v in value_list]
+    return Expr(f.make_map([k.expr for k in key_exprs], [v.expr for v in val_exprs]))
 
 
 def map_keys(map: Expr) -> Expr:
@@ -3381,10 +3398,7 @@ def map_keys(map: Expr) -> Expr:
         >>> df = ctx.from_pydict({"a": [1]})
         >>> result = df.select(
         ...     dfn.functions.map_keys(
-        ...         dfn.functions.make_map(
-        ...             dfn.lit("x"), dfn.lit(1),
-        ...             dfn.lit("y"), dfn.lit(2),
-        ...         )
+        ...         dfn.functions.make_map({"x": 1, "y": 2})
         ...     ).alias("keys"))
         >>> result.collect_column("keys")[0].as_py()
         ['x', 'y']
@@ -3400,10 +3414,7 @@ def map_values(map: Expr) -> Expr:
         >>> df = ctx.from_pydict({"a": [1]})
         >>> result = df.select(
         ...     dfn.functions.map_values(
-        ...         dfn.functions.make_map(
-        ...             dfn.lit("x"), dfn.lit(1),
-        ...             dfn.lit("y"), dfn.lit(2),
-        ...         )
+        ...         dfn.functions.make_map({"x": 1, "y": 2})
         ...     ).alias("vals"))
         >>> result.collect_column("vals")[0].as_py()
         [1, 2]
@@ -3419,10 +3430,7 @@ def map_extract(map: Expr, key: Expr) -> Expr:
         >>> df = ctx.from_pydict({"a": [1]})
         >>> result = df.select(
         ...     dfn.functions.map_extract(
-        ...         dfn.functions.make_map(
-        ...             dfn.lit("x"), dfn.lit(1),
-        ...             dfn.lit("y"), dfn.lit(2),
-        ...         ),
+        ...         dfn.functions.make_map({"x": 1, "y": 2}),
         ...         dfn.lit("x"),
         ...     ).alias("val"))
         >>> result.collect_column("val")[0].as_py()
@@ -3439,10 +3447,7 @@ def map_entries(map: Expr) -> Expr:
         >>> df = ctx.from_pydict({"a": [1]})
         >>> result = df.select(
         ...     dfn.functions.map_entries(
-        ...         dfn.functions.make_map(
-        ...             dfn.lit("x"), dfn.lit(1),
-        ...             dfn.lit("y"), dfn.lit(2),
-        ...         )
+        ...         dfn.functions.make_map({"x": 1, "y": 2})
         ...     ).alias("entries"))
         >>> result.collect_column("entries")[0].as_py()
         [{'key': 'x', 'value': 1}, {'key': 'y', 'value': 2}]

python/tests/test_functions.py

@@ -673,32 +673,50 @@ def test_make_map():
     batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
     df = ctx.create_dataframe([[batch]])
 
+    result = df.select(f.make_map({"x": 1, "y": 2}).alias("map")).collect()[0].column(0)
+    assert result[0].as_py() == [("x", 1), ("y", 2)]
+
+
+def test_make_map_with_expr_values():
+    ctx = SessionContext()
+    batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
+    df = ctx.create_dataframe([[batch]])
+
     result = (
-        df.select(
-            f.make_map(
-                literal("x"),
-                literal(1),
-                literal("y"),
-                literal(2),
-            ).alias("map")
-        )
+        df.select(f.make_map({"x": literal(1), "y": literal(2)}).alias("map"))
         .collect()[0]
         .column(0)
     )
     assert result[0].as_py() == [("x", 1), ("y", 2)]
 
 
-def test_make_map_odd_args():
-    with pytest.raises(ValueError, match="even number of arguments"):
-        f.make_map(literal("x"), literal(1), literal("y"))
+def test_make_map_with_column_data():
+    ctx = SessionContext()
+    batch = pa.RecordBatch.from_arrays(
+        [
+            pa.array(["k1", "k2", "k3"]),
+            pa.array([10, 20, 30]),
+        ],
+        names=["keys", "vals"],
+    )
+    df = ctx.create_dataframe([[batch]])
+
+    m = f.make_map(keys=[column("keys")], values=[column("vals")])
+    result = df.select(f.map_keys(m).alias("k")).collect()[0].column(0)
+    for i, expected in enumerate(["k1", "k2", "k3"]):
+        assert result[i].as_py() == [expected]
+
+    result = df.select(f.map_values(m).alias("v")).collect()[0].column(0)
+    for i, expected in enumerate([10, 20, 30]):
+        assert result[i].as_py() == [expected]
 
 
 def test_map_keys():
     ctx = SessionContext()
     batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
     df = ctx.create_dataframe([[batch]])
 
-    m = f.make_map(literal("x"), literal(1), literal("y"), literal(2))
+    m = f.make_map({"x": 1, "y": 2})
     result = df.select(f.map_keys(m).alias("keys")).collect()[0].column(0)
     assert result[0].as_py() == ["x", "y"]
 
@@ -708,7 +726,7 @@ def test_map_values():
     batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
     df = ctx.create_dataframe([[batch]])
 
-    m = f.make_map(literal("x"), literal(1), literal("y"), literal(2))
+    m = f.make_map({"x": 1, "y": 2})
     result = df.select(f.map_values(m).alias("vals")).collect()[0].column(0)
     assert result[0].as_py() == [1, 2]
 
@@ -718,7 +736,7 @@ def test_map_extract():
     batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
     df = ctx.create_dataframe([[batch]])
 
-    m = f.make_map(literal("x"), literal(1), literal("y"), literal(2))
+    m = f.make_map({"x": 1, "y": 2})
     result = (
         df.select(f.map_extract(m, literal("x")).alias("val")).collect()[0].column(0)
     )
@@ -730,7 +748,7 @@ def test_map_extract_missing_key():
     batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
     df = ctx.create_dataframe([[batch]])
 
-    m = f.make_map(literal("x"), literal(1))
+    m = f.make_map({"x": 1})
     result = (
         df.select(f.map_extract(m, literal("z")).alias("val")).collect()[0].column(0)
     )
@@ -742,7 +760,7 @@ def test_map_entries():
     batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
     df = ctx.create_dataframe([[batch]])
 
-    m = f.make_map(literal("x"), literal(1), literal("y"), literal(2))
+    m = f.make_map({"x": 1, "y": 2})
     result = df.select(f.map_entries(m).alias("entries")).collect()[0].column(0)
     assert result[0].as_py() == [
         {"key": "x", "value": 1},
@@ -755,34 +773,13 @@ def test_element_at():
     batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
     df = ctx.create_dataframe([[batch]])
 
-    m = f.make_map(literal("a"), literal(10), literal("b"), literal(20))
+    m = f.make_map({"a": 10, "b": 20})
     result = (
         df.select(f.element_at(m, literal("b")).alias("val")).collect()[0].column(0)
     )
     assert result[0].as_py() == [20]
 
 
-def test_map_functions_with_column_data():
-    ctx = SessionContext()
-    batch = pa.RecordBatch.from_arrays(
-        [
-            pa.array(["k1", "k2", "k3"]),
-            pa.array([10, 20, 30]),
-        ],
-        names=["keys", "vals"],
-    )
-    df = ctx.create_dataframe([[batch]])
-
-    m = f.make_map(column("keys"), column("vals"))
-    result = df.select(f.map_keys(m).alias("k")).collect()[0].column(0)
-    for i, expected in enumerate(["k1", "k2", "k3"]):
-        assert result[i].as_py() == [expected]
-
-    result = df.select(f.map_values(m).alias("v")).collect()[0].column(0)
-    for i, expected in enumerate([10, 20, 30]):
-        assert result[i].as_py() == [expected]
-
-
 @pytest.mark.parametrize(
     ("function", "expected_result"),
     [