Skip transformation of partition if partition is empty

src/fondant/component/executor.py

@@ -479,6 +479,7 @@ def wrap_transform(
         operation_spec: OperationSpec,
     ) -> t.Callable:
         """Factory that creates a function to wrap the component transform function. The wrapper:
+        - Skips the transformation if the received partition is empty
         - Removes extra columns from the returned dataframe which are not defined in the component
           spec `produces` section
         - Sorts the columns from the returned dataframe according to the order in the component
@@ -491,14 +492,19 @@ def wrap_transform(
         """
 
         def wrapped_transform(dataframe: pd.DataFrame) -> pd.DataFrame:
-            # Call transform method
-            dataframe = transform(dataframe)
-
-            # Drop columns not in specification
+            # Columns of operation specification
             columns = [
                 name for name, field in operation_spec.operation_produces.items()
             ]
 
+            if dataframe.empty:
+                logger.info("Received empty partition, skipping transformation.")
+                return pd.DataFrame(columns=columns)
+
+            # Call transform method
+            dataframe = transform(dataframe)
+
+            # Drop columns not in specification
             return dataframe[columns]
 
         return wrapped_transform

tests/component/test_component.py

@@ -591,3 +591,32 @@ def write(self, dataframe):
     with mock.patch.object(MyWriteComponent, "write", write):
         executor.execute(MyWriteComponent)
         write.mock.assert_called_once()
+
+
+def test_skipping_empty_partition():
+    # Create an empty dataframe to simulate empty partitions
+    input_df = pd.DataFrame.from_dict(
+        {
+            "image_height": [],
+            "image_width": [],
+            "caption_text": [],
+        },
+    )
+
+    def transform(dataframe: pd.DataFrame) -> pd.DataFrame:
+        msg = "This should not be called"
+        raise ValueError(msg)
+
+    wrapped_transform = PandasTransformExecutor.wrap_transform(
+        transform,
+        operation_spec=OperationSpec(
+            ComponentSpec(
+                name="dummy-spec",
+                image="dummy-image",
+                description="dummy-description",
+            ),
+        ),
+    )
+
+    output_df = wrapped_transform(input_df)
+    assert output_df.equals(pd.DataFrame())