Skip to content

Commit

Permalink
Add make_list and tests for make_list, make_array
Browse files Browse the repository at this point in the history
  • Loading branch information
kosiew committed Nov 14, 2024
1 parent 53cdb11 commit ca1cf3a
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 0 deletions.
9 changes: 9 additions & 0 deletions python/datafusion/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,7 @@
"lpad",
"ltrim",
"make_array",
"make_list",
"make_date",
"max",
"md5",
Expand Down Expand Up @@ -1044,6 +1045,14 @@ def make_array(*args: Expr) -> Expr:
return Expr(f.make_array(args))


def make_list(*args: Expr) -> Expr:
"""Returns an array using the specified input expressions.
This is an alias for :py:func:`make_array`.
"""
return make_array(*args)


def array(*args: Expr) -> Expr:
"""Returns an array using the specified input expressions.
Expand Down
31 changes: 31 additions & 0 deletions python/tests/test_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -576,6 +576,37 @@ def test_array_function_cardinality():
)


@pytest.mark.parametrize("make_func", [f.make_array, f.make_list])
def test_make_array_functions(make_func):
ctx = SessionContext()
batch = pa.RecordBatch.from_arrays(
[
pa.array(["Hello", "World", "!"], type=pa.string()),
pa.array([4, 5, 6]),
pa.array(["hello ", " world ", " !"], type=pa.string()),
],
names=["a", "b", "c"],
)
df = ctx.create_dataframe([[batch]])

stmt = make_func(
column("a").cast(pa.string()),
column("b").cast(pa.string()),
column("c").cast(pa.string()),
)
py_expr = [
["Hello", "4", "hello "],
["World", "5", " world "],
["!", "6", " !"],
]

query_result = df.select(stmt).collect()[0].column(0)
for a, b in zip(query_result, py_expr):
np.testing.assert_array_equal(
np.array(a.as_py(), dtype=str), np.array(b, dtype=str)
)


@pytest.mark.parametrize(
("stmt", "py_expr"),
[
Expand Down

0 comments on commit ca1cf3a

Please sign in to comment.