Extension arrays and accessors

pandas · Expert · 20 min

Libraries like GeoPandas and pint-pandas add new column types this way; knowing it lets you build or debug them.

pandas started as a layer over NumPy, which limited it to NumPy's types. The extension array interface lets any array-like type live in a Series or DataFrame column and take part in indexing, take, concat, groupby, isna and I/O.

Two classes define a new type:

  • An `ExtensionDtype` subclass describes it: name, type (the scalar type), na_value, and construct_array_type(). Decorate it with @pd.api.extensions.register_extension_dtype so strings like astype("mytype") work.
  • An `ExtensionArray` subclass holds the data. The required methods are small: _from_sequence, _from_factorized, __getitem__, __len__, __eq__, dtype, nbytes, isna, take, copy and _concat_same_type. Everything else (unique, factorize, value_counts, sorting) has generic defaults built on these, which you can override with faster versions. ExtensionScalarOpsMixin adds arithmetic by applying scalar ops.

pandas' own types use this interface: Categorical, IntegerArray (Int64), BooleanArray, StringArray, DatetimeArray with time zones, PeriodArray, IntervalArray, SparseArray and ArrowExtensionArray. Third-party examples include GeoPandas geometry columns and pint-pandas unit-aware quantities.

In the BlockManager, extension arrays live in 1-D blocks, never consolidated, so a frame with many extension columns has many blocks. Methods that need NumPy call to_numpy(), which may cost a conversion.

Most teams don't need a new dtype; they want domain methods. Accessors cover that cheaply: @pd.api.extensions.register_series_accessor("geo") (or register_dataframe_accessor) attaches a namespace, so s.geo.distance_to(...) works like .str and .dt, with validation in __init__.

Example

import pandas as pd

@pd.api.extensions.register_series_accessor("money")
class MoneyAccessor:
    def __init__(self, s: pd.Series):
        if not pd.api.types.is_numeric_dtype(s):
            raise AttributeError("money accessor needs numeric data")
        self._s = s

    def fmt(self, symbol="EUR"):
        return self._s.map(lambda v: f"{v:,.2f} {symbol}")

    def vat(self, rate=0.2):
        return (self._s * rate).round(2)

prices = pd.Series([1234.5, 99.0])
print(prices.money.fmt().tolist())   # ['1,234.50 EUR', '99.00 EUR']
print(prices.money.vat().tolist())   # [246.9, 19.8]

# Built-in extension arrays behind familiar dtypes
for s in [pd.Series([1, None], dtype="Int64"),
          pd.Series(["a", "b"], dtype="category"),
          pd.Series(pd.date_range("2024", periods=2, tz="UTC"))]:
    print(type(s.array).__name__, s.dtype)
# IntegerArray Int64 / Categorical category / DatetimeArray datetime64[ns, UTC] (us on 3.x)
print(isinstance(pd.Series([1, None], dtype="Int64").array, pd.api.extensions.ExtensionArray))   # True

Remember

  • ExtensionDtype + ExtensionArray let any array type become a pandas column
  • A small set of required methods unlocks indexing, concat, groupby and more
  • Categorical, Int64, tz-aware datetimes and Arrow types are all extension arrays
  • Accessors add domain methods without defining a new dtype

Watch out for

  • Building a full ExtensionArray when an accessor over an existing dtype would meet the need.

Check yourself

  1. What does `s.array` return for a Series with dtype 'category'?
  2. When is a custom accessor a better choice than a new ExtensionArray?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: pandas.pydata.org/docs/development/extending.html