Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enable Workflow.transform to be run with a DataFrame type #1777

Merged
merged 6 commits into from
Mar 21, 2023
Merged

Enable Workflow.transform to be run with a DataFrame type #1777

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
d324acf
Enable `Workflow.transform` to be run with a DataFrame type
oliverholworthy Mar 10, 2023
d65e9ce
Update typehint for transform arguments
oliverholworthy Mar 10, 2023
6a7f43b
Use DataFrameType in Workflow typehints matching existing code style
oliverholworthy Mar 10, 2023
86a42ed
Use singledispatch to split up transform method by type
oliverholworthy Mar 10, 2023
5e73a29
Update docstring of `Workflow.transform` to match new types
oliverholworthy Mar 20, 2023
cadc1b5
Merge branch 'main' into workflow-transform-dataframe
oliverholworthy Mar 20, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions nvtabular/workflow/workflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
import time
import types
import warnings
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING, Optional, Union

import cloudpickle
import fsspec
Expand Down Expand Up @@ -78,7 +78,9 @@ 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:
def transform(
self, dataset: Union[Dataset, "cudf.DataFrame", pd.DataFrame]
karlhigley marked this conversation as resolved.
Show resolved Hide resolved
) -> Union[Dataset, "cudf.DataFrame", pd.DataFrame]:
"""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

Expand All @@ -96,7 +98,15 @@ def transform(self, dataset: Dataset) -> Dataset:
Dataset
Transformed Dataset with the workflow graph applied to it
"""
return self._transform_impl(dataset)
if isinstance(dataset, Dataset):
return self._transform_impl(dataset)
elif isinstance(dataset, pd.DataFrame) or (cudf and isinstance(dataset, cudf.DataFrame)):
return self._transform_df(dataset)
else:
raise ValueError(
"Workflow.transform received an unsupported type: {type(dataset)} "
"Supported types are a `merlin.io.Dataset` or DataFrame (pandas or cudf)"
)
karlhigley marked this conversation as resolved.
Show resolved Hide resolved

def fit_schema(self, input_schema: Schema):
"""Computes input and output schemas for each node in the Workflow graph
Expand Down
9 changes: 9 additions & 0 deletions tests/unit/workflow/test_workflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down