From f7bd12ba5d70103ad6e4145d1ae10b3c8d6decdf Mon Sep 17 00:00:00 2001 From: Benjamin Kietzman Date: Mon, 24 Jun 2019 07:16:56 -0400 Subject: [PATCH] rewrite python Take() test --- python/pyarrow/tests/test_compute.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index 655dd38d1f02e..e356b9d9c9a6a 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -51,7 +51,10 @@ def test_sum(arrow_type): ('double', np.arange(0, 0.5, 0.1)), ('string', ['a', 'b', None, 'ddd', 'ee']), ('binary', [b'a', b'b', b'c', b'ddd', b'ee']), - (pa.binary(3), [b'abc', b'bcd', b'cde', b'def', b'efg']) + (pa.binary(3), [b'abc', b'bcd', b'cde', b'def', b'efg']), + (pa.list_(pa.int8()), [[1, 2], [3, 4], [5, 6], None, [9, 16]]), + (pa.struct([('a', pa.int8()), ('b', pa.int8())]), [ + {'a': 1, 'b': 2}, None, {'a': 3, 'b': 4}, None, {'a': 5, 'b': 6}]), ]) def test_take(ty, values): arr = pa.array(values, type=ty) @@ -100,14 +103,3 @@ def test_take_dictionary(ordered): assert result.to_pylist() == ['a', 'b', 'a'] assert result.dictionary.to_pylist() == ['a', 'b', 'c'] assert result.type.ordered is ordered - - -@pytest.mark.parametrize('array', [ - [[1, 2], [3, 4], [5, 6]], - [{'a': 1, 'b': 2}, None, {'a': 3, 'b': 4}], -], ids=['listarray', 'structarray']) -def test_take_notimplemented(array): - array = pa.array(array) - indices = pa.array([0, 2]) - with pytest.raises(NotImplementedError): - array.take(indices)