feat: Add flatten array function

datafusion/tests/test_functions.py

@@ -207,6 +207,15 @@ def test_array_functions():
     )
     df = ctx.create_dataframe([[batch]])
 
+    def py_flatten(arr):
+        result = []
+        for elem in arr:
+            if isinstance(elem, list):
+                result.extend(py_flatten(elem))
+            else:
+                result.append(elem)
+        return result
+
     col = column("arr")
     test_items = [
         [
@@ -253,6 +262,7 @@ def test_array_functions():
             f.list_length(col),
             lambda: [len(r) for r in data],
         ],
+        [f.flatten(literal(data)), lambda: [py_flatten(data)]],
     ]
 
     for stmt, py_expr in test_items:

src/functions.rs

@@ -369,6 +369,7 @@ scalar_function!(list_element, ArrayElement);
 scalar_function!(list_extract, ArrayElement);
 scalar_function!(array_length, ArrayLength);
 scalar_function!(list_length, ArrayLength);
+scalar_function!(flatten, Flatten);
 
 aggregate_function!(approx_distinct, ApproxDistinct);
 aggregate_function!(approx_median, ApproxMedian);
@@ -572,6 +573,7 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(list_extract))?;
     m.add_wrapped(wrap_pyfunction!(array_length))?;
     m.add_wrapped(wrap_pyfunction!(list_length))?;
+    m.add_wrapped(wrap_pyfunction!(flatten))?;
 
     Ok(())
 }