CoW: add readonly flag to ExtensionArrays, return read-only EA/ndarray in .array/EA.to_numpy() (#61925)

Author: jorisvandenbossche

doc/source/whatsnew/v3.0.0.rst

@@ -1359,6 +1359,10 @@ Other
 - Fixed bug in the :meth:`Series.rank` with object dtype and extremely small float values (:issue:`62036`)
 - Fixed bug where the :class:`DataFrame` constructor misclassified array-like objects with a ``.name`` attribute as :class:`Series` or :class:`Index` (:issue:`61443`)
 - Fixed regression in :meth:`DataFrame.from_records` not initializing subclasses properly (:issue:`57008`)
+- Accessing the underlying NumPy array of a DataFrame or Series will return a read-only
+  array if the array shares data with the original DataFrame or Series (:ref:`copy_on_write_read_only_na`).
+  This logic is expanded to accessing the underlying pandas ExtensionArray
+  through ``.array`` (or ``.values`` depending on the dtype) as well (:issue:`61925`).
 
 .. ***DO NOT USE THIS SECTION***
 

pandas/_libs/ops.pyx

@@ -177,7 +177,7 @@ def vec_compare(ndarray[object] left, ndarray[object] right, object op) -> ndarr
 
 @cython.wraparound(False)
 @cython.boundscheck(False)
-def scalar_binop(object[:] values, object val, object op) -> ndarray:
+def scalar_binop(ndarray[object] values, object val, object op) -> ndarray:
     """
     Apply the given binary operator `op` between each element of the array
     `values` and the scalar `val`.
@@ -214,7 +214,7 @@ def scalar_binop(object[:] values, object val, object op) -> ndarray:
 
 @cython.wraparound(False)
 @cython.boundscheck(False)
-def vec_binop(object[:] left, object[:] right, object op) -> ndarray:
+def vec_binop(ndarray[object] left, ndarray[object] right, object op) -> ndarray:
     """
     Apply the given binary operator `op` pointwise to the elements of
     arrays `left` and `right`.

pandas/core/arrays/_mixins.py

@@ -54,7 +54,10 @@
 from pandas.core.array_algos.transforms import shift
 from pandas.core.arrays.base import ExtensionArray
 from pandas.core.construction import extract_array
-from pandas.core.indexers import check_array_indexer
+from pandas.core.indexers import (
+    check_array_indexer,
+    getitem_returns_view,
+)
 from pandas.core.sorting import nargminmax
 
 if TYPE_CHECKING:
@@ -258,6 +261,9 @@ def shift(self, periods: int = 1, fill_value=None) -> Self:
         return self._from_backing_data(new_values)
 
     def __setitem__(self, key, value) -> None:
+        if self._readonly:
+            raise ValueError("Cannot modify read-only array")
+
         key = check_array_indexer(self, key)
         value = self._validate_setitem_value(value)
         self._ndarray[key] = value
@@ -283,7 +289,10 @@ def __getitem__(
             result = self._ndarray[key]
             if self.ndim == 1:
                 return self._box_func(result)
-            return self._from_backing_data(result)
+            result = self._from_backing_data(result)
+            if getitem_returns_view(self, key):
+                result._readonly = self._readonly
+            return result
 
         # error: Incompatible types in assignment (expression has type "ExtensionArray",
         # variable has type "Union[int, slice, ndarray]")
@@ -294,6 +303,8 @@ def __getitem__(
             return self._box_func(result)
 
         result = self._from_backing_data(result)
+        if getitem_returns_view(self, key):
+            result._readonly = self._readonly
         return result
 
     def _pad_or_backfill(

pandas/core/arrays/arrow/array.py

@@ -78,6 +78,7 @@
 from pandas.core.construction import extract_array
 from pandas.core.indexers import (
     check_array_indexer,
+    getitem_returns_view,
     unpack_tuple_and_ellipses,
     validate_indices,
 )
@@ -790,7 +791,10 @@ def __getitem__(self, item: PositionalIndexer):
 
         value = self._pa_array[item]
         if isinstance(value, pa.ChunkedArray):
-            return self._from_pyarrow_array(value)
+            result = self._from_pyarrow_array(value)
+            if getitem_returns_view(self, item):
+                result._readonly = self._readonly
+            return result
         else:
             pa_type = self._pa_array.type
             scalar = value.as_py()
@@ -2196,6 +2200,9 @@ def __setitem__(self, key, value) -> None:
         -------
         None
         """
+        if self._readonly:
+            raise ValueError("Cannot modify read-only array")
+
         # GH50085: unwrap 1D indexers
         if isinstance(key, tuple) and len(key) == 1:
             key = key[0]

pandas/core/arrays/base.py

@@ -37,6 +37,7 @@
     validate_insert_loc,
 )
 
+from pandas.core.dtypes.astype import astype_is_view
 from pandas.core.dtypes.common import (
     is_list_like,
     is_scalar,
@@ -268,6 +269,8 @@ class ExtensionArray:
     #  strictly less than 2000 to be below Index.__pandas_priority__.
     __pandas_priority__ = 1000
 
+    _readonly = False
+
     # ------------------------------------------------------------------------
     # Constructors
     # ------------------------------------------------------------------------
@@ -454,6 +457,11 @@ def __setitem__(self, key, value) -> None:
         Returns
         -------
         None
+
+        Raises
+        ------
+        ValueError
+            If the array is readonly and modification is attempted.
         """
         # Some notes to the ExtensionArray implementer who may have ended up
         # here. While this method is not required for the interface, if you
@@ -473,6 +481,10 @@ def __setitem__(self, key, value) -> None:
         #   __init__ method coerces that value, then so should __setitem__
         # Note, also, that Series/DataFrame.where internally use __setitem__
         # on a copy of the data.
+        # Check if the array is readonly
+        if self._readonly:
+            raise ValueError("Cannot modify read-only array")
+
         raise NotImplementedError(f"{type(self)} does not implement __setitem__.")
 
     def __len__(self) -> int:
@@ -567,8 +579,14 @@ def to_numpy(
         result = np.asarray(self, dtype=dtype)
         if copy or na_value is not lib.no_default:
             result = result.copy()
+        elif self._readonly and astype_is_view(self.dtype, result.dtype):
+            # If the ExtensionArray is readonly, make the numpy array readonly too
+            result = result.view()
+            result.flags.writeable = False
+
         if na_value is not lib.no_default:
             result[self.isna()] = na_value  # type: ignore[index]
+
         return result
 
     # ------------------------------------------------------------------------

pandas/core/arrays/datetimelike.py

@@ -368,7 +368,12 @@ def __array__(
 
         if copy is True:
             return np.array(self._ndarray, dtype=dtype)
-        return self._ndarray
+
+        result = self._ndarray
+        if self._readonly:
+            result = result.view()
+            result.flags.writeable = False
+        return result
 
     @overload
     def __getitem__(self, key: ScalarIndexer) -> DTScalarOrNaT: ...

pandas/core/arrays/interval.py

@@ -87,7 +87,10 @@
     ensure_wrapped_if_datetimelike,
     extract_array,
 )
-from pandas.core.indexers import check_array_indexer
+from pandas.core.indexers import (
+    check_array_indexer,
+    getitem_returns_view,
+)
 from pandas.core.ops import (
     invalid_comparison,
     unpack_zerodim_and_defer,
@@ -842,9 +845,15 @@ def __getitem__(self, key: PositionalIndexer) -> Self | IntervalOrNA:
         # "Union[Period, Timestamp, Timedelta, NaTType, DatetimeArray, TimedeltaArray,
         # ndarray[Any, Any]]"; expected "Union[Union[DatetimeArray, TimedeltaArray],
         # ndarray[Any, Any]]"
-        return self._simple_new(left, right, dtype=self.dtype)  # type: ignore[arg-type]
+        result = self._simple_new(left, right, dtype=self.dtype)  # type: ignore[arg-type]
+        if getitem_returns_view(self, key):
+            result._readonly = self._readonly
+        return result
 
     def __setitem__(self, key, value) -> None:
+        if self._readonly:
+            raise ValueError("Cannot modify read-only array")
+
         value_left, value_right = self._validate_setitem_value(value)
         key = check_array_indexer(self, key)
 

pandas/core/arrays/masked.py

@@ -26,6 +26,7 @@
 )
 from pandas.errors import AbstractMethodError
 
+from pandas.core.dtypes.astype import astype_is_view
 from pandas.core.dtypes.base import ExtensionDtype
 from pandas.core.dtypes.cast import maybe_downcast_to_dtype
 from pandas.core.dtypes.common import (
@@ -75,7 +76,10 @@
     ensure_wrapped_if_datetimelike,
     extract_array,
 )
-from pandas.core.indexers import check_array_indexer
+from pandas.core.indexers import (
+    check_array_indexer,
+    getitem_returns_view,
+)
 from pandas.core.ops import invalid_comparison
 from pandas.core.util.hashing import hash_array
 
@@ -212,7 +216,10 @@ def __getitem__(self, item: PositionalIndexer) -> Self | Any:
                 return self.dtype.na_value
             return self._data[item]
 
-        return self._simple_new(self._data[item], newmask)
+        result = self._simple_new(self._data[item], newmask)
+        if getitem_returns_view(self, item):
+            result._readonly = self._readonly
+        return result
 
     def _pad_or_backfill(
         self,
@@ -354,6 +361,9 @@ def _validate_setitem_value(self, value):
         raise TypeError(f"Invalid value '{value!s}' for dtype '{self.dtype}'")
 
     def __setitem__(self, key, value) -> None:
+        if self._readonly:
+            raise ValueError("Cannot modify read-only array")
+
         key = check_array_indexer(self, key)
 
         if is_scalar(value):
@@ -566,11 +576,11 @@ def to_numpy(
         hasna = self._hasna
         dtype, na_value = to_numpy_dtype_inference(self, dtype, na_value, hasna)
         if dtype is None:
-            dtype = object
+            dtype = np.dtype(object)
 
         if hasna:
             if (
-                dtype != object
+                dtype != np.dtype(object)
                 and not is_string_dtype(dtype)
                 and na_value is libmissing.NA
             ):
@@ -588,6 +598,9 @@ def to_numpy(
             with warnings.catch_warnings():
                 warnings.filterwarnings("ignore", category=RuntimeWarning)
                 data = self._data.astype(dtype, copy=copy)
+            if self._readonly and not copy and astype_is_view(self.dtype, dtype):
+                data = data.view()
+                data.flags.writeable = False
         return data
 
     def tolist(self) -> list:
@@ -686,7 +699,12 @@ def __array__(
         if copy is False:
             if not self._hasna:
                 # special case, here we can simply return the underlying data
-                return np.array(self._data, dtype=dtype, copy=copy)
+                result = np.array(self._data, dtype=dtype, copy=copy)
+                # If the ExtensionArray is readonly, make the numpy array readonly too
+                if self._readonly:
+                    result = result.view()
+                    result.flags.writeable = False
+                return result
             raise ValueError(
                 "Unable to avoid copy while creating an array as requested."
             )

pandas/core/arrays/numpy_.py

@@ -14,7 +14,10 @@
 from pandas._libs.tslibs import is_supported_dtype
 from pandas.compat.numpy import function as nv
 
-from pandas.core.dtypes.astype import astype_array
+from pandas.core.dtypes.astype import (
+    astype_array,
+    astype_is_view,
+)
 from pandas.core.dtypes.cast import (
     construct_1d_object_array_from_listlike,
     maybe_downcast_to_dtype,
@@ -179,12 +182,23 @@ def dtype(self) -> NumpyEADtype:
     # NumPy Array Interface
 
     def __array__(
-        self, dtype: NpDtype | None = None, copy: bool | None = None
+        self, dtype: np.dtype | None = None, copy: bool | None = None
     ) -> np.ndarray:
         if copy is not None:
             # Note: branch avoids `copy=None` for NumPy 1.x support
-            return np.array(self._ndarray, dtype=dtype, copy=copy)
-        return np.asarray(self._ndarray, dtype=dtype)
+            result = np.array(self._ndarray, dtype=dtype, copy=copy)
+        else:
+            result = np.asarray(self._ndarray, dtype=dtype)
+
+        if (
+            self._readonly
+            and not copy
+            and (dtype is None or astype_is_view(self.dtype, dtype))
+        ):
+            result = result.view()
+            result.flags.writeable = False
+
+        return result
 
     def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
         # Lightly modified version of
@@ -545,6 +559,9 @@ def to_numpy(
             result[mask] = na_value
         else:
             result = self._ndarray
+            if not copy and self._readonly:
+                result = result.view()
+                result.flags.writeable = False
 
         result = np.asarray(result, dtype=dtype)
 

pandas/core/arrays/period.py

@@ -397,7 +397,11 @@ def __array__(
             # For NumPy 1.x compatibility we cannot use copy=None.  And
             # `copy=False` has the meaning of `copy=None` here:
             if not copy:
-                return np.asarray(self.asi8, dtype=dtype)
+                result = np.asarray(self.asi8, dtype=dtype)
+                if self._readonly:
+                    result = result.view()
+                    result.flags.writeable = False
+                return result
             else:
                 return np.array(self.asi8, dtype=dtype)