Skip to content

Commit

Permalink
Fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
icexelloss committed Jul 17, 2023
1 parent 1367519 commit 408dc7e
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 4 deletions.
2 changes: 1 addition & 1 deletion python/pyarrow/_compute.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -2854,7 +2854,7 @@ def register_vector_function(func, function_name, function_doc, in_types, out_ty
>>> func_name = "pct_rank_func"
>>> in_types = {"array": pa.int64()}
>>> out_type = pa.float64()
>>> pc.register_vector_function(add_constant, func_name, func_doc,
>>> pc.register_vector_function(rank_pct, func_name, func_doc,
... in_types, out_type)
>>>
>>> answer = pc.call_function(func_name, [pa.array([20, 30, 40])])
Expand Down
10 changes: 7 additions & 3 deletions python/pyarrow/tests/test_udf.py
Original file line number Diff line number Diff line change
Expand Up @@ -305,7 +305,8 @@ def unary_vector_func_fixture():
Reigster a vector function
"""
def pct_rank(ctx, x):
return pa.array(x.to_pandas().rank(pct=True))
# copy here to get around pandas 1.0 issue
return pa.array(x.to_pandas().copy().rank(pct=True))

func_name = "y=pct_rank(x)"
doc = empty_udf_doc
Expand Down Expand Up @@ -836,20 +837,23 @@ def test_hash_agg_random(sum_agg_func_fixture):
assert result.sort_by('id') == expected.sort_by('id')


@pytest.mark.pandas
def test_vector_basic(unary_vector_func_fixture):
arr = pa.array([10.0, 20.0, 30.0, 40.0, 50.0], pa.float64())
result = pc.call_function("y=pct_rank(x)", [arr])
expected = pa.array(arr.to_pandas().rank(pct=True))
expected = unary_vector_func_fixture(None, arr)
assert result == expected


@pytest.mark.pandas
def test_vector_empty(unary_vector_func_fixture):
arr = pa.array([1], pa.float64())
result = pc.call_function("y=pct_rank(x)", [arr])
expected = pa.array(arr.to_pandas().rank(pct=True))
expected = unary_vector_func_fixture(None, arr)
assert result == expected


@pytest.mark.pandas
def test_vector_struct(struct_vector_func_fixture):
k = pa.array(
[1, 1, 2, 2], pa.int64()
Expand Down

0 comments on commit 408dc7e

Please sign in to comment.