with_columns.py

from hamilton.plugins.h_spark import with_columns

import pyspark.sql as ps
import map_transforms # file defined above

@with_columns(
    map_transforms, 
    columns_to_pass=["column_1_from_dataframe", "column_2_from_dataframe"]
)
def final_result(all_initial_data: ps.DataFrame) -> ps.DataFrame:
    """Gives the final result. This decorator will apply the transformations in the order specified in the DAG.
    Then, the final_result function is called, with the result of the transformations passed in."""
    return _process(all_initial_data)