done

lib/python3.11/site-packages/pandas/tests/extension/array_with_attr/__init__.py

@@ -0,0 +1,6 @@
+from pandas.tests.extension.array_with_attr.array import (
+    FloatAttrArray,
+    FloatAttrDtype,
+)
+
+__all__ = ["FloatAttrArray", "FloatAttrDtype"]

lib/python3.11/site-packages/pandas/tests/extension/array_with_attr/array.py

@@ -0,0 +1,89 @@
+"""
+Test extension array that has custom attribute information (not stored on the dtype).
+
+"""
+from __future__ import annotations
+
+import numbers
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from pandas.core.dtypes.base import ExtensionDtype
+
+import pandas as pd
+from pandas.core.arrays import ExtensionArray
+
+if TYPE_CHECKING:
+    from pandas._typing import type_t
+
+
+class FloatAttrDtype(ExtensionDtype):
+    type = float
+    name = "float_attr"
+    na_value = np.nan
+
+    @classmethod
+    def construct_array_type(cls) -> type_t[FloatAttrArray]:
+        """
+        Return the array type associated with this dtype.
+
+        Returns
+        -------
+        type
+        """
+        return FloatAttrArray
+
+
+class FloatAttrArray(ExtensionArray):
+    dtype = FloatAttrDtype()
+    __array_priority__ = 1000
+
+    def __init__(self, values, attr=None) -> None:
+        if not isinstance(values, np.ndarray):
+            raise TypeError("Need to pass a numpy array of float64 dtype as values")
+        if not values.dtype == "float64":
+            raise TypeError("Need to pass a numpy array of float64 dtype as values")
+        self.data = values
+        self.attr = attr
+
+    @classmethod
+    def _from_sequence(cls, scalars, *, dtype=None, copy=False):
+        if not copy:
+            data = np.asarray(scalars, dtype="float64")
+        else:
+            data = np.array(scalars, dtype="float64", copy=copy)
+        return cls(data)
+
+    def __getitem__(self, item):
+        if isinstance(item, numbers.Integral):
+            return self.data[item]
+        else:
+            # slice, list-like, mask
+            item = pd.api.indexers.check_array_indexer(self, item)
+            return type(self)(self.data[item], self.attr)
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def isna(self):
+        return np.isnan(self.data)
+
+    def take(self, indexer, allow_fill=False, fill_value=None):
+        from pandas.api.extensions import take
+
+        data = self.data
+        if allow_fill and fill_value is None:
+            fill_value = self.dtype.na_value
+
+        result = take(data, indexer, fill_value=fill_value, allow_fill=allow_fill)
+        return type(self)(result, self.attr)
+
+    def copy(self):
+        return type(self)(self.data.copy(), self.attr)
+
+    @classmethod
+    def _concat_same_type(cls, to_concat):
+        data = np.concatenate([x.data for x in to_concat])
+        attr = to_concat[0].attr if len(to_concat) else None
+        return cls(data, attr)

lib/python3.11/site-packages/pandas/tests/extension/array_with_attr/test_array_with_attr.py

@@ -0,0 +1,33 @@
+import numpy as np
+
+import pandas as pd
+import pandas._testing as tm
+from pandas.tests.extension.array_with_attr import FloatAttrArray
+
+
+def test_concat_with_all_na():
+    # https://github.com/pandas-dev/pandas/pull/47762
+    # ensure that attribute of the column array is preserved (when it gets
+    # preserved in reindexing the array) during merge/concat
+    arr = FloatAttrArray(np.array([np.nan, np.nan], dtype="float64"), attr="test")
+
+    df1 = pd.DataFrame({"col": arr, "key": [0, 1]})
+    df2 = pd.DataFrame({"key": [0, 1], "col2": [1, 2]})
+    result = pd.merge(df1, df2, on="key")
+    expected = pd.DataFrame({"col": arr, "key": [0, 1], "col2": [1, 2]})
+    tm.assert_frame_equal(result, expected)
+    assert result["col"].array.attr == "test"
+
+    df1 = pd.DataFrame({"col": arr, "key": [0, 1]})
+    df2 = pd.DataFrame({"key": [0, 2], "col2": [1, 2]})
+    result = pd.merge(df1, df2, on="key")
+    expected = pd.DataFrame({"col": arr.take([0]), "key": [0], "col2": [1]})
+    tm.assert_frame_equal(result, expected)
+    assert result["col"].array.attr == "test"
+
+    result = pd.concat([df1.set_index("key"), df2.set_index("key")], axis=1)
+    expected = pd.DataFrame(
+        {"col": arr.take([0, 1, -1]), "col2": [1, np.nan, 2], "key": [0, 1, 2]}
+    ).set_index("key")
+    tm.assert_frame_equal(result, expected)
+    assert result["col"].array.attr == "test"