Pandas ExtensionArray 的简单示例
在我看来,PandasExtensionArray将是一个简单的入门示例真正有帮助的情况之一。但是,我还没有在任何地方找到足够简单的示例。
创建一个 ExtensionArray
要创建ExtensionArray,您需要
- 创建一个
ExtensionDtype并注册它 -
ExtensionArray通过实现所需的方法来创建一个。
Pandas 文档中还有一个部分提供了简要概述。
示例实现
有很多实现的例子:
- Pandas 自己的内部扩展数组
-
Geopandas “
GeometryArray - Pandas 文档有一个带有扩展数据类型的项目列表
- 例如CyberPandas '
IPArray
- 例如CyberPandas '
- 网络上的许多其他人,例如Fletcher的
StringSupportingExtensionArray
题
尽管研究了以上所有内容,我仍然发现扩展数组难以理解。所有示例都有很多细节和自定义功能,这使得很难确定实际需要什么。我怀疑很多人都遇到过类似的问题。
因此,我要求提供一个简单且最小的工作ExtensionArray.
举一个具体的例子,假设我想扩展ExtensionArray以获得一个能够保存 NA 值的整数数组。这本质上是IntegerArray,但剥离了超出ExtensionArray.
回答
准系统示例 #1: NullableIntArray
Pandas 包装了 numpy,在 numpy 中混合数据类型需要这样做dtype=object效率不高。IntegerArray跳过箍以避免object效率,但我们只对简单感兴趣,所以:
- 为了去除复杂性,忘记效率,只使用一个 numpy
object数组。 - 为了简化空感知计算,将此
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)
- 很好的解释!