Implementation of new pipeline interface

src/fondant/component/component.py

@@ -16,7 +16,12 @@ class BaseComponent:
         **kwargs: The provided user arguments are passed in as keyword arguments
     """
 
-    def __init__(self, spec: ComponentSpec, **kwargs):
+    def __init__(
+        self,
+        spec: ComponentSpec,
+        schema: t.Optional[t.Dict[str, str]] = None,
+        **kwargs,
+    ):
         pass
 
 

src/fondant/component/data_io.py

@@ -28,9 +28,11 @@ def __init__(
         manifest: Manifest,
         component_spec: ComponentSpec,
         input_partition_rows: t.Optional[int] = None,
+        consumes: t.Optional[t.Dict[str, str]] = None,
     ):
         super().__init__(manifest=manifest, component_spec=component_spec)
         self.input_partition_rows = input_partition_rows
+        self.consumes = consumes
 
     def partition_loaded_dataframe(self, dataframe: dd.DataFrame) -> dd.DataFrame:
         """
@@ -131,6 +133,10 @@ def load_dataframe(self) -> dd.DataFrame:
 
         logging.info(f"Columns of dataframe: {list(dataframe.columns)}")
 
+        # Renaming dataframe columns
+        if self.consumes:
+            dataframe = dataframe.rename(columns=self.consumes)
+
         return dataframe
 
 
@@ -140,8 +146,10 @@ def __init__(
         *,
         manifest: Manifest,
         component_spec: ComponentSpec,
+        produces: t.Optional[t.Dict[str, str]] = None,
     ):
         super().__init__(manifest=manifest, component_spec=component_spec)
+        self.produces = produces
 
     def write_dataframe(
         self,
@@ -158,6 +166,11 @@ def write_dataframe(
         self.validate_dataframe_columns(dataframe, columns_to_produce)
 
         dataframe = dataframe[columns_to_produce]
+
+        # Renaming dataframe produces
+        if self.produces:
+            dataframe = dataframe.rename(columns=self.produces)
+
         write_task = self._write_dataframe(dataframe)
 
         with ProgressBar():

src/fondant/component/executor.py

@@ -67,6 +67,9 @@ def __init__(
         input_partition_rows: int,
         cluster_type: t.Optional[str] = None,
         client_kwargs: t.Optional[dict] = None,
+        schema: t.Optional[t.Dict[str, str]] = None,
+        consumes: t.Optional[t.Dict[str, str]] = None,
+        produces: t.Optional[t.Dict[str, str]] = None,
     ) -> None:
         self.spec = spec
         self.cache = cache
@@ -75,6 +78,10 @@ def __init__(
         self.metadata = Metadata.from_dict(metadata)
         self.user_arguments = user_arguments
         self.input_partition_rows = input_partition_rows
+        self.schema = schema
+
+        self.consumes = consumes
+        self.produces = produces
 
         if cluster_type == "local":
             client_kwargs = client_kwargs or {
@@ -112,6 +119,10 @@ def from_args(cls) -> "Executor":
         parser.add_argument("--input_partition_rows", type=int)
         parser.add_argument("--cluster_type", type=str)
         parser.add_argument("--client_kwargs", type=json.loads)
+        parser.add_argument("--schema", type=json.loads)
+        parser.add_argument("--consumes", type=json.loads)
+        parser.add_argument("--produces", type=json.loads)
+
         args, _ = parser.parse_known_args()
 
         if "component_spec" not in args:
@@ -123,13 +134,19 @@ def from_args(cls) -> "Executor":
         cache = args.cache
         cluster_type = args.cluster_type
         client_kwargs = args.client_kwargs
+        schema = args.schema
+        consumes = args.consumes
+        produces = args.produces
 
         return cls.from_spec(
             component_spec,
             cache=cache,
             input_partition_rows=input_partition_rows,
             cluster_type=cluster_type,
             client_kwargs=client_kwargs,
+            schema=schema,
+            consumes=consumes,
+            produces=produces,
         )
 
     @classmethod
@@ -141,6 +158,9 @@ def from_spec(
         input_partition_rows: int,
         cluster_type: t.Optional[str],
         client_kwargs: t.Optional[dict],
+        schema: t.Optional[dict],
+        consumes: t.Optional[dict],
+        produces: t.Optional[dict],
     ) -> "Executor":
         """Create an executor from a component spec."""
         args_dict = vars(cls._add_and_parse_args(component_spec))
@@ -160,6 +180,15 @@ def from_spec(
         if "client_kwargs" in args_dict:
             args_dict.pop("client_kwargs")
 
+        if "schema" in args_dict:
+            args_dict.pop("schema")
+
+        if "consumes" in args_dict:
+            args_dict.pop("consumes")
+
+        if "produces" in args_dict:
+            args_dict.pop("produces")
+
         input_manifest_path = args_dict.pop("input_manifest_path")
         output_manifest_path = args_dict.pop("output_manifest_path")
         metadata = args_dict.pop("metadata")
@@ -175,6 +204,9 @@ def from_spec(
             input_partition_rows=input_partition_rows,
             cluster_type=cluster_type,
             client_kwargs=client_kwargs,
+            schema=schema,
+            consumes=consumes,
+            produces=produces,
         )
 
     @classmethod
@@ -251,11 +283,18 @@ def _execute_component(
             A Dask DataFrame containing the output data
         """
 
-    def _write_data(self, dataframe: dd.DataFrame, *, manifest: Manifest):
+    def _write_data(
+        self,
+        dataframe: dd.DataFrame,
+        *,
+        manifest: Manifest,
+        produces: t.Optional[t.Dict[str, str]],
+    ):
         """Create a data writer given a manifest and writes out the index and subsets."""
         data_writer = DaskDataWriter(
             manifest=manifest,
             component_spec=self.spec,
+            produces=produces,
         )
 
         data_writer.write_dataframe(dataframe, self.client)
@@ -331,7 +370,7 @@ def _run_execution(
         input_manifest: Manifest,
     ) -> Manifest:
         logging.info("Executing component")
-        component = component_cls(self.spec, **self.user_arguments)
+        component = component_cls(self.spec, self.schema, **self.user_arguments)
         output_df = self._execute_component(
             component,
             manifest=input_manifest,
@@ -340,7 +379,11 @@ def _run_execution(
             component_spec=self.spec,
             run_id=self.metadata.run_id,
         )
-        self._write_data(dataframe=output_df, manifest=output_manifest)
+        self._write_data(
+            dataframe=output_df,
+            manifest=output_manifest,
+            produces=self.produces,
+        )
 
         return output_manifest
 
@@ -478,6 +521,7 @@ def _execute_component(
             manifest=manifest,
             component_spec=self.spec,
             input_partition_rows=self.input_partition_rows,
+            consumes=self.consumes,
         )
         dataframe = data_loader.load_dataframe()
         return component.transform(dataframe)
@@ -530,6 +574,7 @@ def _execute_component(
             manifest=manifest,
             component_spec=self.spec,
             input_partition_rows=self.input_partition_rows,
+            consumes=self.consumes,
         )
         dataframe = data_loader.load_dataframe()
 
@@ -574,6 +619,7 @@ def _execute_component(
             manifest=manifest,
             component_spec=self.spec,
             input_partition_rows=self.input_partition_rows,
+            consumes=self.consumes,
         )
         dataframe = data_loader.load_dataframe()
         component.write(dataframe)