Repository: refuel-ai/autolabel

Issue:
[Bug]: Send only needed columns in transforms
**Describe the bug**
Currently, in the webpage transform, even if a subset of columns are specified in the output columns, all columns are sent to the final dataset object.

**To Reproduce**
Run the example_webpage_transform and look at the ds.df

**Expected behavior**
We should just return the columns which are present in the output_columns dictionary.



Discussion:

Resolving patch:
diff --git a/src/autolabel/transforms/base.py b/src/autolabel/transforms/base.py
--- a/src/autolabel/transforms/base.py
+++ b/src/autolabel/transforms/base.py
@@ -29,14 +29,13 @@ def name() -> str:
         pass
 
     @property
-    @abstractmethod
     def output_columns(self) -> Dict[str, Any]:
         """
         Returns a dictionary of output columns. The keys are the names of the output columns
         as expected by the transform. The values are the column names they should be mapped to in
         the dataset.
         """
-        return self._output_columns
+        return {k: self._output_columns.get(k, None) for k in self.COLUMN_NAMES}
 
     @abstractmethod
     async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:
@@ -79,3 +78,15 @@ async def apply(self, row: Dict[str, Any]) -> Dict[str, Any]:
             self.cache.update(cache_entry)
 
         return output
+
+    def _return_output_row(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Returns the output row with the correct column names.
+        Args:
+            row: The output row.
+        Returns:
+            The output row with the correct column names.
+        """
+        # remove null key
+        row.pop(None, None)
+        return row
diff --git a/src/autolabel/transforms/pdf.py b/src/autolabel/transforms/pdf.py
--- a/src/autolabel/transforms/pdf.py
+++ b/src/autolabel/transforms/pdf.py
@@ -6,6 +6,11 @@
 
 
 class PDFTransform(BaseTransform):
+    COLUMN_NAMES = [
+        "content_column",
+        "metadata_column",
+    ]
+
     def __init__(
         self,
         cache: BaseCache,
@@ -52,14 +57,6 @@ def __init__(
     def name() -> str:
         return TransformType.PDF
 
-    @property
-    def output_columns(self) -> Dict[str, Any]:
-        COLUMN_NAMES = [
-            "content_column",
-            "metadata_column",
-        ]
-        return {k: self._output_columns.get(k, k) for k in COLUMN_NAMES}
-
     def get_page_texts(self, row: Dict[str, Any]) -> List[str]:
         """This function gets the text from each page of a PDF file.
         If OCR is enabled, it uses the pdf2image library to convert the PDF into images and then uses
@@ -97,7 +94,7 @@ async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:
             self.output_columns["content_column"]: output,
             self.output_columns["metadata_column"]: {"num_pages": len(texts)},
         }
-        return transformed_row
+        return self._return_output_row(transformed_row)
 
     def params(self):
         return {
diff --git a/src/autolabel/transforms/webpage_transform.py b/src/autolabel/transforms/webpage_transform.py
--- a/src/autolabel/transforms/webpage_transform.py
+++ b/src/autolabel/transforms/webpage_transform.py
@@ -19,6 +19,13 @@
 
 
 class WebpageTransform(BaseTransform):
+    COLUMN_NAMES = [
+        "content_column",
+        "content_in_bytes_column",
+        "soup_column",
+        "metadata_column",
+    ]
+
     def __init__(
         self,
         cache: BaseCache,
@@ -63,16 +70,6 @@ def __init__(
     def name(self) -> str:
         return TransformType.WEBPAGE_TRANSFORM
 
-    @property
-    def output_columns(self) -> Dict[str, Any]:
-        COLUMN_NAMES = [
-            "content_column",
-            "content_in_bytes_column",
-            "soup_column",
-            "metadata_column",
-        ]
-        return {k: self._output_columns.get(k, k) for k in COLUMN_NAMES}
-
     def _load_metadata(self, url, soup) -> Dict[str, Any]:
         metadata = {"url": url}
         if soup.find("title"):
@@ -138,7 +135,7 @@ async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:
             self.output_columns["metadata_column"]: url_response_data.get("metadata"),
         }
 
-        return transformed_row
+        return self._return_output_row(transformed_row)
 
     def params(self):
         return {


Test patch:
diff --git a/tests/unit/transforms/test_webpage_transform.py b/tests/unit/transforms/test_webpage_transform.py
--- a/tests/unit/transforms/test_webpage_transform.py
+++ b/tests/unit/transforms/test_webpage_transform.py
@@ -21,9 +21,7 @@ async def test_webpage_transform():
     # Transform the row
     transformed_row = await transform.apply(row)
     # Check the output
-    assert set(transformed_row.keys()) == set(
-        ["webpage_content", "metadata", "content_in_bytes_column", "soup_column"]
-    )
+    assert set(transformed_row.keys()) == set(["webpage_content", "metadata"])
     assert isinstance(transformed_row["webpage_content"], str)
     assert isinstance(transformed_row["metadata"], dict)
     assert len(transformed_row["webpage_content"]) > 0