This is an automated email from the ASF dual-hosted git repository.
apitrou pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new 82f8826deb ARROW-16174: [Python] Fix FixedSizeListArray.flatten() on
sliced input (#14000)
82f8826deb is described below
commit 82f8826deb8dec86f5b5fde6b56cc212798f0c41
Author: Miles Granger <[email protected]>
AuthorDate: Wed Sep 21 16:37:49 2022 +0200
ARROW-16174: [Python] Fix FixedSizeListArray.flatten() on sliced input
(#14000)
[ARROW-16174](https://issues.apache.org/jira/browse/ARROW-16174)
Current behavior
```python
import pyarrow as pa
array = pa.array([[1], [2], [3]], type=pa.list_(pa.int64(), list_size=1))
array[2:].flatten().to_pylist()
[1, 2, 3]
```
After this patch
```python
import pyarrow as pa
array = pa.array([[1], [2], [3]], type=pa.list_(pa.int64(), list_size=1))
array[2:].flatten().to_pylist()
[3]
```
Authored-by: Miles Granger <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
python/pyarrow/array.pxi | 12 +------
python/pyarrow/lib.pxd | 2 +-
python/pyarrow/tests/test_array.py | 64 +++++++++++++++++++++++++-------------
3 files changed, 45 insertions(+), 33 deletions(-)
diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
index 4b35697874..40b7bc22c8 100644
--- a/python/pyarrow/array.pxi
+++ b/python/pyarrow/array.pxi
@@ -2125,7 +2125,7 @@ cdef class MapArray(ListArray):
return pyarrow_wrap_array((<CMapArray*> self.ap).items())
-cdef class FixedSizeListArray(Array):
+cdef class FixedSizeListArray(BaseListArray):
"""
Concrete class for Arrow arrays of a fixed size list data type.
"""
@@ -2212,16 +2212,6 @@ cdef class FixedSizeListArray(Array):
@property
def values(self):
- return self.flatten()
-
- def flatten(self):
- """
- Unnest this FixedSizeListArray by one level.
-
- Returns
- -------
- result : Array
- """
cdef CFixedSizeListArray* arr = <CFixedSizeListArray*> self.ap
return pyarrow_wrap_array(arr.values())
diff --git a/python/pyarrow/lib.pxd b/python/pyarrow/lib.pxd
index 67db3d2ffb..25725a4957 100644
--- a/python/pyarrow/lib.pxd
+++ b/python/pyarrow/lib.pxd
@@ -383,7 +383,7 @@ cdef class MapArray(ListArray):
pass
-cdef class FixedSizeListArray(Array):
+cdef class FixedSizeListArray(BaseListArray):
pass
diff --git a/python/pyarrow/tests/test_array.py
b/python/pyarrow/tests/test_array.py
index 39049a7859..154ca348d2 100644
--- a/python/pyarrow/tests/test_array.py
+++ b/python/pyarrow/tests/test_array.py
@@ -2685,30 +2685,49 @@ def test_list_array_flatten(offset_type,
list_type_factory):
assert arr2.values.values.equals(arr0)
[email protected]('list_type_factory', [pa.list_, pa.large_list])
-def test_list_value_parent_indices(list_type_factory):
[email protected]('list_type', [
+ pa.list_(pa.int32()),
+ pa.list_(pa.int32(), list_size=2),
+ pa.large_list(pa.int32())])
+def test_list_value_parent_indices(list_type):
arr = pa.array(
[
- [0, 1, 2],
+ [0, 1],
None,
- [],
+ [None, None],
[3, 4]
- ], type=list_type_factory(pa.int32()))
- expected = pa.array([0, 0, 0, 3, 3], type=pa.int64())
+ ], type=list_type)
+ expected = pa.array([0, 0, 2, 2, 3, 3], type=pa.int64())
assert arr.value_parent_indices().equals(expected)
[email protected](('offset_type', 'list_type_factory'),
- [(pa.int32(), pa.list_), (pa.int64(), pa.large_list)])
-def test_list_value_lengths(offset_type, list_type_factory):
- arr = pa.array(
- [
- [0, 1, 2],
- None,
- [],
- [3, 4]
- ], type=list_type_factory(pa.int32()))
- expected = pa.array([3, None, 0, 2], type=offset_type)
[email protected](('offset_type', 'list_type'),
+ [(pa.int32(), pa.list_(pa.int32())),
+ (pa.int32(), pa.list_(pa.int32(), list_size=2)),
+ (pa.int64(), pa.large_list(pa.int32()))])
+def test_list_value_lengths(offset_type, list_type):
+
+ # FixedSizeListArray needs fixed list sizes
+ if getattr(list_type, "list_size", None):
+ arr = pa.array(
+ [
+ [0, 1],
+ None,
+ [None, None],
+ [3, 4]
+ ], type=list_type)
+ expected = pa.array([2, None, 2, 2], type=offset_type)
+
+ # Otherwise create variable list sizes
+ else:
+ arr = pa.array(
+ [
+ [0, 1, 2],
+ None,
+ [],
+ [3, 4]
+ ], type=list_type)
+ expected = pa.array([3, None, 0, 2], type=offset_type)
assert arr.value_lengths().equals(expected)
@@ -2771,7 +2790,6 @@ def test_fixed_size_list_array_flatten():
typ1 = pa.list_(pa.int64(), 2)
arr1 = pa.array([
[1, 2], [3, 4], [5, 6],
- None, None, None,
[7, None], None, [8, 9]
], type=typ1)
assert arr1.type.equals(typ1)
@@ -2779,15 +2797,19 @@ def test_fixed_size_list_array_flatten():
typ0 = pa.int64()
arr0 = pa.array([
- 1, 2, 3, 4, 5, 6,
- None, None, None, None, None, None,
- 7, None, None, None, 8, 9,
+ 1, 2, 3, 4, 5, 6, 7, None, 8, 9,
], type=typ0)
assert arr0.type.equals(typ0)
assert arr1.flatten().equals(arr0)
assert arr2.flatten().flatten().equals(arr0)
+def test_fixed_size_list_array_flatten_with_slice():
+ array = pa.array([[1], [2], [3]],
+ type=pa.list_(pa.float64(), list_size=1))
+ assert array[2:].flatten() == pa.array([3], type=pa.float64())
+
+
def test_map_array_values_offsets():
ty = pa.map_(pa.utf8(), pa.int32())
ty_values = pa.struct([pa.field("key", pa.utf8(), nullable=False),
