Enable Workflow.transform to be run with a DataFrame type

nvtabular/workflow/workflow.py

@@ -21,6 +21,7 @@
 import time
 import types
 import warnings
+from functools import singledispatchmethod
 from typing import TYPE_CHECKING, Optional
 
 import cloudpickle
@@ -78,7 +79,8 @@ def __init__(self, output_node: WorkflowNode, client: Optional["distributed.Clie
         self.graph = Graph(output_node)
         self.executor = DaskExecutor(client)
 
-    def transform(self, dataset: Dataset) -> Dataset:
+    @singledispatchmethod
+    def transform(self, data):
         """Transforms the dataset by applying the graph of operators to it. Requires the ``fit``
         method to have already been called, or calculated statistics to be loaded from disk
 
@@ -88,16 +90,33 @@ def transform(self, dataset: Dataset) -> Dataset:
 
         Parameters
         -----------
-        dataset: Dataset
-            Input dataset to transform
+        data: Union[Dataset, DataFrameType]
+            Input dataset or dataframe to transform
 
         Returns
         -------
-        Dataset
-            Transformed Dataset with the workflow graph applied to it
+        Dataset or DataFrame
+            Transformed Dataset or DataFrame with the workflow graph applied to it
         """
+        raise NotImplementedError(
+            "Workflow.transform received an unsupported type: {type(dataset)} "
+            "Supported types are a `merlin.io.Dataset` or DataFrame (pandas or cudf)"
+        )
+
+    @transform.register
+    def _(self, dataset: Dataset) -> Dataset:
         return self._transform_impl(dataset)
 
+    @transform.register
+    def _(self, dataframe: pd.DataFrame) -> pd.DataFrame:
+        return self._transform_df(dataframe)
+
+    if cudf:
+
+        @transform.register
+        def _(self, dataframe: cudf.DataFrame) -> cudf.DataFrame:
+            return self._transform_df(dataframe)
+
     def fit_schema(self, input_schema: Schema):
         """Computes input and output schemas for each node in the Workflow graph
 

tests/unit/workflow/test_workflow.py

@@ -53,6 +53,15 @@ def test_workflow_double_fit():
         workflow.transform(df_event).to_ddf().compute()
 
 
+def test_workflow_transform_df():
+    df = make_df({"user_session": ["1", "2", "4", "4", "5"]})
+    ops = ["user_session"] >> nvt.ops.Categorify()
+    dataset = nvt.Dataset(df)
+    workflow = nvt.Workflow(ops)
+    workflow.fit(dataset)
+    assert isinstance(workflow.transform(df), type(df))
+
+
 @pytest.mark.parametrize("engine", ["parquet"])
 def test_workflow_fit_op_rename(tmpdir, dataset, engine):
     # NVT