Make map the primary function with make_map as alias

map() now supports three calling conventions matching upstream:
- map({"a": 1, "b": 2}) — from a Python dictionary
- map([keys], [values]) — two lists that get zipped
- map(k1, v1, k2, v2, ...) — variadic key-value pairs

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

@@ -3383,49 +3383,57 @@ def empty(array: Expr) -> Expr:
 # map functions
 
 
-def make_map(
-    data: dict[Any, Any] | None = None,
-    keys: list[Any] | None = None,
-    values: list[Any] | None = None,
-) -> Expr:
+def map(*args: Any) -> 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.
+    Supports three calling conventions:
 
-    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``).
+    - ``map({"a": 1, "b": 2})`` — from a Python dictionary.
+    - ``map([keys], [values])`` — two lists that get zipped.
+    - ``map(k1, v1, k2, v2, ...)`` — variadic key-value pairs.
+
+    Keys and values that are not already :py:class:`~datafusion.expr.Expr`
+    are automatically converted to literal expressions.
 
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"a": [1]})
         >>> result = df.select(
-        ...     dfn.functions.make_map({"a": 1, "b": 2}).alias("map"))
-        >>> result.collect_column("map")[0].as_py()
+        ...     dfn.functions.map({"a": 1, "b": 2}).alias("m"))
+        >>> result.collect_column("m")[0].as_py()
         [('a', 1), ('b', 2)]
     """
-    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
+    if len(args) == 1 and isinstance(args[0], dict):
+        key_list = list(args[0].keys())
+        value_list = list(args[0].values())
+    elif (
+        len(args) == 2  # noqa: PLR2004
+        and isinstance(args[0], list)
+        and isinstance(args[1], list)
+    ):
+        key_list = args[0]
+        value_list = args[1]
+    elif len(args) >= 2 and len(args) % 2 == 0:  # noqa: PLR2004
+        key_list = list(args[0::2])
+        value_list = list(args[1::2])
     else:
-        msg = "Must specify either data or both keys and values"
+        msg = "map expects a dict, two lists, or an even number of key-value arguments"
         raise ValueError(msg)
 
     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 make_map(*args: Any) -> Expr:
+    """Returns a map expression.
+
+    See Also:
+        This is an alias for :py:func:`map`.
+    """
+    return map(*args)
+
+
 def map_keys(map: Expr) -> Expr:
     """Returns a list of all keys in the map.
 
@@ -3434,7 +3442,7 @@ def map_keys(map: Expr) -> Expr:
         >>> df = ctx.from_pydict({"a": [1]})
         >>> result = df.select(
         ...     dfn.functions.map_keys(
-        ...         dfn.functions.make_map({"x": 1, "y": 2})
+        ...         dfn.functions.map({"x": 1, "y": 2})
         ...     ).alias("keys"))
         >>> result.collect_column("keys")[0].as_py()
         ['x', 'y']
@@ -3450,7 +3458,7 @@ def map_values(map: Expr) -> Expr:
         >>> df = ctx.from_pydict({"a": [1]})
         >>> result = df.select(
         ...     dfn.functions.map_values(
-        ...         dfn.functions.make_map({"x": 1, "y": 2})
+        ...         dfn.functions.map({"x": 1, "y": 2})
         ...     ).alias("vals"))
         >>> result.collect_column("vals")[0].as_py()
         [1, 2]
@@ -3466,7 +3474,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({"x": 1, "y": 2}),
+        ...         dfn.functions.map({"x": 1, "y": 2}),
         ...         dfn.lit("x"),
         ...     ).alias("val"))
         >>> result.collect_column("val")[0].as_py()
@@ -3483,7 +3491,7 @@ def map_entries(map: Expr) -> Expr:
         >>> df = ctx.from_pydict({"a": [1]})
         >>> result = df.select(
         ...     dfn.functions.map_entries(
-        ...         dfn.functions.make_map({"x": 1, "y": 2})
+        ...         dfn.functions.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

@@ -668,29 +668,29 @@ def test_array_function_obj_tests(stmt, py_expr):
         assert a == b
 
 
-def test_make_map():
+def test_map_from_dict():
     ctx = SessionContext()
     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)
+    result = df.select(f.map({"x": 1, "y": 2}).alias("m")).collect()[0].column(0)
     assert result[0].as_py() == [("x", 1), ("y", 2)]
 
 
-def test_make_map_with_expr_values():
+def test_map_from_dict_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({"x": literal(1), "y": literal(2)}).alias("map"))
+        df.select(f.map({"x": literal(1), "y": literal(2)}).alias("m"))
         .collect()[0]
         .column(0)
     )
     assert result[0].as_py() == [("x", 1), ("y", 2)]
 
 
-def test_make_map_with_column_data():
+def test_map_from_two_lists():
     ctx = SessionContext()
     batch = pa.RecordBatch.from_arrays(
         [
@@ -701,7 +701,7 @@ def test_make_map_with_column_data():
     )
     df = ctx.create_dataframe([[batch]])
 
-    m = f.make_map(keys=[column("keys")], values=[column("vals")])
+    m = f.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]
@@ -711,12 +711,48 @@ def test_make_map_with_column_data():
         assert result[i].as_py() == [expected]
 
 
+def test_map_from_variadic_pairs():
+    ctx = SessionContext()
+    batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
+    df = ctx.create_dataframe([[batch]])
+
+    result = df.select(f.map("x", 1, "y", 2).alias("m")).collect()[0].column(0)
+    assert result[0].as_py() == [("x", 1), ("y", 2)]
+
+
+def test_map_variadic_with_exprs():
+    ctx = SessionContext()
+    batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
+    df = ctx.create_dataframe([[batch]])
+
+    result = (
+        df.select(f.map(literal("x"), literal(1), literal("y"), literal(2)).alias("m"))
+        .collect()[0]
+        .column(0)
+    )
+    assert result[0].as_py() == [("x", 1), ("y", 2)]
+
+
+def test_map_odd_args_raises():
+    with pytest.raises(ValueError, match="map expects"):
+        f.map("x", 1, "y")
+
+
+def test_make_map_is_alias():
+    ctx = SessionContext()
+    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("m")).collect()[0].column(0)
+    assert result[0].as_py() == [("x", 1), ("y", 2)]
+
+
 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({"x": 1, "y": 2})
+    m = f.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"]
 
@@ -726,7 +762,7 @@ def test_map_values():
     batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
     df = ctx.create_dataframe([[batch]])
 
-    m = f.make_map({"x": 1, "y": 2})
+    m = f.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]
 
@@ -736,7 +772,7 @@ def test_map_extract():
     batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
     df = ctx.create_dataframe([[batch]])
 
-    m = f.make_map({"x": 1, "y": 2})
+    m = f.map({"x": 1, "y": 2})
     result = (
         df.select(f.map_extract(m, literal("x")).alias("val")).collect()[0].column(0)
     )
@@ -748,7 +784,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({"x": 1})
+    m = f.map({"x": 1})
     result = (
         df.select(f.map_extract(m, literal("z")).alias("val")).collect()[0].column(0)
     )
@@ -760,7 +796,7 @@ def test_map_entries():
     batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
     df = ctx.create_dataframe([[batch]])
 
-    m = f.make_map({"x": 1, "y": 2})
+    m = f.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},
@@ -773,7 +809,7 @@ def test_element_at():
     batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
     df = ctx.create_dataframe([[batch]])
 
-    m = f.make_map({"a": 10, "b": 20})
+    m = f.map({"a": 10, "b": 20})
     result = (
         df.select(f.element_at(m, literal("b")).alias("val")).collect()[0].column(0)
     )