34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
from __future__ import annotations
|
|
from typing import Dict
|
|
import polars as pl
|
|
from tools.base import BaseTool
|
|
|
|
|
|
class FormulaTool(BaseTool):
|
|
def execute(self, inputs: Dict[str, pl.DataFrame]) -> Dict[str, pl.DataFrame]:
|
|
df = inputs.get("Input", pl.DataFrame())
|
|
|
|
if self.config is None:
|
|
return {"Output": df}
|
|
|
|
for ff in self.config.findall("FormulaFields/FormulaField"):
|
|
expr_text = ff.attrib.get("expression", "")
|
|
field = ff.attrib.get("field", "")
|
|
alteryx_type = ff.attrib.get("type", "V_WString")
|
|
size = ff.attrib.get("size")
|
|
dtype = self.ctx.type_mapper.map(alteryx_type, size)
|
|
|
|
if not field:
|
|
continue
|
|
|
|
try:
|
|
series = self.ctx.transpiler.eval_series(df, expr_text, field, dtype)
|
|
except Exception as e:
|
|
raise RuntimeError(
|
|
f"Formula field {field!r} expression {expr_text!r} failed: {e}"
|
|
) from e
|
|
|
|
df = df.with_columns(series.alias(field))
|
|
|
|
return {"Output": df}
|