feat: expose named_struct in python (apache#700)

Ref apache#692
mesejo · May 15, 2024 · d6c42b4 · d6c42b4
1 parent 344ebad
commit d6c42b4
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/python/datafusion/tests/test_functions.py b/python/datafusion/tests/test_functions.py
@@ -50,6 +50,32 @@ def df():
     return ctx.create_dataframe([[batch]])
 
 
+def test_named_struct(df):
+    df = df.with_column(
+        "d",
+        f.named_struct(
+            literal("a"),
+            column("a"),
+            literal("b"),
+            column("b"),
+            literal("c"),
+            column("c"),
+        ),
+    )
+
+    expected = """DataFrame()
++-------+---+---------+------------------------------+
+| a     | b | c       | d                            |
++-------+---+---------+------------------------------+
+| Hello | 4 | hello   | {a: Hello, b: 4, c: hello }  |
+| World | 5 |  world  | {a: World, b: 5, c:  world } |
+| !     | 6 |  !      | {a: !, b: 6, c:  !}          |
++-------+---+---------+------------------------------+
+""".strip()
+
+    assert str(df) == expected
+
+
 def test_literal(df):
     df = df.select(
         literal(1),

diff --git a/src/functions.rs b/src/functions.rs
@@ -503,6 +503,7 @@ expr_fn_vec!(trunc);
 expr_fn!(upper, arg1, "Converts the string to all upper case.");
 expr_fn!(uuid);
 expr_fn_vec!(r#struct); // Use raw identifier since struct is a keyword
+expr_fn_vec!(named_struct);
 expr_fn!(from_unixtime, unixtime);
 expr_fn!(arrow_typeof, arg_1);
 expr_fn!(random);
@@ -680,6 +681,7 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(mean))?;
     m.add_wrapped(wrap_pyfunction!(median))?;
     m.add_wrapped(wrap_pyfunction!(min))?;
+    m.add_wrapped(wrap_pyfunction!(named_struct))?;
     m.add_wrapped(wrap_pyfunction!(nanvl))?;
     m.add_wrapped(wrap_pyfunction!(now))?;
     m.add_wrapped(wrap_pyfunction!(nullif))?;