25 lines
879 B
Python
25 lines
879 B
Python
from __future__ import annotations
|
|
from typing import Dict
|
|
import polars as pl
|
|
from tools.base import BaseTool
|
|
|
|
|
|
class RecordIDTool(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}
|
|
|
|
field = self._cfg("Field", "RecordID") or "RecordID"
|
|
start_str = self._cfg("StartValue", "1") or "1"
|
|
start = int(start_str)
|
|
alteryx_type = self._cfg("FieldType", "Int32") or "Int32"
|
|
dtype = self.ctx.type_mapper.map(alteryx_type)
|
|
|
|
df = df.with_row_index(name=field, offset=start)
|
|
df = df.with_columns(pl.col(field).cast(dtype))
|
|
|
|
# Move the ID column to first position
|
|
cols = [field] + [c for c in df.columns if c != field]
|
|
return {"Output": df.select(cols)}
|