Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add missing array functions #551

Merged
merged 7 commits into from
Dec 28, 2023
Merged

Add missing array functions #551

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
a487d83
Add array_append, array_concat and array_cat
ongchi Dec 21, 2023
7f6afcf
Add tests for array functions array_append, array_concat and array_cat
ongchi Dec 21, 2023
102d307
Add array_dims and list_dims
ongchi Dec 21, 2023
c4d2734
Add tests for array_dims and list_dims
ongchi Dec 21, 2023
9f1580b
Add array_element, array_extract, list_element and list_extract
ongchi Dec 21, 2023
d06303f
Add tests for array_element, array_extract, list_element and list_ext…
ongchi Dec 21, 2023
48566e5
Add array_length and list_length
ongchi Dec 25, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 64 additions & 0 deletions datafusion/tests/test_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
from datafusion import functions as f
from datafusion import literal

np.seterr(invalid="ignore")
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's reasonable to turn off warnings of invalid floating-point errors, which can safely be ignored in test cases.



@pytest.fixture
def df():
Expand Down Expand Up @@ -197,6 +199,68 @@ def test_math_functions():
)


def test_array_functions():
data = [[1.0, 2.0, 3.0], [4.0, 5.0], [6.0]]
ctx = SessionContext()
batch = pa.RecordBatch.from_arrays(
[np.array(data, dtype=object)], names=["arr"]
)
df = ctx.create_dataframe([[batch]])

col = column("arr")
test_items = [
[
f.array_append(col, literal(99.0)),
lambda: [np.append(arr, 99.0) for arr in data],
],
[
f.array_concat(col, col),
lambda: [np.concatenate([arr, arr]) for arr in data],
],
[
f.array_cat(col, col),
lambda: [np.concatenate([arr, arr]) for arr in data],
],
[
f.array_dims(col),
lambda: [[len(r)] for r in data],
],
[
f.list_dims(col),
lambda: [[len(r)] for r in data],
],
[
f.array_element(col, literal(1)),
lambda: [r[0] for r in data],
],
[
f.array_extract(col, literal(1)),
lambda: [r[0] for r in data],
],
[
f.list_element(col, literal(1)),
lambda: [r[0] for r in data],
],
[
f.list_extract(col, literal(1)),
lambda: [r[0] for r in data],
],
[
f.array_length(col),
lambda: [len(r) for r in data],
],
[
f.list_length(col),
lambda: [len(r) for r in data],
],
]

for stmt, py_expr in test_items:
query_result = df.select(stmt).collect()[0].column(0).tolist()
for a, b in zip(query_result, py_expr()):
np.testing.assert_array_almost_equal(a, b)


def test_string_functions(df):
df = df.select(
f.ascii(column("a")),
Expand Down
27 changes: 27 additions & 0 deletions src/functions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -350,6 +350,19 @@ scalar_function!(random, Random);
scalar_function!(encode, Encode);
scalar_function!(decode, Decode);

// Array Functions
scalar_function!(array_append, ArrayAppend);
scalar_function!(array_concat, ArrayConcat);
scalar_function!(array_cat, ArrayConcat);
scalar_function!(array_dims, ArrayDims);
scalar_function!(list_dims, ArrayDims);
scalar_function!(array_element, ArrayElement);
scalar_function!(array_extract, ArrayElement);
scalar_function!(list_element, ArrayElement);
scalar_function!(list_extract, ArrayElement);
scalar_function!(array_length, ArrayLength);
scalar_function!(list_length, ArrayLength);

aggregate_function!(approx_distinct, ApproxDistinct);
aggregate_function!(approx_median, ApproxMedian);
aggregate_function!(approx_percentile_cont, ApproxPercentileCont);
Expand Down Expand Up @@ -539,5 +552,19 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
//Binary String Functions
m.add_wrapped(wrap_pyfunction!(encode))?;
m.add_wrapped(wrap_pyfunction!(decode))?;

// Array Functions
m.add_wrapped(wrap_pyfunction!(array_append))?;
m.add_wrapped(wrap_pyfunction!(array_concat))?;
m.add_wrapped(wrap_pyfunction!(array_cat))?;
m.add_wrapped(wrap_pyfunction!(array_dims))?;
m.add_wrapped(wrap_pyfunction!(list_dims))?;
m.add_wrapped(wrap_pyfunction!(array_element))?;
m.add_wrapped(wrap_pyfunction!(array_extract))?;
m.add_wrapped(wrap_pyfunction!(list_element))?;
m.add_wrapped(wrap_pyfunction!(list_extract))?;
m.add_wrapped(wrap_pyfunction!(array_length))?;
m.add_wrapped(wrap_pyfunction!(list_length))?;

Ok(())
}
Loading