如何为Polars Series关联元数据并实现其保留与合并？

一、Polars原生元数据支持现状 #

Polars本身没有内置的Series级元数据存储机制，需要通过扩展类型或自定义DataFrame/Series来实现你的需求。

二、方案1：自定义DataType（推荐） #

利用Polars的扩展类型系统，创建包含元数据（如单位、变量类型）的自定义数据类型，元数据会和Series的类型绑定，自动跟随所有操作（重命名、排序、选择等），同时可以定义运算时的元数据合并规则。

示例：带单位的数值类型 #

import polars as pl
from polars.datatypes import DataType, Float64

class UnitFloat(DataType):
    def __init__(self, unit: str):
        self.unit = unit
        self.base_type = Float64
    
    def __eq__(self, other):
        return isinstance(other, UnitFloat) and self.unit == other.unit
    
    def __str__(self):
        return f"UnitFloat[{self.unit}]"
    
    # 定义除法运算的元数据合并规则
    def div(self, other):
        if isinstance(other, UnitFloat):
            new_unit = f"{self.unit}/{other.unit}"
            return UnitFloat(new_unit)
        raise TypeError("仅支持UnitFloat类型间的除法运算")

# 创建带单位的Series
s_m = pl.Series([10, 20, 30], dtype=UnitFloat("m"))
s_s = pl.Series([2, 4, 6], dtype=UnitFloat("s"))

# 除法运算自动生成新单位
s_ms = s_m / s_s
print(s_ms.dtype)  # 输出 UnitFloat[m/s]

这种方式的优势是元数据完全和数据类型绑定，所有Polars原生操作（重命名、排序、筛选）都会自动保留类型信息，无需额外处理。

三、方案2：继承DataFrame/Series自定义元数据管理 #

如果需要更灵活的元数据（比如同时存储单位、变量类型等多维度信息），可以继承Polars的DataFrame和Series，自行维护元数据字典，并拦截操作方法来同步元数据，同时获取列映射关系。

核心实现要点 #

• 自定义类维护metadata字典，键为列名，值为元数据（如{"unit": "m", "var_type": "independent"}）
• 重写Polars的操作方法（如rename、select、sort等），执行原生操作后同步更新元数据，并记录列映射关系

示例代码 #

import polars as pl

class MetaDataFrame(pl.DataFrame):
    def __init__(self, data=None, metadata=None, **kwargs):
        super().__init__(data, **kwargs)
        self.metadata = metadata or {}
        self._sync_metadata()
    
    def _sync_metadata(self):
        # 移除不存在列的元数据
        for col in list(self.metadata.keys()):
            if col not in self.columns:
                del self.metadata[col]
        # 给新增列添加默认元数据（可选）
        for col in self.columns:
            if col not in self.metadata:
                self.metadata[col] = {"unit": None, "var_type": None}
    
    def rename(self, mapping, **kwargs):
        # 执行原生重命名
        new_df = super().rename(mapping, **kwargs)
        # 生成旧列名→新列名的映射
        col_mapping = {old: new for old, new in mapping.items()}
        # 补充未重命名的列
        for col in self.columns:
            if col not in col_mapping:
                col_mapping[col] = col
        # 更新元数据
        new_metadata = {col_mapping[old]: meta for old, meta in self.metadata.items()}
        # 返回自定义DataFrame实例
        return MetaDataFrame(new_df, metadata=new_metadata)
    
    def select(self, *args, **kwargs):
        new_df = super().select(*args, **kwargs)
        # 筛选元数据
        new_metadata = {col: self.metadata[col] for col in new_df.columns}
        return MetaDataFrame(new_df, metadata=new_metadata)
    
    # 自定义列间运算，合并元数据
    def divide_columns(self, col1: str, col2: str, new_col: str):
        result_series = self[col1] / self[col2]
        # 合并单位元数据
        unit1 = self.metadata[col1]["unit"]
        unit2 = self.metadata[col2]["unit"]
        new_unit = f"{unit1}/{unit2}" if unit1 and unit2 else None
        # 创建新DataFrame并添加元数据
        new_df = self.with_columns(result_series.alias(new_col))
        new_df.metadata[new_col] = {"unit": new_unit, "var_type": "derived"}
        return new_df

# 使用示例
df = MetaDataFrame({
    "distance": [10, 20, 30],
    "time": [2, 4, 6]
}, metadata={
    "distance": {"unit": "m", "var_type": "independent"},
    "time": {"unit": "s", "var_type": "independent"}
})

# 重命名操作，元数据自动同步
renamed_df = df.rename({"distance": "dist"})
print(renamed_df.metadata)
# 输出：{'dist': {'unit': 'm', 'var_type': 'independent'}, 'time': {'unit': 's', 'var_type': 'independent'}}

# 执行除法运算生成带新单位的列
speed_df = renamed_df.divide_columns("dist", "time", "speed")
print(speed_df.metadata["speed"])
# 输出：{'unit': 'm/s', 'var_type': 'derived'}

获取列映射关系 #

在自定义操作方法中可直接生成并返回映射：

• 重命名操作：传入的mapping参数本身就是旧列名→新列名的映射
• 选择/筛选操作：新列列表与原列名的对应关系为原列名→原列名（仅保留选中列）
• 排序操作：列名不变仅顺序改变，映射关系为原列名→原列名，顺序对应新DataFrame的columns属性

你可以修改重写的方法来返回映射，比如调整rename方法：

def rename(self, mapping, **kwargs):
    new_df = super().rename(mapping, **kwargs)
    col_mapping = {old: new for old, new in mapping.items()}
    for col in self.columns:
        if col not in col_mapping:
            col_mapping[col] = col
    new_metadata = {col_mapping[old]: meta for old, meta in self.metadata.items()}
    # 同时返回DataFrame和列映射关系
    return MetaDataFrame(new_df, metadata=new_metadata), col_mapping

四、注意事项 #

• 自定义类需确保所有Polars操作都返回自定义实例，否则元数据会丢失
• 对于LazyFrame操作，需额外自定义LazyFrame类并拦截计划中的操作
• 自定义DataType需遵循Polars扩展类型规范，部分复杂操作可能需要额外适配

内容的提问来源于stack exchange，提问作者Max Tyler