Add missing map functions

crates/core/src/functions.rs

@@ -93,6 +93,13 @@ fn array_cat(exprs: Vec<PyExpr>) -> PyExpr {
  array_concat(exprs)
 }
 
+#[pyfunction]
+fn make_map(keys: Vec<PyExpr>, values: Vec<PyExpr>) -> PyExpr {
+ let keys = keys.into_iter().map(|x| x.into()).collect();
+ let values = values.into_iter().map(|x| x.into()).collect();
+ datafusion::functions_nested::map::map(keys, values).into()
+}
+
 #[pyfunction]
 #[pyo3(signature = (array, element, index=None))]
 fn array_position(array: PyExpr, element: PyExpr, index: Option<i64>) -> PyExpr {
@@ -666,6 +673,12 @@ array_fn!(cardinality, array);
 array_fn!(flatten, array);
 array_fn!(range, start stop step);
 
+// Map Functions
+array_fn!(map_keys, map);
+array_fn!(map_values, map);
+array_fn!(map_extract, map key);
+array_fn!(map_entries, map);
+
 aggregate_function!(array_agg);
 aggregate_function!(max);
 aggregate_function!(min);
@@ -1126,6 +1139,13 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
  m.add_wrapped(wrap_pyfunction!(flatten))?;
  m.add_wrapped(wrap_pyfunction!(cardinality))?;
 
+ // Map Functions
+ m.add_wrapped(wrap_pyfunction!(make_map))?;
+ m.add_wrapped(wrap_pyfunction!(map_keys))?;
+ m.add_wrapped(wrap_pyfunction!(map_values))?;
+ m.add_wrapped(wrap_pyfunction!(map_extract))?;
+ m.add_wrapped(wrap_pyfunction!(map_entries))?;
+
  // Window Functions
  m.add_wrapped(wrap_pyfunction!(lead))?;
  m.add_wrapped(wrap_pyfunction!(lag))?;

python/datafusion/functions.py

@@ -139,6 +139,7 @@
  "degrees",
  "dense_rank",
  "digest",
+ "element_at",
  "empty",
  "encode",
  "ends_with",
@@ -202,7 +203,12 @@
  "make_array",
  "make_date",
  "make_list",
+ "make_map",
  "make_time",
+ "map_entries",
+ "map_extract",
+ "map_keys",
+ "map_values",
  "max",
  "md5",
  "mean",
@@ -3374,6 +3380,158 @@ def empty(array: Expr) -> Expr:
  return array_empty(array)
 
 
+# map functions
+
+
+def make_map(*args: Any) -> Expr:
+ """Returns a map expression.
+
+ Supports three calling conventions:
+
+ - ``make_map({"a": 1, "b": 2})`` — from a Python dictionary.
+ - ``make_map([keys], [values])`` — from a list of keys and a list of
+ their associated values. Both lists must be the same length.
+ - ``make_map(k1, v1, k2, v2, ...)`` — from alternating keys and their
+ associated values.
+
+ Keys and values that are not already :py:class:`~datafusion.expr.Expr`
+ are automatically converted to literal expressions.
+
+ Examples:
+ From a dictionary:
+
+ >>> ctx = dfn.SessionContext()
+ >>> df = ctx.from_pydict({"a": [1]})
+ >>> result = df.select(
+ ... dfn.functions.make_map({"a": 1, "b": 2}).alias("m"))
+ >>> result.collect_column("m")[0].as_py()
+ [('a', 1), ('b', 2)]
+
+ From two lists:
+
+ >>> df = ctx.from_pydict({"key": ["x", "y"], "val": [10, 20]})
+ >>> df = df.select(
+ ... dfn.functions.make_map(
+ ... [dfn.col("key")], [dfn.col("val")]
+ ... ).alias("m"))
+ >>> df.collect_column("m")[0].as_py()
+ [('x', 10)]
+
+ From alternating keys and values:
+
+ >>> df = ctx.from_pydict({"a": [1]})
+ >>> result = df.select(
+ ... dfn.functions.make_map("x", 1, "y", 2).alias("m"))
+ >>> result.collect_column("m")[0].as_py()
+ [('x', 1), ('y', 2)]
+ """
+ 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)
+ ):
+ if len(args[0]) != len(args[1]):
+ msg = "make_map requires key and value lists to be the same length"
+ raise ValueError(msg)
+ 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 = (
+ "make_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 map_keys(map: Expr) -> Expr:
+ """Returns a list of all keys in the map.
+
+ Examples:
+ >>> ctx = dfn.SessionContext()
+ >>> df = ctx.from_pydict({"a": [1]})
+ >>> df = df.select(
+ ... dfn.functions.make_map({"x": 1, "y": 2}).alias("m"))
+ >>> result = df.select(
+ ... dfn.functions.map_keys(dfn.col("m")).alias("keys"))
+ >>> result.collect_column("keys")[0].as_py()
+ ['x', 'y']
+ """
+ return Expr(f.map_keys(map.expr))
+
+
+def map_values(map: Expr) -> Expr:
+ """Returns a list of all values in the map.
+
+ Examples:
+ >>> ctx = dfn.SessionContext()
+ >>> df = ctx.from_pydict({"a": [1]})
+ >>> df = df.select(
+ ... dfn.functions.make_map({"x": 1, "y": 2}).alias("m"))
+ >>> result = df.select(
+ ... dfn.functions.map_values(dfn.col("m")).alias("vals"))
+ >>> result.collect_column("vals")[0].as_py()
+ [1, 2]
+ """
+ return Expr(f.map_values(map.expr))
+
+
+def map_extract(map: Expr, key: Expr) -> Expr:
+ """Returns the value for a given key in the map.
+
+ Returns ``[None]`` if the key is absent.
+
+ Examples:
+ >>> ctx = dfn.SessionContext()
+ >>> df = ctx.from_pydict({"a": [1]})
+ >>> df = df.select(
+ ... dfn.functions.make_map({"x": 1, "y": 2}).alias("m"))
+ >>> result = df.select(
+ ... dfn.functions.map_extract(
+ ... dfn.col("m"), dfn.lit("x")
+ ... ).alias("val"))
+ >>> result.collect_column("val")[0].as_py()
+ [1]
+ """
+ return Expr(f.map_extract(map.expr, key.expr))
+
+
+def map_entries(map: Expr) -> Expr:
+ """Returns a list of all entries (key-value struct pairs) in the map.
+
+ Examples:
+ >>> ctx = dfn.SessionContext()
+ >>> df = ctx.from_pydict({"a": [1]})
+ >>> df = df.select(
+ ... dfn.functions.make_map({"x": 1, "y": 2}).alias("m"))
+ >>> result = df.select(
+ ... dfn.functions.map_entries(dfn.col("m")).alias("entries"))
+ >>> result.collect_column("entries")[0].as_py()
+ [{'key': 'x', 'value': 1}, {'key': 'y', 'value': 2}]
+ """
+ return Expr(f.map_entries(map.expr))
+
+
+def element_at(map: Expr, key: Expr) -> Expr:
+ """Returns the value for a given key in the map.
+
+ Returns ``[None]`` if the key is absent.
+
+ See Also:
+ This is an alias for :py:func:`map_extract`.
+ """
+ return map_extract(map, key)
+
+
 # aggregate functions
 def approx_distinct(
  expression: Expr,

python/tests/test_functions.py

@@ -668,6 +668,106 @@ def test_array_function_obj_tests(stmt, py_expr):
  assert a == b
 
 
+@pytest.mark.parametrize(
+ ("args", "expected"),
+ [
+ pytest.param(
+ ({"x": 1, "y": 2},),
+ [("x", 1), ("y", 2)],
+ id="dict",
+ ),
+ pytest.param(
+ ({"x": literal(1), "y": literal(2)},),
+ [("x", 1), ("y", 2)],
+ id="dict_with_exprs",
+ ),
+ pytest.param(
+ ("x", 1, "y", 2),
+ [("x", 1), ("y", 2)],
+ id="variadic_pairs",
+ ),
+ pytest.param(
+ (literal("x"), literal(1), literal("y"), literal(2)),
+ [("x", 1), ("y", 2)],
+ id="variadic_with_exprs",
+ ),
+ ],
+)
+def test_make_map(args, expected):
+ ctx = SessionContext()
+ batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"])
+ df = ctx.create_dataframe([[batch]])
+
+ result = df.select(f.make_map(*args).alias("m")).collect()[0].column(0)
+ assert result[0].as_py() == expected
+
+
+def test_make_map_from_two_lists():
+ 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)
+ assert result.to_pylist() == [["k1"], ["k2"], ["k3"]]
+
+ result = df.select(f.map_values(m).alias("v")).collect()[0].column(0)
+ assert result.to_pylist() == [[10], [20], [30]]
+
+
+def test_make_map_odd_args_raises():
+ with pytest.raises(ValueError, match="make_map expects"):
+ f.make_map("x", 1, "y")
+
+
+def test_make_map_mismatched_lengths():
+ with pytest.raises(ValueError, match="same length"):
+ f.make_map(["a", "b"], [1])
+
+
+@pytest.mark.parametrize(
+ ("func", "expected"),
+ [
+ pytest.param(f.map_keys, ["x", "y"], id="map_keys"),
+ pytest.param(f.map_values, [1, 2], id="map_values"),
+ pytest.param(
+ lambda m: f.map_extract(m, literal("x")),
+ [1],
+ id="map_extract",
+ ),
+ pytest.param(
+ lambda m: f.map_extract(m, literal("z")),
+ [None],
+ id="map_extract_missing_key",
+ ),
+ pytest.param(
+ f.map_entries,
+ [{"key": "x", "value": 1}, {"key": "y", "value": 2}],
+ id="map_entries",
+ ),
+ pytest.param(
+ lambda m: f.element_at(m, literal("y")),
+ [2],
+ id="element_at",
+ ),
+ ],
+)
+def test_map_functions(func, expected):
+ 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})
+ result = df.select(func(m).alias("out")).collect()[0].column(0)
+ assert result[0].as_py() == expected
+
+
 @pytest.mark.parametrize(
  ("function", "expected_result"),
  [