diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index 5fc74969abfd..fb004243a209 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -1864,7 +1864,17 @@ 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() + # TODO(GH-50448): convert per range instead of per element to cut + # the per-element call overhead further. + return [self._getitem_py(i, maps_as_pydicts) for i in range(n)] + + cdef object _getitem_py(self, int64_t i, object maps_as_pydicts): + # 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(maps_as_pydicts=maps_as_pydicts) def tolist(self): """ @@ -2444,6 +2454,12 @@ cdef class BooleanArray(Array): """ Concrete class for Arrow arrays of boolean data type. """ + + cdef object _getitem_py(self, int64_t i, object maps_as_pydicts): + if self.ap.IsNull(i): + return None + return ( self.ap).Value(i) + @property def false_count(self): return ( self.ap).false_count() @@ -2458,6 +2474,34 @@ cdef class NumericArray(Array): A base class for Arrow numeric arrays. """ + cdef object _getitem_py(self, int64_t i, object maps_as_pydicts): + cdef Type tid = self.ap.type_id() + if self.ap.IsNull(i): + return None + if tid == _Type_INT8: + return ( self.ap).Value(i) + elif tid == _Type_INT16: + return ( self.ap).Value(i) + elif tid == _Type_INT32: + return ( self.ap).Value(i) + elif tid == _Type_INT64: + return ( self.ap).Value(i) + elif tid == _Type_UINT8: + return ( self.ap).Value(i) + elif tid == _Type_UINT16: + return ( self.ap).Value(i) + elif tid == _Type_UINT32: + return ( self.ap).Value(i) + elif tid == _Type_UINT64: + return ( self.ap).Value(i) + elif tid == _Type_FLOAT: + return ( self.ap).Value(i) + elif tid == _Type_DOUBLE: + return ( 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, maps_as_pydicts) + cdef class IntegerArray(NumericArray): """ @@ -2776,6 +2820,16 @@ cdef class ListArray(BaseListArray): Concrete class for Arrow arrays of a list data type. """ + cdef object _getitem_py(self, int64_t i, object maps_as_pydicts): + cdef CListArray* arr = 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 = self._children_cache + cdef int64_t j, start = arr.value_offset(i), end = arr.value_offset(i + 1) + return [values._getitem_py(j, maps_as_pydicts) for j in range(start, end)] + @staticmethod def from_arrays(offsets, values, DataType type=None, MemoryPool pool=None, mask=None): """ @@ -2961,6 +3015,16 @@ cdef class LargeListArray(BaseListArray): Identical to ListArray, but 64-bit offsets. """ + cdef object _getitem_py(self, int64_t i, object maps_as_pydicts): + cdef CLargeListArray* arr = 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 = self._children_cache + cdef int64_t j, start = arr.value_offset(i), end = arr.value_offset(i + 1) + return [values._getitem_py(j, maps_as_pydicts) for j in range(start, end)] + @staticmethod def from_arrays(offsets, values, DataType type=None, MemoryPool pool=None, mask=None): """ @@ -3551,6 +3615,57 @@ cdef class MapArray(ListArray): Concrete class for Arrow arrays of a map data type. """ + cdef object _getitem_py(self, int64_t i, object maps_as_pydicts): + cdef CListArray* arr = self.ap + cdef bint as_dicts = maps_as_pydicts is not None + if as_dicts and maps_as_pydicts != "lossy" and maps_as_pydicts != "strict": + # Matches MapScalar.as_py, which validates before the null check. + raise ValueError( + "Invalid value for 'maps_as_pydicts': " + + "valid values are 'lossy', 'strict' or `None` (default). " + + f"Received {maps_as_pydicts!r}." + ) + if arr.IsNull(i): + return None + if self._children_cache is None: + self._children_cache = (self.keys, self.items) + cdef Array keys = ( self._children_cache)[0] + cdef Array items = ( self._children_cache)[1] + cdef int64_t j, start = arr.value_offset(i), end = arr.value_offset(i + 1) + if not as_dicts: + # Matches MapScalar.as_py with the default maps_as_pydicts=None: + # an association list of (key, value) tuples. + return [ + (keys._getitem_py(j, None), items._getitem_py(j, maps_as_pydicts)) + for j in range(start, end) + ] + # Convert all keys first (as MapScalar.as_py does via keys()) and + # detect duplicates before converting any value, so that the 'lossy' + # warnings and the 'strict' KeyError are emitted at the same point as + # in MapScalar.as_py even when a later value conversion raises. + cdef int64_t count = end - start + cdef list keys_py = [keys._getitem_py(j, None) for j in range(start, end)] + cdef dict result = {} + cdef int64_t k + if len(set(keys_py)) == count: + for k in range(count): + result[keys_py[k]] = items._getitem_py(start + k, maps_as_pydicts) + return result + # Duplicate keys: per-key loop matching MapScalar.as_py exactly. + for k in range(count): + key = keys_py[k] + if key in result: + if maps_as_pydicts == "strict": + raise KeyError( + "Converting to Python dictionary is not supported in strict mode " + f"when duplicate keys are present (duplicate key was '{key}')." + ) + else: + warnings.warn( + f"Encountered key '{key}' which was already encountered.") + result[key] = items._getitem_py(start + k, maps_as_pydicts) + return result + @staticmethod def from_arrays(offsets, keys, items, DataType type=None, MemoryPool pool=None, mask=None): """ @@ -3688,6 +3803,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, object maps_as_pydicts): + cdef CFixedSizeListArray* arr = 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 = self._children_cache + cdef int64_t j, start = arr.value_offset(i), end = arr.value_offset(i + 1) + return [values._getitem_py(j, maps_as_pydicts) for j in range(start, end)] + @staticmethod def from_arrays(values, list_size=None, DataType type=None, mask=None): """ @@ -3974,6 +4099,13 @@ cdef class StringArray(Array): Concrete class for Arrow arrays of string (or utf8) data type. """ + cdef object _getitem_py(self, int64_t i, object maps_as_pydicts): + if self.ap.IsNull(i): + return None + cdef cpp_string_view view = ( self.ap).GetView(i) + # Matches StringScalar.as_py, which is str(buf, 'utf8'). + return cp.PyUnicode_DecodeUTF8(view.data(), view.size(), NULL) + @staticmethod def from_buffers(int length, Buffer value_offsets, Buffer data, Buffer null_bitmap=None, int null_count=-1, @@ -4006,6 +4138,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, object maps_as_pydicts): + if self.ap.IsNull(i): + return None + cdef cpp_string_view view = ( self.ap).GetView(i) + return cp.PyUnicode_DecodeUTF8(view.data(), view.size(), NULL) + @staticmethod def from_buffers(int length, Buffer value_offsets, Buffer data, Buffer null_bitmap=None, int null_count=-1, @@ -4038,11 +4176,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, object maps_as_pydicts): + if self.ap.IsNull(i): + return None + cdef cpp_string_view view = ( self.ap).GetView(i) + return cp.PyUnicode_DecodeUTF8(view.data(), 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, object maps_as_pydicts): + if self.ap.IsNull(i): + return None + cdef cpp_string_view view = ( self.ap).GetView(i) + return cp.PyBytes_FromStringAndSize(view.data(), view.size()) + @property def total_values_length(self): """ @@ -4056,6 +4207,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, object maps_as_pydicts): + if self.ap.IsNull(i): + return None + cdef cpp_string_view view = ( self.ap).GetView(i) + return cp.PyBytes_FromStringAndSize(view.data(), view.size()) + @property def total_values_length(self): """ @@ -4070,6 +4228,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, object maps_as_pydicts): + if self.ap.IsNull(i): + return None + cdef cpp_string_view view = ( self.ap).GetView(i) + return cp.PyBytes_FromStringAndSize(view.data(), view.size()) + cdef class DictionaryArray(Array): """ @@ -4229,6 +4393,28 @@ cdef class StructArray(Array): Concrete class for Arrow arrays of a struct data type. """ + cdef object _getitem_py(self, int64_t i, object maps_as_pydicts): + 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 = ( self._children_cache)[0] + fields = ( self._children_cache)[1] + cdef Array field_arr + result = {} + for k in range(num_fields): + field_arr = fields[k] + result[names[k]] = field_arr._getitem_py(i, maps_as_pydicts) + return result + def field(self, index): """ Retrieves the child array belonging to field. diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 8b4786ecbf13..05467441ee19 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -946,6 +946,7 @@ 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) @@ -953,6 +954,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil: 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) @@ -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, diff --git a/python/pyarrow/lib.pxd b/python/pyarrow/lib.pxd index 683faa7855c5..8e954323f8a9 100644 --- a/python/pyarrow/lib.pxd +++ b/python/pyarrow/lib.pxd @@ -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, object maps_as_pydicts) cdef int64_t length(self) cdef void _assert_cpu(self) except * diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index adc3e097b54a..9128647e3dc7 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -465,6 +465,97 @@ 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())), + ] + 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): + dup.to_pylist() + + +def test_to_pylist_maps_as_pydicts(): + # GH-50429: maps_as_pydicts converts through the scalar-free path; the + # semantics must match MapScalar.as_py exactly. + map_type = pa.map_(pa.string(), pa.int32()) + arrays = [ + pa.array([[("k1", 1), ("k2", None)], None, []], type=map_type), + pa.array([[[("k", 1)], None], None], type=pa.list_(map_type)), + pa.array([[("o", [("i", 5)])]], + type=pa.map_(pa.string(), map_type)), + pa.array([{"m": [("k", 1)]}, None], + type=pa.struct([("m", map_type)])), + ] + for arr in arrays: + for mode in ("lossy", "strict"): + for view in (arr, arr.slice(1)): + expected = [x.as_py(maps_as_pydicts=mode) for x in view] + assert view.to_pylist(maps_as_pydicts=mode) == expected + + dup = pa.array([[("k", 1), ("k", 2)]], type=map_type) + with pytest.warns(UserWarning, match="already encountered"): + assert dup.to_pylist(maps_as_pydicts="lossy") == [{"k": 2}] + with pytest.raises(KeyError, match="strict mode"): + dup.to_pylist(maps_as_pydicts="strict") + + # Duplicate keys must be detected before converting values: with a + # poison value *after* the duplicate, strict mode raises the outer + # duplicate-key error, and lossy mode still emits its warning first. + nested_map = pa.map_(pa.string(), map_type) + poison = pa.array( + [[("k1", [("a", 1)]), ("k1", [("d", 1), ("d", 2)])]], type=nested_map) + with pytest.raises(KeyError, match="duplicate key was 'k1'"): + poison.to_pylist(maps_as_pydicts="strict") + with pytest.warns(UserWarning, match="Encountered key 'k1'"): + assert poison.to_pylist(maps_as_pydicts="lossy") == [{"k1": {"d": 2}}] + with pytest.raises(ValueError, match="Invalid value for 'maps_as_pydicts'"): + dup.to_pylist(maps_as_pydicts="bogus") + # invalid values are only rejected when a map value is converted, + # matching the Scalar-based behavior + assert pa.array([1, 2]).to_pylist(maps_as_pydicts="bogus") == [1, 2] + + def test_array_slice(): arr = pa.array(range(10))