Skip to content
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 152 additions & 1 deletion python/pyarrow/array.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -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):

Copy link
Copy Markdown
Member

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?

Copy link
Copy Markdown
Member Author

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.

# 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):
"""
Expand Down Expand Up @@ -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()
Expand All @@ -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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The 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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reordered to the regular int8..uint64, float, double sequence in 3d303ce.

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):
"""
Expand Down Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You may be able to call GetView(i) which will give you a std::string_view.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice — done. Switched string/binary (+ large variants) to GetView(i); it also made StringViewArray/BinaryViewArray specializations trivial, so those are included now too.

return cp.PyUnicode_DecodeUTF8(<const char*> data, length, NULL)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PyArrow only supports 64-bit platforms, where Py_ssize_t and int64_t have the same width; and on a hypothetical 32-bit build a single value longer than PY_SSIZE_T_MAX could not exist in the first place, since the child data buffer itself is bounded by the process address space. So the implicit conversion cannot truncate in practice.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The 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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The 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,
Expand Down Expand Up @@ -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):
"""
Expand All @@ -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):
"""
Expand Down Expand Up @@ -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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not raise ValueError here instead of adding this weird fallback path?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done — it now raises the same ValueError as StructScalar.as_py directly (the cache/fallback dance is gone), and the test asserts the message.

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.
Expand Down
3 changes: 3 additions & 0 deletions python/pyarrow/lib.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The 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 type/_name offsets stay stable and the new field follows the append-only convention that Cython's size check assumes.


cdef readonly:
DataType type
Expand All @@ -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 *

Expand Down
32 changes: 32 additions & 0 deletions python/pyarrow/tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The 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 ValueError.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we do pa.binary_view as well?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added binary_view and string_view cases — and since GetView covers the view types, they now take real fast paths instead of the Scalar fallback.

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))

Expand Down
Loading