Skip to content
Open
Show file tree
Hide file tree
Changes from all 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
154 changes: 153 additions & 1 deletion python/pyarrow/array.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -1864,7 +1864,21 @@ 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]
# TODO(GH-50448): convert per range instead of per element to cut
# the per-element call overhead further.
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 +2458,12 @@ 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 +2478,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_INT8:
return (<CInt8Array*> self.ap).Value(i)
elif tid == _Type_INT16:
return (<CInt16Array*> self.ap).Value(i)
elif tid == _Type_INT32:
return (<CInt32Array*> self.ap).Value(i)
elif tid == _Type_INT64:
return (<CInt64Array*> self.ap).Value(i)
elif tid == _Type_UINT8:
return (<CUInt8Array*> self.ap).Value(i)
elif tid == _Type_UINT16:
return (<CUInt16Array*> self.ap).Value(i)
elif tid == _Type_UINT32:
return (<CUInt32Array*> self.ap).Value(i)
elif tid == _Type_UINT64:
return (<CUInt64Array*> self.ap).Value(i)
elif tid == _Type_FLOAT:
return (<CFloatArray*> self.ap).Value(i)
elif tid == _Type_DOUBLE:
return (<CDoubleArray*> 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 +2824,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 +3019,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 +3619,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 +3769,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 +4065,13 @@ cdef class StringArray(Array):
Concrete class for Arrow arrays of string (or utf8) data type.
"""

cdef object _getitem_py(self, int64_t i):
if self.ap.IsNull(i):
return None
cdef cpp_string_view view = (<CBinaryArray*> self.ap).GetView(i)
# Matches StringScalar.as_py, which is str(buf, 'utf8').
return cp.PyUnicode_DecodeUTF8(view.data(), <Py_ssize_t> view.size(), 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,12 @@ cdef class LargeStringArray(Array):
Concrete class for Arrow arrays of large string (or utf8) data type.
"""

cdef object _getitem_py(self, int64_t i):
if self.ap.IsNull(i):
return None
cdef cpp_string_view view = (<CLargeBinaryArray*> self.ap).GetView(i)
return cp.PyUnicode_DecodeUTF8(view.data(), <Py_ssize_t> view.size(), NULL)

@staticmethod
def from_buffers(int length, Buffer value_offsets, Buffer data,
Buffer null_bitmap=None, int null_count=-1,
Expand Down Expand Up @@ -4038,11 +4142,24 @@ cdef class StringViewArray(Array):
Concrete class for Arrow arrays of string (or utf8) view data type.
"""

cdef object _getitem_py(self, int64_t i):
if self.ap.IsNull(i):
return None
cdef cpp_string_view view = (<CBinaryViewArray*> self.ap).GetView(i)
return cp.PyUnicode_DecodeUTF8(view.data(), <Py_ssize_t> view.size(), NULL)


cdef class BinaryArray(Array):
"""
Concrete class for Arrow arrays of variable-sized binary data type.
"""

cdef object _getitem_py(self, int64_t i):
if self.ap.IsNull(i):
return None
cdef cpp_string_view view = (<CBinaryArray*> self.ap).GetView(i)
return cp.PyBytes_FromStringAndSize(view.data(), <Py_ssize_t> view.size())

@property
def total_values_length(self):
"""
Expand All @@ -4056,6 +4173,13 @@ cdef class LargeBinaryArray(Array):
"""
Concrete class for Arrow arrays of large variable-sized binary data type.
"""

cdef object _getitem_py(self, int64_t i):
if self.ap.IsNull(i):
return None
cdef cpp_string_view view = (<CLargeBinaryArray*> self.ap).GetView(i)
return cp.PyBytes_FromStringAndSize(view.data(), <Py_ssize_t> view.size())

@property
def total_values_length(self):
"""
Expand All @@ -4070,6 +4194,12 @@ cdef class BinaryViewArray(Array):
Concrete class for Arrow arrays of variable-sized binary view data type.
"""

cdef object _getitem_py(self, int64_t i):
if self.ap.IsNull(i):
return None
cdef cpp_string_view view = (<CBinaryViewArray*> self.ap).GetView(i)
return cp.PyBytes_FromStringAndSize(view.data(), <Py_ssize_t> view.size())


cdef class DictionaryArray(Array):
"""
Expand Down Expand Up @@ -4229,6 +4359,28 @@ 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):
# Matches StructScalar.as_py
raise ValueError(
"Converting to Python dictionary is not supported when "
"duplicate field names are present")
self._children_cache = (
names, [self.field(k) for k in range(num_fields)])
names = (<tuple> self._children_cache)[0]
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
8 changes: 8 additions & 0 deletions python/pyarrow/includes/libarrow.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -946,13 +946,15 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:

cdef cppclass CBinaryArray" arrow::BinaryArray"(CArray):
const uint8_t* GetValue(int i, int32_t* length)
cpp_string_view GetView(int64_t i)
shared_ptr[CBuffer] value_data()
int32_t value_offset(int64_t i)
int32_t value_length(int64_t i)
int32_t total_values_length()

cdef cppclass CLargeBinaryArray" arrow::LargeBinaryArray"(CArray):
const uint8_t* GetValue(int i, int64_t* length)
cpp_string_view GetView(int64_t i)
shared_ptr[CBuffer] value_data()
int64_t value_offset(int64_t i)
int64_t value_length(int64_t i)
Expand All @@ -977,6 +979,12 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:

c_string GetString(int i)

cdef cppclass CBinaryViewArray" arrow::BinaryViewArray"(CArray):
cpp_string_view GetView(int64_t i)

cdef cppclass CStringViewArray" arrow::StringViewArray"(CBinaryViewArray):
pass

cdef cppclass CStructArray" arrow::StructArray"(CArray):
CStructArray(shared_ptr[CDataType]& type, int64_t length,
vector[shared_ptr[CArray]]& children,
Expand Down
7 changes: 7 additions & 0 deletions python/pyarrow/lib.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -288,8 +288,15 @@ cdef class Array(_PandasConvertible):
# To allow Table to propagate metadata to pandas.Series
object _name

cdef:
# Lazily wrapped child array(s) reused by _getitem_py (see GH-50326).
# Appended after the pre-existing attributes to keep their offsets
# stable for extensions compiled against an older pyarrow.
object _children_cache

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
50 changes: 50 additions & 0 deletions python/pyarrow/tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -465,6 +465,56 @@ def test_array_getitem_numpy_scalars():
assert arr[np.int32(idx)].as_py() == lst[idx]


def test_to_pylist_bulk_paths():
# GH-50326: to_pylist converts through scalar-free _getitem_py
# specializations; 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([b"a\x00b", None, b"", b"\xff"], type=pa.binary()),
pa.array([b"a\x00b", None, b""], type=pa.large_binary()),
pa.array([[b"x", None, b"\x00y"], None, []],
type=pa.list_(pa.binary())),
pa.array([1, None, -(2**62), 2**62], type=pa.int64()),
pa.array([0, None, 2**63 + 7], type=pa.uint64()),
pa.array([-128, 127, None], type=pa.int8()),
pa.array([1.5, None, -0.5], type=pa.float64()),
pa.array([1.5, None], type=pa.float32()),
pa.array([True, None, False], type=pa.bool_()),
pa.array([{"a": 1, "b": "x"}, None, {"a": None, "b": None}],
type=pa.struct([("a", pa.int32()), ("b", pa.string())])),
pa.array([], type=pa.list_(pa.int32())),
pa.array([None, None], type=pa.list_(pa.string())),
]

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]

# Duplicate struct field names raise like StructScalar.as_py does
dup = pa.StructArray.from_arrays(
[pa.array([1, 2]), pa.array(["a", "b"])], names=["x", "x"])
with pytest.raises(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 you test (part of) the error message as well? Something like:

Suggested change
with pytest.raises(ValueError):
with pytest.raises(ValueError, match='some regex'):

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.

dup.to_pylist()


def test_array_slice():
arr = pa.array(range(10))

Expand Down
Loading