-
Notifications
You must be signed in to change notification settings - Fork 4.2k
GH-50326: [Python] Convert arrays to Python objects without per-element Scalars in to_pylist #50327
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.
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
base: main
Are you sure you want to change the base?
Changes from 2 commits
12b921e
7a15077
985a2e3
728584f
3d303ce
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -1864,7 +1864,19 @@ cdef class Array(_PandasConvertible): | |
| lst : list | ||
| """ | ||
| self._assert_cpu() | ||
| return [x.as_py(maps_as_pydicts=maps_as_pydicts) for x in self] | ||
| cdef int64_t i, n = self.length() | ||
| if maps_as_pydicts is not None: | ||
| # Converting maps to dicts has per-entry semantics (duplicate-key | ||
| # detection); use the Scalar-based conversion for exact behavior. | ||
| return [x.as_py(maps_as_pydicts=maps_as_pydicts) for x in self] | ||
| return [self._getitem_py(i) for i in range(n)] | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| # Return self[i] as a Python object, without creating a Python Scalar | ||
| # (nor, for nested types, per-row Array wrappers) where a subclass | ||
| # provides a specialization; this base implementation goes through | ||
| # Scalar.as_py and thus preserves its semantics exactly (see GH-50326). | ||
| return self.getitem(i).as_py() | ||
|
|
||
| def tolist(self): | ||
| """ | ||
|
|
@@ -2444,6 +2456,11 @@ cdef class BooleanArray(Array): | |
| """ | ||
| Concrete class for Arrow arrays of boolean data type. | ||
| """ | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| if self.ap.IsNull(i): | ||
| return None | ||
| return (<CBooleanArray*> self.ap).Value(i) | ||
| @property | ||
| def false_count(self): | ||
| return (<CBooleanArray*> self.ap).false_count() | ||
|
|
@@ -2458,6 +2475,34 @@ cdef class NumericArray(Array): | |
| A base class for Arrow numeric arrays. | ||
| """ | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| cdef Type tid = self.ap.type_id() | ||
| if self.ap.IsNull(i): | ||
| return None | ||
| if tid == _Type_INT64: | ||
| return (<CInt64Array*> self.ap).Value(i) | ||
| elif tid == _Type_INT32: | ||
| return (<CInt32Array*> self.ap).Value(i) | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why the peculiar ordering of types? I would rather have something more regular for readability.
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Reordered to the regular |
||
| elif tid == _Type_DOUBLE: | ||
| return (<CDoubleArray*> self.ap).Value(i) | ||
| elif tid == _Type_FLOAT: | ||
| return (<CFloatArray*> self.ap).Value(i) | ||
| elif tid == _Type_INT16: | ||
| return (<CInt16Array*> self.ap).Value(i) | ||
| elif tid == _Type_INT8: | ||
| return (<CInt8Array*> self.ap).Value(i) | ||
| elif tid == _Type_UINT64: | ||
| return (<CUInt64Array*> self.ap).Value(i) | ||
| elif tid == _Type_UINT32: | ||
| return (<CUInt32Array*> self.ap).Value(i) | ||
| elif tid == _Type_UINT16: | ||
| return (<CUInt16Array*> self.ap).Value(i) | ||
| elif tid == _Type_UINT8: | ||
| return (<CUInt8Array*> self.ap).Value(i) | ||
| # Subclasses whose as_py returns non-primitive objects (dates, times, | ||
| # timestamps, durations, half floats, ...) use the exact Scalar path. | ||
| return Array._getitem_py(self, i) | ||
|
|
||
|
|
||
| cdef class IntegerArray(NumericArray): | ||
| """ | ||
|
|
@@ -2776,6 +2821,16 @@ cdef class ListArray(BaseListArray): | |
| Concrete class for Arrow arrays of a list data type. | ||
| """ | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| cdef CListArray* arr = <CListArray*> self.ap | ||
| if arr.IsNull(i): | ||
| return None | ||
| if self._children_cache is None: | ||
| self._children_cache = pyarrow_wrap_array(arr.values()) | ||
| cdef Array values = <Array> self._children_cache | ||
| cdef int64_t j, start = arr.value_offset(i), end = arr.value_offset(i + 1) | ||
| return [values._getitem_py(j) for j in range(start, end)] | ||
|
|
||
| @staticmethod | ||
| def from_arrays(offsets, values, DataType type=None, MemoryPool pool=None, mask=None): | ||
| """ | ||
|
|
@@ -2961,6 +3016,16 @@ cdef class LargeListArray(BaseListArray): | |
| Identical to ListArray, but 64-bit offsets. | ||
| """ | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| cdef CLargeListArray* arr = <CLargeListArray*> self.ap | ||
| if arr.IsNull(i): | ||
| return None | ||
| if self._children_cache is None: | ||
| self._children_cache = pyarrow_wrap_array(arr.values()) | ||
| cdef Array values = <Array> self._children_cache | ||
| cdef int64_t j, start = arr.value_offset(i), end = arr.value_offset(i + 1) | ||
| return [values._getitem_py(j) for j in range(start, end)] | ||
|
|
||
| @staticmethod | ||
| def from_arrays(offsets, values, DataType type=None, MemoryPool pool=None, mask=None): | ||
| """ | ||
|
|
@@ -3551,6 +3616,19 @@ cdef class MapArray(ListArray): | |
| Concrete class for Arrow arrays of a map data type. | ||
| """ | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| cdef CListArray* arr = <CListArray*> self.ap | ||
| if arr.IsNull(i): | ||
| return None | ||
| if self._children_cache is None: | ||
| self._children_cache = (self.keys, self.items) | ||
| cdef Array keys = <Array> (<tuple> self._children_cache)[0] | ||
| cdef Array items = <Array> (<tuple> self._children_cache)[1] | ||
| cdef int64_t j, start = arr.value_offset(i), end = arr.value_offset(i + 1) | ||
| # Matches MapScalar.as_py with the default maps_as_pydicts=None: | ||
| # an association list of (key, value) tuples. | ||
| return [(keys._getitem_py(j), items._getitem_py(j)) for j in range(start, end)] | ||
|
|
||
| @staticmethod | ||
| def from_arrays(offsets, keys, items, DataType type=None, MemoryPool pool=None, mask=None): | ||
| """ | ||
|
|
@@ -3688,6 +3766,16 @@ cdef class FixedSizeListArray(BaseListArray): | |
| Concrete class for Arrow arrays of a fixed size list data type. | ||
| """ | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| cdef CFixedSizeListArray* arr = <CFixedSizeListArray*> self.ap | ||
| if arr.IsNull(i): | ||
| return None | ||
| if self._children_cache is None: | ||
| self._children_cache = pyarrow_wrap_array(arr.values()) | ||
| cdef Array values = <Array> self._children_cache | ||
| cdef int64_t j, start = arr.value_offset(i), end = arr.value_offset(i + 1) | ||
| return [values._getitem_py(j) for j in range(start, end)] | ||
|
|
||
| @staticmethod | ||
| def from_arrays(values, list_size=None, DataType type=None, mask=None): | ||
| """ | ||
|
|
@@ -3974,6 +4062,16 @@ cdef class StringArray(Array): | |
| Concrete class for Arrow arrays of string (or utf8) data type. | ||
| """ | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| cdef: | ||
| int32_t length | ||
| const uint8_t* data | ||
| if self.ap.IsNull(i): | ||
| return None | ||
| data = (<CStringArray*> self.ap).GetValue(i, &length) | ||
| # Matches StringScalar.as_py, which is str(buf, 'utf8'). | ||
| return cp.PyUnicode_DecodeUTF8(<const char*> data, length, NULL) | ||
|
|
||
| @staticmethod | ||
| def from_buffers(int length, Buffer value_offsets, Buffer data, | ||
| Buffer null_bitmap=None, int null_count=-1, | ||
|
|
@@ -4006,6 +4104,15 @@ cdef class LargeStringArray(Array): | |
| Concrete class for Arrow arrays of large string (or utf8) data type. | ||
| """ | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| cdef: | ||
| int64_t length | ||
| const uint8_t* data | ||
| if self.ap.IsNull(i): | ||
| return None | ||
| data = (<CLargeStringArray*> self.ap).GetValue(i, &length) | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You may be able to call
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nice — done. Switched string/binary (+ large variants) to |
||
| return cp.PyUnicode_DecodeUTF8(<const char*> data, length, NULL) | ||
|
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. PyArrow only supports 64-bit platforms, where
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We do support 32-bit platforms actually. But obviously a string size couldn't be larger than 4GiB on such a platform anyway, so the AI's comment is moot.
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks for the correction — I stand corrected on 32-bit support; agreed it's moot either way since such a value can't exist in a 32-bit address space. |
||
|
|
||
| @staticmethod | ||
| def from_buffers(int length, Buffer value_offsets, Buffer data, | ||
| Buffer null_bitmap=None, int null_count=-1, | ||
|
|
@@ -4043,6 +4150,16 @@ cdef class BinaryArray(Array): | |
| """ | ||
| Concrete class for Arrow arrays of variable-sized binary data type. | ||
| """ | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| cdef: | ||
| int32_t length | ||
| const uint8_t* data | ||
| if self.ap.IsNull(i): | ||
| return None | ||
| data = (<CBinaryArray*> self.ap).GetValue(i, &length) | ||
| return cp.PyBytes_FromStringAndSize(<const char*> data, length) | ||
|
|
||
| @property | ||
| def total_values_length(self): | ||
| """ | ||
|
|
@@ -4056,6 +4173,16 @@ cdef class LargeBinaryArray(Array): | |
| """ | ||
| Concrete class for Arrow arrays of large variable-sized binary data type. | ||
| """ | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| cdef: | ||
| int64_t length | ||
| const uint8_t* data | ||
| if self.ap.IsNull(i): | ||
| return None | ||
| data = (<CLargeBinaryArray*> self.ap).GetValue(i, &length) | ||
| return cp.PyBytes_FromStringAndSize(<const char*> data, length) | ||
|
|
||
| @property | ||
| def total_values_length(self): | ||
| """ | ||
|
|
@@ -4229,6 +4356,30 @@ cdef class StructArray(Array): | |
| Concrete class for Arrow arrays of a struct data type. | ||
| """ | ||
|
|
||
| cdef object _getitem_py(self, int64_t i): | ||
| if self.ap.IsNull(i): | ||
| return None | ||
| cdef int64_t k, num_fields = self.type.num_fields | ||
| if self._children_cache is None: | ||
| names = [self.type.field(k).name for k in range(num_fields)] | ||
| if len(set(names)) != len(names): | ||
| # StructScalar.as_py raises ValueError on duplicate field | ||
| # names; mark the cache so we take the Scalar path below. | ||
| self._children_cache = (None, None) | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why not raise
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done — it now raises the same |
||
| else: | ||
| self._children_cache = ( | ||
| names, [self.field(k) for k in range(num_fields)]) | ||
| names = (<tuple> self._children_cache)[0] | ||
| if names is None: | ||
| return Array._getitem_py(self, i) | ||
| fields = (<tuple> self._children_cache)[1] | ||
| cdef Array field_arr | ||
| result = {} | ||
| for k in range(num_fields): | ||
| field_arr = <Array> fields[k] | ||
| result[names[k]] = field_arr._getitem_py(i) | ||
| return result | ||
|
|
||
| def field(self, index): | ||
| """ | ||
| Retrieves the child array belonging to field. | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -282,6 +282,8 @@ cdef class Array(_PandasConvertible): | |
| cdef: | ||
| shared_ptr[CArray] sp_array | ||
| CArray* ap | ||
| # Lazily wrapped child array(s) reused by _getitem_py (see GH-50326) | ||
| object _children_cache | ||
|
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Good point — moved the declaration after the pre-existing attributes in 728584f, so |
||
|
|
||
| cdef readonly: | ||
| DataType type | ||
|
|
@@ -290,6 +292,7 @@ cdef class Array(_PandasConvertible): | |
|
|
||
| cdef void init(self, const shared_ptr[CArray]& sp_array) except * | ||
| cdef getitem(self, int64_t i) | ||
| cdef object _getitem_py(self, int64_t i) | ||
| cdef int64_t length(self) | ||
| cdef void _assert_cpu(self) except * | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -465,6 +465,38 @@ def test_array_getitem_numpy_scalars(): | |
| assert arr[np.int32(idx)].as_py() == lst[idx] | ||
|
|
||
|
|
||
| def test_to_pylist_bulk_paths(): | ||
| # GH-50326: list-like and string arrays convert to Python objects in | ||
| # bulk instead of going through one Scalar per element; the result must | ||
| # match the per-scalar conversion exactly. | ||
| arrays = [ | ||
| pa.array([[1, None, 3], None, [], [4]], type=pa.list_(pa.int32())), | ||
| pa.array([["a", None], None, [], ["bcd", ""]], | ||
| type=pa.list_(pa.string())), | ||
| pa.array([["a", None], None, [], ["bcd", ""]], | ||
| type=pa.large_list(pa.large_string())), | ||
| pa.array([[1, None], None, [3, 4]], type=pa.list_(pa.int32(), 2)), | ||
| pa.array([[[1], [2, None]], None, [None, [3]]], | ||
| type=pa.list_(pa.list_(pa.int32()))), | ||
| pa.array([[("k1", 1), ("k2", None)], None, []], | ||
| type=pa.map_(pa.string(), pa.int32())), | ||
| pa.array(["a", None, "", "\N{GRINNING FACE} \N{SNOWMAN}"], | ||
| type=pa.string()), | ||
| pa.array(["a", None, "", "\N{GRINNING FACE} \N{SNOWMAN}"], | ||
| type=pa.large_string()), | ||
| pa.array([], type=pa.list_(pa.int32())), | ||
| pa.array([None, None], type=pa.list_(pa.string())), | ||
| ] | ||
|
Comment on lines
+472
to
+501
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Added binary/large_binary (including embedded NUL bytes), list, wide-range integers, floats, boolean and struct coverage in 728584f, plus an assertion that duplicate struct field names still raise
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can we do
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Added |
||
| for arr in arrays: | ||
| for view in (arr, arr.slice(1), arr.slice(0, 2), arr.slice(2)): | ||
| assert view.to_pylist() == [x.as_py() for x in view] | ||
|
|
||
| # Values inside numeric lists must stay Python ints/None, never floats | ||
| result = pa.array([[1, None, 3]], type=pa.list_(pa.int32())).to_pylist() | ||
| assert result == [[1, None, 3]] | ||
| assert [type(x) for x in result[0]] == [int, type(None), int] | ||
|
|
||
|
|
||
| def test_array_slice(): | ||
| arr = pa.array(range(10)) | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not necessary for this PR, but as the AI hinted it might be better to replace this with a
cdef list _getitem_range_py(self, int64_t offset, int64_t length). This would cut down on function call and prologue overhead.Perhaps add a TODO or open a separate issue?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Filed GH-50448 for the per-range conversion (folding in the dispatch-hoisting and null-check ideas from this thread as well), and added a TODO pointing at it.