From f141daec9b102275304ce4c5490ca4a7d16ce480 Mon Sep 17 00:00:00 2001 From: Eddy Oyieko Date: Tue, 16 Jan 2024 11:10:46 +0000 Subject: [PATCH 1/4] Updated functions.rs, test_functions.py - Flatten --- datafusion/tests/test_functions.py | 13 +++++++++++++ src/functions.rs | 2 ++ 2 files changed, 15 insertions(+) diff --git a/datafusion/tests/test_functions.py b/datafusion/tests/test_functions.py index d0514f892..4be847722 100644 --- a/datafusion/tests/test_functions.py +++ b/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,10 @@ def test_array_functions(): f.list_length(col), lambda: [len(r) for r in data], ], + [ + f.flatten(col), + lambda: [py_flatten(data)] + ], ] for stmt, py_expr in test_items: diff --git a/src/functions.rs b/src/functions.rs index 3dc5322aa..44573d076 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -249,6 +249,7 @@ scalar_function!(degrees, Degrees); scalar_function!(exp, Exp); scalar_function!(factorial, Factorial); scalar_function!(floor, Floor); +scalar_function!(flatten, Flatten); scalar_function!(gcd, Gcd); scalar_function!(initcap, InitCap, "Converts the first letter of each word to upper case and the rest to lower case. Words are sequences of alphanumeric characters separated by non-alphanumeric characters."); scalar_function!(isnan, Isnan); @@ -464,6 +465,7 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(exp))?; m.add_wrapped(wrap_pyfunction!(factorial))?; m.add_wrapped(wrap_pyfunction!(floor))?; + m.add_wrapped(wrap_pyfunction!(flatten))?; m.add_wrapped(wrap_pyfunction!(from_unixtime))?; m.add_wrapped(wrap_pyfunction!(gcd))?; m.add_wrapped(wrap_pyfunction!(grouping))?; From 514a5ded932db440308dd7939af83e5aa2e5c39b Mon Sep 17 00:00:00 2001 From: Eddy Oyieko Date: Tue, 16 Jan 2024 13:35:07 +0000 Subject: [PATCH 2/4] Updated functions.rs - Formatting --- src/functions.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/functions.rs b/src/functions.rs index 44573d076..def9ff57f 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -249,7 +249,6 @@ scalar_function!(degrees, Degrees); scalar_function!(exp, Exp); scalar_function!(factorial, Factorial); scalar_function!(floor, Floor); -scalar_function!(flatten, Flatten); scalar_function!(gcd, Gcd); scalar_function!(initcap, InitCap, "Converts the first letter of each word to upper case and the rest to lower case. Words are sequences of alphanumeric characters separated by non-alphanumeric characters."); scalar_function!(isnan, Isnan); @@ -370,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); @@ -465,7 +465,6 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(exp))?; m.add_wrapped(wrap_pyfunction!(factorial))?; m.add_wrapped(wrap_pyfunction!(floor))?; - m.add_wrapped(wrap_pyfunction!(flatten))?; m.add_wrapped(wrap_pyfunction!(from_unixtime))?; m.add_wrapped(wrap_pyfunction!(gcd))?; m.add_wrapped(wrap_pyfunction!(grouping))?; @@ -574,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(()) } From 9fec85b8193e960fe23047a8007920dd8a4b81fc Mon Sep 17 00:00:00 2001 From: Eddy Oyieko <67474838+mobley-trent@users.noreply.github.com> Date: Mon, 12 Feb 2024 14:52:38 +0000 Subject: [PATCH 3/4] Updated test_functions.py - Converted test array to a literal for testing --- datafusion/tests/test_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/tests/test_functions.py b/datafusion/tests/test_functions.py index 4be847722..1cd9188e4 100644 --- a/datafusion/tests/test_functions.py +++ b/datafusion/tests/test_functions.py @@ -263,7 +263,7 @@ def py_flatten(arr): lambda: [len(r) for r in data], ], [ - f.flatten(col), + f.flatten(literal(data)), lambda: [py_flatten(data)] ], ] From 872110320a9559664dc394530c521e5926a1a703 Mon Sep 17 00:00:00 2001 From: Eddy Oyieko <67474838+mobley-trent@users.noreply.github.com> Date: Tue, 13 Feb 2024 07:47:19 +0000 Subject: [PATCH 4/4] Updated test_functions.py - Linting --- datafusion/tests/test_functions.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/datafusion/tests/test_functions.py b/datafusion/tests/test_functions.py index 1cd9188e4..f97862a2b 100644 --- a/datafusion/tests/test_functions.py +++ b/datafusion/tests/test_functions.py @@ -262,10 +262,7 @@ def py_flatten(arr): f.list_length(col), lambda: [len(r) for r in data], ], - [ - f.flatten(literal(data)), - lambda: [py_flatten(data)] - ], + [f.flatten(literal(data)), lambda: [py_flatten(data)]], ] for stmt, py_expr in test_items: