from ppc_robot_lib.steps import AbstractStep
from ppc_robot_lib.tasks import TaskContextInterface, StepPerformance
[docs]
class ProjectionStep(AbstractStep):
"""
Projection creates a new table containing only the given columns, in the given order.
**Example:**
Let's assume we have a table with columns A, B, C, D, in that order. By executing::
>>> from ppc_robot_lib.steps.transformations import ProjectionStep
>>> ProjectionStep("in_table", columns=['C', 'B', 'D'], output_table="out_table")
we would get a table with columns C, B and D. Column A would be dropped.
"""
def __init__(self, input_table: str, columns: list[str], output_table: str = None):
"""
:param input_table: Input table.
:param columns: List of columns to transfer to the new table.
:param output_table: Output table. If no table is given, input table is used.
"""
self.input_table = input_table
self.columns = columns
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)
new_table = table[self.columns]
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)
rows = len(new_table.index)
return StepPerformance(new_table, rows_in=rows, rows_out=rows)