Skip transformation of partition if partition is empty

ml6team · Mar 18, 2024 · 442501c · 442501c
1 parent c569c49
commit 442501c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/src/fondant/component/executor.py b/src/fondant/component/executor.py
@@ -492,18 +492,19 @@ def wrap_transform(
         """
 
         def wrapped_transform(dataframe: pd.DataFrame) -> pd.DataFrame:
-            # Call transform method
+            # 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 dataframe
+                return pd.DataFrame(columns=columns)
 
+            # Call transform method
             dataframe = transform(dataframe)
 
             # Drop columns not in specification
-            columns = [
-                name for name, field in operation_spec.operation_produces.items()
-            ]
-
             return dataframe[columns]
 
         return wrapped_transform

diff --git a/tests/component/test_component.py b/tests/component/test_component.py
@@ -617,5 +617,6 @@ def transform(dataframe: pd.DataFrame) -> pd.DataFrame:
             ),
         ),
     )
+
     output_df = wrapped_transform(input_df)
-    assert input_df.equals(output_df)
+    assert output_df.equals(pd.DataFrame())