ENH: implement tm.shares_memory (#44747)

Author: jbrockmendel

pandas/_testing/__init__.py

@@ -110,12 +110,15 @@
     UInt64Index,
 )
 from pandas.core.arrays import (
+    BaseMaskedArray,
     DatetimeArray,
+    ExtensionArray,
     PandasArray,
     PeriodArray,
     TimedeltaArray,
     period_array,
 )
+from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
 
 if TYPE_CHECKING:
     from pandas import (
@@ -1050,3 +1053,51 @@ def at(x):
 
 def iat(x):
     return x.iat
+
+
+# -----------------------------------------------------------------------------
+
+
+def shares_memory(left, right) -> bool:
+    """
+    Pandas-compat for np.shares_memory.
+    """
+    if isinstance(left, np.ndarray) and isinstance(right, np.ndarray):
+        return np.shares_memory(left, right)
+    elif isinstance(left, np.ndarray):
+        # Call with reversed args to get to unpacking logic below.
+        return shares_memory(right, left)
+
+    if isinstance(left, RangeIndex):
+        return False
+    if isinstance(left, MultiIndex):
+        return shares_memory(left._codes, right)
+    if isinstance(left, (Index, Series)):
+        return shares_memory(left._values, right)
+
+    if isinstance(left, NDArrayBackedExtensionArray):
+        return shares_memory(left._ndarray, right)
+    if isinstance(left, pd.SparseArray):
+        return shares_memory(left.sp_values, right)
+
+    if isinstance(left, ExtensionArray) and left.dtype == "string[pyarrow]":
+        # https://github.com/pandas-dev/pandas/pull/43930#discussion_r736862669
+        if isinstance(right, ExtensionArray) and right.dtype == "string[pyarrow]":
+            left_pa_data = left._data
+            right_pa_data = right._data
+            left_buf1 = left_pa_data.chunk(0).buffers()[1]
+            right_buf1 = right_pa_data.chunk(0).buffers()[1]
+            return left_buf1 == right_buf1
+
+    if isinstance(left, BaseMaskedArray) and isinstance(right, BaseMaskedArray):
+        # By convention, we'll say these share memory if they share *either*
+        #  the _data or the _mask
+        return np.shares_memory(left._data, right._data) or np.shares_memory(
+            left._mask, right._mask
+        )
+
+    if isinstance(left, DataFrame) and len(left._mgr.arrays) == 1:
+        arr = left._mgr.arrays[0]
+        return shares_memory(arr, right)
+
+    raise NotImplementedError(type(left), type(right))

pandas/tests/arrays/categorical/test_constructors.py

@@ -738,7 +738,7 @@ def test_from_sequence_copy(self):
 
         result = Categorical._from_sequence(cat, dtype=None, copy=True)
 
-        assert not np.shares_memory(result._codes, cat._codes)
+        assert not tm.shares_memory(result, cat)
 
     @pytest.mark.xfail(
         not IS64 or is_platform_windows(),

pandas/tests/arrays/floating/test_arithmetic.py

@@ -199,6 +199,5 @@ def test_unary_float_operators(float_ea_dtype, source, neg_target, abs_target):
 
     tm.assert_extension_array_equal(neg_result, neg_target)
     tm.assert_extension_array_equal(pos_result, arr)
-    assert not np.shares_memory(pos_result._data, arr._data)
-    assert not np.shares_memory(pos_result._mask, arr._mask)
+    assert not tm.shares_memory(pos_result, arr)
     tm.assert_extension_array_equal(abs_result, abs_target)

pandas/tests/arrays/floating/test_astype.py

@@ -78,8 +78,7 @@ def test_astype_copy():
     # copy=True -> ensure both data and mask are actual copies
     result = arr.astype("Float64", copy=True)
     assert result is not arr
-    assert not np.shares_memory(result._data, arr._data)
-    assert not np.shares_memory(result._mask, arr._mask)
+    assert not tm.shares_memory(result, arr)
     result[0] = 10
     tm.assert_extension_array_equal(arr, orig)
     result[0] = pd.NA
@@ -101,8 +100,7 @@ def test_astype_copy():
     orig = pd.array([0.1, 0.2, None], dtype="Float64")
 
     result = arr.astype("Float32", copy=False)
-    assert not np.shares_memory(result._data, arr._data)
-    assert not np.shares_memory(result._mask, arr._mask)
+    assert not tm.shares_memory(result, arr)
     result[0] = 10
     tm.assert_extension_array_equal(arr, orig)
     result[0] = pd.NA

pandas/tests/arrays/integer/test_arithmetic.py

@@ -299,6 +299,5 @@ def test_unary_int_operators(any_signed_int_ea_dtype, source, neg_target, abs_ta
 
     tm.assert_extension_array_equal(neg_result, neg_target)
     tm.assert_extension_array_equal(pos_result, arr)
-    assert not np.shares_memory(pos_result._data, arr._data)
-    assert not np.shares_memory(pos_result._mask, arr._mask)
+    assert not tm.shares_memory(pos_result, arr)
     tm.assert_extension_array_equal(abs_result, abs_target)

pandas/tests/arrays/integer/test_dtypes.py

@@ -153,8 +153,7 @@ def test_astype_copy():
     # copy=True -> ensure both data and mask are actual copies
     result = arr.astype("Int64", copy=True)
     assert result is not arr
-    assert not np.shares_memory(result._data, arr._data)
-    assert not np.shares_memory(result._mask, arr._mask)
+    assert not tm.shares_memory(result, arr)
     result[0] = 10
     tm.assert_extension_array_equal(arr, orig)
     result[0] = pd.NA
@@ -176,8 +175,7 @@ def test_astype_copy():
     orig = pd.array([1, 2, 3, None], dtype="Int64")
 
     result = arr.astype("Int32", copy=False)
-    assert not np.shares_memory(result._data, arr._data)
-    assert not np.shares_memory(result._mask, arr._mask)
+    assert not tm.shares_memory(result, arr)
     result[0] = 10
     tm.assert_extension_array_equal(arr, orig)
     result[0] = pd.NA

pandas/tests/arrays/masked_shared.py

@@ -1,7 +1,6 @@
 """
 Tests shared by MaskedArray subclasses.
 """
-import numpy as np
 
 import pandas as pd
 import pandas._testing as tm
@@ -56,7 +55,7 @@ class NumericOps:
 
     def test_no_shared_mask(self, data):
         result = data + 1
-        assert np.shares_memory(result._mask, data._mask) is False
+        assert not tm.shares_memory(result, data)
 
     def test_array(self, comparison_op, dtype):
         op = comparison_op

pandas/tests/arrays/test_array.py

@@ -184,15 +184,15 @@ def test_array_copy():
     a = np.array([1, 2])
     # default is to copy
     b = pd.array(a, dtype=a.dtype)
-    assert np.shares_memory(a, b._ndarray) is False
+    assert not tm.shares_memory(a, b)
 
     # copy=True
     b = pd.array(a, dtype=a.dtype, copy=True)
-    assert np.shares_memory(a, b._ndarray) is False
+    assert not tm.shares_memory(a, b)
 
     # copy=False
     b = pd.array(a, dtype=a.dtype, copy=False)
-    assert np.shares_memory(a, b._ndarray) is True
+    assert tm.shares_memory(a, b)
 
 
 cet = pytz.timezone("CET")

pandas/tests/arrays/test_numpy.py

@@ -128,7 +128,7 @@ def test_constructor_copy():
     arr = np.array([0, 1])
     result = PandasArray(arr, copy=True)
 
-    assert np.shares_memory(result._ndarray, arr) is False
+    assert not tm.shares_memory(result, arr)
 
 
 def test_constructor_with_data(any_numpy_array):

pandas/tests/arrays/test_timedeltas.py

@@ -99,11 +99,11 @@ def test_pos(self):
 
         result = +arr
         tm.assert_timedelta_array_equal(result, arr)
-        assert not np.shares_memory(result._ndarray, arr._ndarray)
+        assert not tm.shares_memory(result, arr)
 
         result2 = np.positive(arr)
         tm.assert_timedelta_array_equal(result2, arr)
-        assert not np.shares_memory(result2._ndarray, arr._ndarray)
+        assert not tm.shares_memory(result2, arr)
 
     def test_neg(self):
         vals = np.array([-3600 * 10 ** 9, "NaT", 7200 * 10 ** 9], dtype="m8[ns]")

pandas/tests/series/indexing/test_setitem.py

@@ -562,7 +562,7 @@ def _check_inplace(self, is_inplace, orig, arr, obj):
             if arr.dtype.kind in ["m", "M"]:
                 # We may not have the same DTA/TDA, but will have the same
                 #  underlying data
-                assert arr._data is obj._values._data
+                assert arr._ndarray is obj._values._ndarray
             else:
                 assert obj._values is arr
         else: