Polars 2.0.0でMap型が導入されました。
本記事では、Map型とネームスペースの拡張を活用して一次式を保持・計算してみました。
ここでいう一次式とは、$ 3x + 2y + 1 $のような式です。
前提
- Python: 3.14
- Polars: 2.0.0
一次式の計算
Polarsで下記を計算してみましょう。
E1 = x
E2 = x + y
A = 2
B = 1
E1 + E2 * A + B → 3x + 2y + 1
このような計算は1次元のベクトルでできますが、0が多いとムダな計算をします。
ここでは、0の係数を保持しない疎な表現としてMap型を使います。
なお、今回の実装ではmap_elements()を介してPython側で計算を行っているため、実測パフォーマンスの向上というよりは「スパースな一次式をDataFrame上で直感的に扱えるDSL（拡張表現）のプロトタイプ」という位置づけになります。
E1列、E2列、A列、B列を持つ1行のDataFrame（df）を作成します（map_exprモジュールについては後述します）。
import polars as pl
from map_expr import MAP_EXPR_DTYPE
df = pl.DataFrame(
{
"E1": [{"x": 1}],
"E2": [{"x": 1, "y": 1}],
"A": [2],
"B": [1],
},
schema_overrides={"E1": MAP_EXPR_DTYPE, "E2": MAP_EXPR_DTYPE},
)
print(df)
出力
shape: (1, 4)
┌───────────────┬──────────────────────┬─────┬─────┐
│ E1 ┆ E2 ┆ A ┆ B │
│ --- ┆ --- ┆ --- ┆ --- │
│ map[str, f64] ┆ map[str, f64] ┆ i64 ┆ i64 │
╞═══════════════╪══════════════════════╪═════╪═════╡
│ {"x": 1.0} ┆ {"x": 1.0, "y": 1.0} ┆ 2 ┆ 1 │
└───────────────┴──────────────────────┴─────┴─────┘
$ x + y $という一次式を{"x": 1.0, "y": 1.0}という辞書で保持します。
キーが変数名を、値が係数を意味します。
MAP_EXPR_DTYPEはpl.Map(pl.String, pl.Float64)です。
pl.Map(pl.String, pl.Float64)は、キーが文字列で値が実数を意味します。
続いて、E1 + E2 * A + Bを計算しましょう。しかし、Map型には四則演算がありません。
ここではlinというネームスペースを用意し、加減とスカラー倍・スカラー除算を定義しました。
次のように演算できます。
E1 = pl.col("E1")
E2 = pl.col("E2")
A = pl.col("A")
B = pl.col("B")
expr = E1.lin + E2.lin * A + B
print(df.select(Result=expr))
出力
shape: (1, 1)
┌───────────────────────────────┐
│ Result │
│ --- │
│ map[str, f64] │
╞═══════════════════════════════╡
│ {"x": 3.0, "y": 2.0, "": 1.0} │
└───────────────────────────────┘
定数項は変数名が存在しないため、特殊なキー""で表現します。
したがって{"x": 3.0, "y": 2.0, "": 1.0}は、$ 3 x + 2 y + 1 $を表します。
to_formula()を使うと数式で表現します。
print(df.select(Result=expr.to_formula()))
出力
shape: (1, 1)
┌─────────────────────────┐
│ Result │
│ --- │
│ str │
╞═════════════════════════╡
│ 3.0 * x + 2.0 * y + 1.0 │
└─────────────────────────┘
E1 + E2 * A + Bが3.0 * x + 2.0 * y + 1.0と表現できました。
まとめ
Polars 2.0.0で追加されたMap型を使うことで、0の係数を持たない疎な一次式をMapとして表現できます。
さらに、ネームスペースを利用して一次式用の演算を定義することで、Map型を使った一次式を通常の式のように計算できます。
今回の例では、E1 + E2 * A + Bを計算すると、{"x": 3.0, "y": 2.0, "": 1.0}というMapになり、to_formula()を使うことで3.0 * x + 2.0 * y + 1.0という数式として表現できました。
このように、PolarsのMap型とネームスペースを組み合わせることで、一次式を扱う独自のデータ表現を作ることができます。
参考：map_expr.pyの実装
map_expr.pyの実装は次の通りです。あくまで、スパースな一次式をDataFrameで扱うためのプロトタイプであることにご留意ください。
この実装では、Map同士の加減と、一次式のスカラー倍・スカラー除算を扱います。一般の二次式のような変数同士の乗算は扱いません。また、""を定数項に使用するため、変数名として空文字は使用できません。
また、この実装ではpl.Exprのprivateメソッドを使っています。この部分はPolarsのバージョンによって変更される可能性があります。
0除算をするとエラーではなくnullになります。
"""Mapを使った一次式の計算"""
import math
from numbers import Real
from typing import Any, Self
import polars as pl
MAP_EXPR_DTYPE = pl.Map(pl.String, pl.Float64)
def _to_dict_operand(val: Any) -> dict[str, float] | None:
"""行要素を一次式 dict に正規化する"""
if val is None:
return None
if isinstance(val, Real):
return {"": float(val)}
if isinstance(val, dict):
return {k: float(v) for k, v in val.items()}
return None
def _simplify(value: dict[str, float]) -> dict[str, float]:
return {k: v for k, v in value.items() if not math.isclose(v, 0.0, abs_tol=1e-9)}
def _add(row: dict[str, Any], *, is_add: bool = True) -> dict[str, float] | None:
left = _to_dict_operand(row["left"])
right = _to_dict_operand(row["right"])
if left is None or right is None:
return None
result = dict(left)
for key, value in right.items():
result[key] = result.get(key, 0.0) + float(value if is_add else -value)
return _simplify(result)
def _extract_scalar(val: Any) -> float | None:
"""スカラー値(実数または定数項のみのMap)を取り出す"""
if isinstance(val, Real):
return float(val)
if isinstance(val, dict) and tuple(val) in {(), ("",)}:
return float(val.get("", 0.0))
return None
def _mul(row: dict[str, Any], *, is_mul: bool = True) -> dict[str, float] | None:
left = row["left"]
right = row["right"]
if left is None or right is None:
return None
if is_mul:
scale = _extract_scalar(right)
vec_dict = _to_dict_operand(left)
if scale is None:
scale = _extract_scalar(left)
vec_dict = _to_dict_operand(right)
if scale is None or vec_dict is None:
raise ValueError("乗算の片方は定数(数値または定数項のみのMap)である必要があります")
return _simplify({k: v * scale for k, v in vec_dict.items()})
scale = _extract_scalar(right)
vec_dict = _to_dict_operand(left)
if scale is None or vec_dict is None:
raise ValueError("除算の右辺は定数である必要があります")
if math.isclose(scale, 0.0, abs_tol=1e-12):
return None # 0除算は null 扱い
return _simplify({k: v / scale for k, v in vec_dict.items()})
def _format_formula(value: dict[str, float] | None) -> str | None:
if value is None:
return None
# 変数はアルファベット順、定数項("")は最後
var_keys = sorted([k for k in value if k])
keys = var_keys + ([""] if "" in value else [])
terms = []
for key in keys:
coefficient = value[key]
if math.isclose(coefficient, 0.0, abs_tol=1e-9):
continue
magnitude = abs(coefficient)
if not key:
term = str(magnitude)
elif math.isclose(magnitude, 1.0, abs_tol=1e-9):
term = key
else:
term = f"{magnitude} * {key}"
terms.append((coefficient < 0, term))
if not terms:
return "0.0"
negative, first = terms[0]
formula = f"- {first}" if negative else first
for negative, term in terms[1:]:
formula += f" {'-' if negative else '+'} {term}"
return formula
class MapExpr(pl.Expr):
"""Map式をラップし、四則演算と数式変換を提供するExpr"""
def __new__(cls, expr: pl.Expr) -> Self:
obj = super().__new__(cls)
obj._pyexpr = expr._pyexpr # ruff: ignore[private-member-access]
return obj
@property
def expr(self) -> pl.Expr:
return pl.Expr._from_pyexpr(self._pyexpr) # ruff: ignore[private-member-access]
@classmethod
def _coerce(cls, value: Any) -> pl.Expr:
if isinstance(value, MapExpr):
return value.expr
if isinstance(value, MapExprNamespace):
return value.expr.expr
if isinstance(value, pl.Expr):
return value
if isinstance(value, Real):
return pl.lit({"": float(value)}, dtype=MAP_EXPR_DTYPE)
raise TypeError(f"未対応の型です: {type(value)}")
def __add__(self, other: Any) -> MapExpr:
expr = pl.struct(left=self.expr, right=self._coerce(other))
return MapExpr(expr.map_elements(_add, return_dtype=MAP_EXPR_DTYPE))
def __radd__(self, other: Any) -> MapExpr:
return self.__add__(other)
def __sub__(self, other: Any) -> MapExpr:
expr = pl.struct(left=self.expr, right=self._coerce(other))
return MapExpr(
expr.map_elements(lambda r: _add(r, is_add=False), return_dtype=MAP_EXPR_DTYPE),
)
def __rsub__(self, other: Any) -> MapExpr:
expr = pl.struct(left=self._coerce(other), right=self.expr)
return MapExpr(
expr.map_elements(lambda r: _add(r, is_add=False), return_dtype=MAP_EXPR_DTYPE),
)
def __mul__(self, other: Any) -> MapExpr:
expr = pl.struct(left=self.expr, right=self._coerce(other))
return MapExpr(expr.map_elements(_mul, return_dtype=MAP_EXPR_DTYPE))
def __rmul__(self, other: Any) -> MapExpr:
return self.__mul__(other)
def __truediv__(self, other: Any) -> MapExpr:
expr = pl.struct(left=self.expr, right=self._coerce(other))
return MapExpr(
expr.map_elements(lambda r: _mul(r, is_mul=False), return_dtype=MAP_EXPR_DTYPE),
)
def to_formula(self) -> pl.Expr:
return self.expr.map_elements(_format_formula, return_dtype=pl.String)
@pl.api.register_expr_namespace("lin")
class MapExprNamespace:
"""``pl.col("X").lin``のエントリポイント"""
def __init__(self, expr: pl.Expr) -> None:
self._map_expr = MapExpr(expr)
@property
def expr(self) -> MapExpr:
return self._map_expr
def __add__(self, other: Any) -> MapExpr:
return self._map_expr + other
def __radd__(self, other: Any) -> MapExpr:
return self._map_expr.__radd__(other)
def __sub__(self, other: Any) -> MapExpr:
return self._map_expr - other
def __rsub__(self, other: Any) -> MapExpr:
return self._map_expr.__rsub__(other)
def __mul__(self, other: Any) -> MapExpr:
return self._map_expr * other
def __rmul__(self, other: Any) -> MapExpr:
return self._map_expr.__rmul__(other)
def __truediv__(self, other: Any) -> MapExpr:
return self._map_expr / other
def to_formula(self) -> pl.Expr:
return self._map_expr.to_formula()
参考