38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
from __future__ import annotations
|
|
from typing import Dict
|
|
import polars as pl
|
|
from tools.base import BaseTool
|
|
|
|
|
|
class UnionTool(BaseTool):
|
|
def execute(self, inputs: Dict[str, pl.DataFrame]) -> Dict[str, pl.DataFrame]:
|
|
if not inputs:
|
|
return {"Output": pl.DataFrame()}
|
|
|
|
mode = self._cfg("Mode", "Auto") or "Auto"
|
|
|
|
# Sort inputs: Input, Input1, Input2... (Alteryx uses 1-based suffixes)
|
|
def _sort_key(k: str) -> int:
|
|
suffix = k.replace("Input", "").lstrip("#")
|
|
return int(suffix) if suffix.isdigit() else 0
|
|
|
|
sorted_keys = sorted(inputs.keys(), key=_sort_key)
|
|
dfs = [inputs[k] for k in sorted_keys if inputs[k] is not None]
|
|
|
|
if not dfs:
|
|
return {"Output": pl.DataFrame()}
|
|
|
|
if mode == "ByPosition":
|
|
names = dfs[0].columns
|
|
renamed = []
|
|
for df in dfs:
|
|
if len(df.columns) == len(names):
|
|
renamed.append(df.rename(dict(zip(df.columns, names))))
|
|
else:
|
|
renamed.append(df)
|
|
result = pl.concat(renamed, how="diagonal_relaxed")
|
|
else:
|
|
result = pl.concat(dfs, how="diagonal_relaxed")
|
|
|
|
return {"Output": result}
|