ark-kun / pipeline_components Goto Github PK

from kfp.components import InputPath, OutputPath, create_component_from_func


def select_columns_using_Pandas_on_CSV_data(
    table_path: InputPath("CSV"),
    transformed_table_path: OutputPath("CSV"),
    column_names: list,
):
    """Selects columns from a data table.
    Args:
        table_path: Input data table.
        transformed_table_path: Transformed data table that only has the chosen columns.
        column_names: Names of the columns to select from the table.
    """
    import pandas

    df = pandas.read_csv(
        table_path,
        dtype="string",
    )
    print(df.head(5))
    print(InputPath)
    df = df[column_names]

    df.to_csv(transformed_table_path, index=False)



select_columns_using_Pandas_on_CSV_data_op = create_component_from_func(
        select_columns_using_Pandas_on_CSV_data,
        output_component_file="component.yaml",
        base_image="python:3.9",
        packages_to_install=["pandas==1.4.2",],
        annotations={
            "author": "Alexey Volkov <[email protected]>",
            "canonical_location": "https://raw.githubusercontent.com/Ark-kun/pipeline_components/master/components/pandas/Select_columns/in_CSV_format/component.yaml",
        },
    )
@kfp.dsl.pipeline(name='my-pipeline')
def my_pipeline(
    table_path: str,
    transform_code: str
):
    transform_dataframe_task = select_columns_using_Pandas_on_CSV_data_op('/home/jovyan/iris-1.csv',['setosa'])
    
# Compile and run the pipeline
kfp.compiler.Compiler().compile(my_pipeline, 'my-pipeline.tar.gz')