Start from dataset schema for lightweight python component consumes

examples/sample_pipeline/pipeline.py

@@ -59,7 +59,6 @@ def transform(self, dataframe: pd.DataFrame) -> pd.DataFrame:
 
 _ = dataset.apply(
     ref=CalculateChunkLength,
-    consumes={"text": pa.string()},
     produces={"chunk_length": pa.int32()},
     arguments={"arg_x": "value_x"},
 )

src/fondant/core/component_spec.py

@@ -94,8 +94,8 @@ def __init__(
         image: str,
         *,
         description: t.Optional[str] = None,
-        consumes: t.Optional[t.Dict[str, t.Union[str, pa.DataType, bool]]] = None,
-        produces: t.Optional[t.Dict[str, t.Union[str, pa.DataType, bool]]] = None,
+        consumes: t.Optional[t.Mapping[str, t.Union[str, pa.DataType, bool]]] = None,
+        produces: t.Optional[t.Mapping[str, t.Union[str, pa.DataType, bool]]] = None,
         previous_index: t.Optional[str] = None,
         args: t.Optional[t.Dict[str, t.Any]] = None,
         tags: t.Optional[t.List[str]] = None,
@@ -462,7 +462,7 @@ def _inner_mapping(self, name: str) -> t.Mapping[str, Field]:
                 else:
                     msg = (
                         f"Received pyarrow DataType value {value} for key {key} in the "
-                        f"`{name}` argument passed to the operation, but {key} is "
+                        f"`{name}` argument passed to the operation, but `{key}` is "
                         f"already defined in the `{name}` section of the component spec "
                         f"with type {spec_type}"
                     )
@@ -494,8 +494,8 @@ def _outer_mapping(self, name: str) -> t.Mapping[str, Field]:
                 mapping[value] = Field(name=value, type=mapping.pop(key).type)
             else:
                 msg = (
-                    f"Received a string value for key {key} in the `{name}` "
-                    f"argument passed to the operation, but {key} is not defined in "
+                    f"Received a string value for key `{key}` in the `{name}` "
+                    f"argument passed to the operation, but `{key}` is not defined in "
                     f"the `{name}` section of the component spec."
                 )
                 raise InvalidPipelineDefinition(msg)

src/fondant/pipeline/lightweight_component.py

@@ -1,11 +1,15 @@
 import inspect
 import itertools
+import logging
 import textwrap
 import typing as t
 from dataclasses import asdict, dataclass
 from functools import wraps
 
 from fondant.component import BaseComponent, Component
+from fondant.core.schema import Field, Type
+
+logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -28,11 +32,42 @@ class PythonComponent(BaseComponent):
     def image(cls) -> Image:
         raise NotImplementedError
 
+    @classmethod
+    def consumes(cls) -> t.Optional[t.Dict[str, t.Any]]:
+        pass
+
+    @classmethod
+    def get_consumes_spec(
+        cls,
+        dataset_fields: t.Mapping[str, Field],
+    ):
+        consumes = cls.consumes()
+
+        if consumes is None:
+            # Get consumes spec from the dataset
+            spec_consumes = {k: v.type.to_dict() for k, v in dataset_fields.items()}
+
+            logger.warning(
+                "No consumes defined. Consumes will be inferred from the dataset."
+                "All field will be consumed which may lead to additional computation,"
+                "Consider defining consumes in the component.\n Consumes: %s",
+                spec_consumes,
+            )
+
+        elif consumes == {"additionalProperties": True}:
+            spec_consumes = consumes
+
+        else:
+            spec_consumes = {k: Type(v).to_dict() for k, v in consumes.items()}
+
+        return spec_consumes
+
 
 def lightweight_component(
     *args,
     extra_requires: t.Optional[t.List[str]] = None,
     base_image: t.Optional[str] = None,
+    consumes: t.Optional[t.Dict[str, t.Any]] = None,
 ):
     """Decorator to enable a python component."""
 
@@ -121,6 +156,10 @@ class PythonComponentOp(cls, PythonComponent):
             def image(cls) -> Image:
                 return image
 
+            @classmethod
+            def consumes(cls) -> t.Optional[t.Dict[str, t.Dict[t.Any, t.Any]]]:
+                return consumes
+
         return PythonComponentOp
 
     # Call wrapper with function (`args[0]`) when no additional arguments were passed

src/fondant/pipeline/pipeline.py

@@ -206,22 +206,39 @@ def from_component_yaml(cls, path, **kwargs) -> "ComponentOp":
         )
 
     @classmethod
-    def from_ref(cls, ref: t.Any, **kwargs) -> "ComponentOp":
+    def from_ref(
+        cls,
+        ref: t.Any,
+        fields: t.Optional[t.Mapping[str, Field]] = None,
+        **kwargs,
+    ) -> "ComponentOp":
         """Create a ComponentOp from a reference. The reference can
         be a reusable component name, a path to a custom component,
         or a python component class.
+
+        Args:
+            ref: The name of a reusable component, or the path to the directory containing
+                a custom component, or a python component class.
+            fields: The fields of the dataset available to the component.
+            **kwargs: The provided user arguments are passed in as keyword arguments
         """
         if inspect.isclass(ref) and issubclass(ref, BaseComponent):
             if issubclass(ref, PythonComponent):
                 name = ref.__name__
                 image = ref.image()
                 description = ref.__doc__ or "python component"
 
+                consumes_spec = (
+                    ref.get_consumes_spec(fields)
+                    if fields
+                    else {"additionalProperties": True}
+                )
+
                 component_spec = ComponentSpec(
                     name,
                     image.base_image,
                     description=description,
-                    consumes={"additionalProperties": True},
+                    consumes=consumes_spec,
                     produces={"additionalProperties": True},
                     args={
                         name: arg.to_spec()
@@ -724,6 +741,7 @@ def apply(
         """
         operation = ComponentOp.from_ref(
             ref,
+            fields=self.fields,
             produces=produces,
             consumes=consumes,
             arguments=arguments,
@@ -771,6 +789,7 @@ def write(
         """
         operation = ComponentOp.from_ref(
             ref,
+            fields=self.fields,
             consumes=consumes,
             arguments=arguments,
             input_partition_rows=input_partition_rows,