[AIP-5284] Integration with Transforms

Author: smatzana

datasets/datasets_decorator.py

@@ -2,11 +2,13 @@
 import keyword
 from typing import Callable, Optional
 
-from datasets import DatasetPlugin
 from datasets.context import Context
 from datasets.utils import _pascal_to_snake_case
+from datasets.txf_integration.txf_utils import add_txf_attributes
+from .dataset_plugin import DatasetPlugin
 
 
+# flake8: noqa: C901
 def dataset(
     name: str = None,
     field_name: Optional[str] = None,
@@ -27,6 +29,19 @@ def step_wrapper(*args, **kwargs):
                 _snake_name = _pascal_to_snake_case(dataset.name)
                 setattr(self, _snake_name, dataset)
 
+            # Transformation integration
+            if type(self).__name__ in ["FlowSpec", "OnlineFlowSpec"]:
+                add_txf_attributes(self)
+                if not self._txf_registered_datasets.get(self.name, None):
+                    self._txf_registered_datasets[self.name] = [dataset]
+                else:
+                    self._txf_registered_datasets[self.name].append(dataset)
+
+                if len(self._txf_callbacks):
+                    for _, callback in self._txf_callbacks.items():
+                        callback()
+            # End Transformation integration
+
             func(*args, **kwargs)
 
         return step_wrapper

datasets/tests/resources/txf_test_flow.py

@@ -0,0 +1,40 @@
+import pandas as pd
+from metaflow import FlowSpec, step
+
+from datasets import Mode, dataset
+from datasets.txf_integration.txf_utils import (
+    TXF_REGISTERED_DATASETS_ATTRIBUTE,
+    TXF_CALLBACKS_ATTRIBUTE,
+    TXF_METADATA_ATTRIBUTE,
+    add_callback,
+)
+
+
+def udf():
+    pass
+
+
+class TxTestFlow(FlowSpec):
+    @step
+    def start(self):
+        add_callback(self, "udf", udf)
+        self.next(self.ds)
+
+    @dataset(name="TxfDataset", partition_by="region", mode=Mode.WRITE)
+    @step
+    def ds(self):
+        df = pd.DataFrame({"region": ["A", "A", "A", "B", "B", "B"], "home_id": [1, 2, 3, 4, 5, 6]})
+        self.txf_dataset.write(df)
+        assert hasattr(self, TXF_REGISTERED_DATASETS_ATTRIBUTE)
+        assert hasattr(self, TXF_CALLBACKS_ATTRIBUTE)
+        assert hasattr(self, TXF_METADATA_ATTRIBUTE)
+        assert len(getattr(self, TXF_CALLBACKS_ATTRIBUTE)) == 1
+        self.next(self.end)
+
+    @step
+    def end(self):
+        pass
+
+
+if __name__ == "__main__":
+    TxTestFlow()

datasets/tests/test_txf.py

@@ -0,0 +1,5 @@
+from datasets.tests.test_tutorials import run_flow
+
+
+def test_input_output_flow():
+    run_flow("tests/resources/txf_test_flow.py")

datasets/txf_integration/__init__.py

@@ -0,0 +1,19 @@
+from typing import Callable
+
+TXF_REGISTERED_DATASETS_ATTRIBUTE = "_txf_registered_datasets"
+TXF_CALLBACKS_ATTRIBUTE = "_txf_callbacks"
+TXF_METADATA_ATTRIBUTE = "_txf_metadata_accumulator"
+
+
+def add_txf_attributes(flow):
+    if not hasattr(flow, TXF_REGISTERED_DATASETS_ATTRIBUTE):
+        setattr(flow, TXF_REGISTERED_DATASETS_ATTRIBUTE, dict())
+    if not hasattr(flow, TXF_CALLBACKS_ATTRIBUTE):
+        setattr(flow, TXF_CALLBACKS_ATTRIBUTE, {})
+    if not hasattr(flow, TXF_METADATA_ATTRIBUTE):
+        setattr(flow, TXF_METADATA_ATTRIBUTE, {})
+
+
+def add_callback(flow, caller: str, func: Callable):
+    add_txf_attributes(flow)
+    flow._txf_callbacks[caller] = func