Start from dataset schema for lightweight python component consumes

examples/sample_pipeline/pipeline.py

@@ -50,13 +50,17 @@
     ],
 )
 class CalculateChunkLength(PandasTransformComponent):
+    def __init__(self, feature_name: str, **kwargs):
+        self.feature_name = feature_name
+
     def transform(self, dataframe: pd.DataFrame) -> pd.DataFrame:
-        dataframe["chunk_length"] = dataframe["text"].apply(len)
+        dataframe[self.feature_name] = dataframe["chunk"].apply(len)
         return dataframe
 
 
 _ = dataset.apply(
     ref=CalculateChunkLength,
-    consumes={"text": pa.string()},
+    consumes={"chunk": "text"},
     produces={"chunk_length": pa.int32()},
+    arguments={"feature_name": "chunk_length"},
 )

src/fondant/core/component_spec.py

@@ -55,6 +55,17 @@ def kubeflow_type(self) -> str:
         }
         return lookup[self.type]
 
+    def to_spec(self):
+        return {
+            k: v
+            for k, v in {
+                "description": self.description,
+                "type": self.type.__name__,
+                "default": self.default,
+            }.items()
+            if v is not None
+        }
+
 
 class ComponentSpec:
     """
@@ -83,8 +94,8 @@ def __init__(
         image: str,
         *,
         description: t.Optional[str] = None,
-        consumes: t.Optional[t.Dict[str, t.Union[str, pa.DataType]]] = None,
-        produces: t.Optional[t.Dict[str, t.Union[str, pa.DataType]]] = None,
+        consumes: t.Optional[t.Mapping[str, t.Union[str, pa.DataType]]] = None,
+        produces: t.Optional[t.Mapping[str, t.Union[str, pa.DataType]]] = None,
         previous_index: t.Optional[str] = None,
         args: t.Optional[t.Dict[str, t.Any]] = None,
         tags: t.Optional[t.List[str]] = None,
@@ -252,9 +263,9 @@ def args(self) -> t.Mapping[str, Argument]:
             {
                 name: Argument(
                     name=name,
-                    description=arg_info["description"],
+                    description=arg_info.get("description"),
                     type=pydoc.locate(arg_info["type"]),  # type: ignore
-                    default=arg_info["default"] if "default" in arg_info else None,
+                    default=arg_info.get("default"),
                     optional=arg_info.get("default") == "None",
                 )
                 for name, arg_info in self._specification.get("args", {}).items()

src/fondant/pipeline/lightweight_component.py

@@ -5,7 +5,7 @@
 from dataclasses import dataclass
 from functools import wraps
 
-from fondant.component import Component
+from fondant.component import BaseComponent, Component
 
 
 @dataclass
@@ -20,7 +20,7 @@ def __post_init__(self):
             self.base_image = "fondant:latest"
 
 
-class PythonComponent:
+class PythonComponent(BaseComponent):
     @classmethod
     def image(cls) -> Image:
         raise NotImplementedError
@@ -42,12 +42,12 @@ def wrapper(cls):
 
         # updated=() is needed to prevent an attempt to update the class's __dict__
         @wraps(cls, updated=())
-        class AppliedPythonComponent(cls, PythonComponent):
+        class PythonComponentOp(cls, PythonComponent):
             @classmethod
             def image(cls) -> Image:
                 return image
 
-        return AppliedPythonComponent
+        return PythonComponentOp
 
     return wrapper
 

src/fondant/pipeline/pipeline.py

@@ -22,6 +22,7 @@
 from fondant.core.manifest import Manifest
 from fondant.core.schema import Field
 from fondant.pipeline import Image, PythonComponent
+from fondant.pipeline.argument_inference import infer_arguments
 
 logger = logging.getLogger(__name__)
 
@@ -395,8 +396,8 @@ def read(
                 name,
                 image.base_image,  # TODO: revisit
                 description=description,
-                consumes={"additionalProperties": True},
                 produces={"additionalProperties": True},
+                args={k: v.to_spec() for k, v in infer_arguments(ref).items()},
             )
 
             operation = ComponentOp(
@@ -702,12 +703,18 @@ def apply(
             image = ref.image()
             description = ref.__doc__ or "python component"
 
+            consumes_spec = {k: v.type.to_json() for k, v in self.fields.items()}
+            if consumes:
+                for k, v in consumes.items():
+                    consumes_spec[k] = consumes_spec[v]
+
             component_spec = ComponentSpec(
                 name,
                 image.base_image,  # TODO: revisit
                 description=description,
-                consumes={"additionalProperties": True},
+                consumes=consumes_spec,
                 produces={"additionalProperties": True},
+                args={k: v.to_spec() for k, v in infer_arguments(ref).items()},
             )
 
             operation = ComponentOp(
@@ -777,12 +784,18 @@ def write(
             image = ref.image()
             description = ref.__doc__ or "python component"
 
+            consumes_spec = {k: v.type.to_json() for k, v in self.fields.items()}
+            if consumes:
+                for k, v in consumes.items():
+                    consumes_spec[k] = consumes_spec[v]
+
             component_spec = ComponentSpec(
                 name,
                 image.base_image,  # TODO: revisit
                 description=description,
-                consumes={"additionalProperties": True},
+                consumes=consumes_spec,
                 produces={"additionalProperties": True},
+                args={k: v.to_spec() for k, v in infer_arguments(ref).items()},
             )
 
             operation = ComponentOp(

tests/pipeline/test_python_component.py

@@ -5,6 +5,7 @@
 import pyarrow as pa
 from fondant.component import DaskLoadComponent, PandasTransformComponent
 from fondant.pipeline import Pipeline, lightweight_component
+from fondant.pipeline.compiler import DockerCompiler
 
 
 def test_build_python_script():
@@ -46,15 +47,17 @@ def load(self) -> dd.DataFrame:
     )
 
 
-def test_lightweight_component(tmp_path_factory):
+def test_compile_lightweight_component(tmp_path_factory):
     pipeline = Pipeline(
         name="dummy-pipeline",
         base_path="./data",
     )
 
     @lightweight_component(
-        base_image="python:3.8-slim-buster",
-        extra_requires=["pandas", "dask"],
+        base_image="python:3.8",
+        extra_requires=[
+            "fondant[component]@git+https://github.com/ml6team/fondant@main",
+        ],
     )
     class CreateData(DaskLoadComponent):
         def load(self) -> dd.DataFrame:
@@ -72,9 +75,14 @@ def load(self) -> dd.DataFrame:
         produces={"x": pa.int32(), "y": pa.int32()},
     )
 
-    @lightweight_component()
+    @lightweight_component(
+        base_image="python:3.8",
+        extra_requires=[
+            "fondant[component]@git+https://github.com/ml6team/fondant@main",
+        ],
+    )
     class AddN(PandasTransformComponent):
-        def __init__(self, n: int):
+        def __init__(self, n: int, **kwargs):
             self.n = n
 
         def transform(self, dataframe: pd.DataFrame) -> pd.DataFrame:
@@ -85,4 +93,59 @@ def transform(self, dataframe: pd.DataFrame) -> pd.DataFrame:
         ref=AddN,
         produces={"x": pa.int32(), "y": pa.int32()},
         consumes={"x": pa.int32(), "y": pa.int32()},
+        arguments={"n": 1},
     )
+
+    DockerCompiler().compile(pipeline)
+
+
+def test_consumes_mapping():
+    pipeline = Pipeline(
+        name="dummy-pipeline",
+        base_path="./data",
+    )
+
+    @lightweight_component(
+        base_image="python:3.8",
+        extra_requires=[
+            "fondant[component]@git+https://github.com/ml6team/fondant@main",
+        ],
+    )
+    class CreateData(DaskLoadComponent):
+        def load(self) -> dd.DataFrame:
+            df = pd.DataFrame(
+                {
+                    "x": [1, 2, 3],
+                    "y": [4, 5, 6],
+                },
+                index=pd.Index(["a", "b", "c"], name="id"),
+            )
+            return dd.from_pandas(df, npartitions=1)
+
+    dataset = pipeline.read(
+        ref=CreateData,
+        produces={"x": pa.int32(), "y": pa.int32()},
+    )
+
+    @lightweight_component(
+        base_image="python:3.8",
+        extra_requires=[
+            "fondant[component]@git+https://github.com/ml6team/fondant@main",
+        ],
+    )
+    class AddN(PandasTransformComponent):
+        def __init__(self, n: int, **kwargs):
+            self.n = n
+
+        def transform(self, dataframe: pd.DataFrame) -> pd.DataFrame:
+            dataframe["a"] = dataframe["a"].map(lambda x: x + self.n)
+            return dataframe
+
+    _ = dataset.apply(
+        ref=AddN,
+        consumes={"a": "x"},
+        produces={"a": pa.int32()},
+        arguments={"n": 1},
+    )
+
+    DockerCompiler().compile(pipeline)