from ppc_robot_lib.steps import AbstractStep
from ppc_robot_lib.tasks import TaskContextInterface, StepPerformance
[docs]
class SuffixColumnsStep(AbstractStep):
"""
Suffixes column names with the given ``suffix``. If a column name is listed in ``except_cols``, its name
is kept as-is.
**Example:**
>>> from ppc_robot_lib.steps.transformations import SuffixColumnsStep
>>> SuffixColumnsStep('performance_yday', '_yday', except_cols=['Id', 'CampaignName'])
"""
def __init__(self, input_table: str, suffix: str, except_cols: list[str] = None, output_table: str = None):
"""
:param input_table: Input table.
:param suffix: Suffix to add to column names.
:param except_cols: List of collumn names that will not be suffixed.
:param output_table: Output table. If none is given, input table is used.
"""
self.input_table = input_table
self.suffix = suffix
if except_cols is not None:
self.except_cols = except_cols
else:
self.except_cols = []
if output_table is not None:
self.output_table = output_table
else:
self.output_table = input_table
def execute(self, task_ctx: TaskContextInterface) -> StepPerformance:
table = task_ctx.work_set.get_table(self.input_table)
rename = {}
for column in table.columns:
if column not in self.except_cols:
rename[column] = column + self.suffix
inplace = self.input_table == self.output_table
new_table = table.rename(columns=rename, inplace=inplace)
if not inplace:
if self.output_table in task_ctx.work_set:
task_ctx.work_set.delete_table(self.output_table)
task_ctx.work_set.set_table(self.output_table, new_table)
else:
new_table = table
rows = len(new_table.index)
return StepPerformance(new_table, rows_in=rows, rows_out=rows)