Pandas ExtensionArray 的简单示例

Question

在我看来，PandasExtensionArray将是一个简单的入门示例真正有帮助的情况之一。但是，我还没有在任何地方找到足够简单的示例。

创建一个 ExtensionArray

要创建ExtensionArray，您需要

• 创建一个ExtensionDtype并注册它
• ExtensionArray通过实现所需的方法来创建一个。

Pandas 文档中还有一个部分提供了简要概述。

示例实现

有很多实现的例子：

• Pandas 自己的内部扩展数组
• Geopandas “GeometryArray
• Pandas 文档有一个带有扩展数据类型的项目列表
  • 例如Cyber​​Pandas 'IPArray
• 网络上的许多其他人，例如Fletcher的StringSupportingExtensionArray

题

尽管研究了以上所有内容，我仍然发现扩展数组难以理解。所有示例都有很多细节和自定义功能，这使得很难确定实际需要什么。我怀疑很多人都遇到过类似的问题。

因此，我要求提供一个简单且最小的工作ExtensionArray.

举一个具体的例子，假设我想扩展ExtensionArray以获得一个能够保存 NA 值的整数数组。这本质上是IntegerArray，但剥离了超出ExtensionArray.

Answer 1

准系统示例 #1： NullableIntArray

假设我想要 ... 本质上IntegerArray，但剥离了超出ExtensionArray.

Pandas 包装了 numpy，在 numpy 中混合数据类型需要这样做dtype=object效率不高。IntegerArray跳过箍以避免object效率，但我们只对简单感兴趣，所以：

1. 为了去除复杂性，忘记效率，只使用一个 numpyobject数组。
2. 为了简化空感知计算，将此object数组包装在 numpy 中MaskedArray。

方法
构造函数	将传入数据转换object为int+数组nan，然后存储为MaskedArrayin self._ma。
吸气剂	返回self._ma[index]给定单个索引或NullableIntArray(self._ma[index])给定类似列表的索引。
二传手	设置self._ma[index]为（转换后的）传入值并相应地更新它mask。
计算器	调用各自的MaskedArray计算方法（自动处理空值）。

class NullableIntArray(pd.api.extensions.ExtensionArray):
    def __init__(self, data):
        isnull = pd.isnull(data)
        data = np.array(data, dtype=object)         # use object array to mix int/nan
        data[~isnull] = data[~isnull].astype(int)   # convert non-nan to int
        self._ma = np.ma.MaskedArray(data, isnull)  # store in MaskedArray for nan-awareness
    
    def __getitem__(self, index):
        if isinstance(index, int):
            item = self._ma.data[index]             # get value at index
        else:
            item = type(self)(self._ma.data[index]) # get NullableIntArray of subset
        return item
    
    def __setitem__(self, index: int, value) -> None:
        isnull = pd.isnull(value)
        self._ma.data[index] = value if isnull else int(value)
        self._ma.mask[index] = isnull
    
    def __repr__(self):
        return f'{type(self).__name__}({self._ma.data.tolist()}, dtype={self.dtype})'
    
    def sum(self) -> int:
        return self._ma.sum()  # MaskedArray.sum automatically handles nulls
    
    def mean(self) -> np.float64:
        return self._ma.mean() # MaskedArray.mean automatically handles nulls
    
    @property
    def dtype(self):
        return 'NullableInt'

使用NullableIntArray：

>>> a = NullableIntArray([1, np.nan, 5.0, np.nan, 9])
# NullableIntArray([1, nan, 5, nan, 9], dtype=NullableInt)

>>> a[0]
# 1

>>> a[1]
# nan

>>> a[1:]
# NullableIntArray([nan, 5, nan, 9], dtype=NullableInt)

>>> a[1] = 3.0
# NullableIntArray([1, 3, 5, nan, 9], dtype=NullableInt)

>>> a.sum() # 1+3+5+9
# 18

>>> a.mean() # 18/4
# 4.5

准系统示例#2： AutoNullArray

AutoNullArray自动用np.nan. 在这里，我们将数据存储在一个object数组中，并在构造函数/setter 中自动取消传入值：

class AutoNullArray(pd.api.extensions.ExtensionArray):
    def __init__(self, values, na_values=None):
        if na_values is None:
            na_values = ['', '<NA>', 'N/A', 'NA', 'NULL', 'NaN', 'n/a', 'nan', 'null']
        self._na_values = na_values
        
        self._values = np.array(values, dtype=object)     # use object array for simplicity
        self._values[np.isin(values, na_values)] = np.nan # replace na_values with nan
    
    def __getitem__(self, index):
        if isinstance(index, int):
            item = self._values[index]             # get value at index
        else:
            item = type(self)(self._values[index]) # get AutoNullArray of subset
        return item
    
    def __setitem__(self, index: int, value) -> None:
        if value in self._na_values:
            value = np.nan               # replace na_values with nan
        self._values[index] = value      # set to auto-nullified value

    def __repr__(self):
        return f'{type(self).__name__}({self._values.tolist()}, dtype={self.dtype})'
    
    def fillna(self, value=None) -> np.ndarray:
        isnull = pd.isnull(self._values)
        new_values = self._values.copy() # copy to avoid in-place modification
        new_values[isnull] = value       # fill null values with incoming value
        return type(self)(new_values)    # return result as AutoNullArray
    
    @property
    def dtype(self):
        return 'AutoNull'
    
    @property
    def na_values(self):
        return self._na_values

使用AutoNullArray：

>>> data = ['foo', 'NaN', 'bar', 'NULL', 42, '', -123.45, np.nan]
>>> a = AutoNullArray(data)
# AutoNullArray(['foo', nan, 'bar', nan, 42, nan, -123.45, nan], dtype=AutoNull)

>>> a.na_values
# ['', '<NA>', 'N/A', 'NA', 'NULL', 'NaN', 'n/a', 'nan', 'null']

>>> a[2]
# 'bar'

>>> a[-1]
# nan

>>> a[[2, 3, -2]]
# AutoNullArray(['bar', nan, -123.45], dtype=AutoNull)

>>> a[0] = 'hello'
# AutoNullArray(['hello', nan, 'bar', nan, 42, nan, -123.45, nan], dtype=AutoNull)

>>> a[0] = '<NA>'
# AutoNullArray([nan, nan, 'bar', nan, 42, nan, -123.45, nan], dtype=AutoNull)

>>> a.fillna('SO')
# AutoNullArray(['SO', 'SO', 'bar', 'SO', 42, 'SO', -123.45, 'SO'], dtype=AutoNull)