Skip to content

Commit

Permalink
add dependencies method (#498)
Browse files Browse the repository at this point in the history
  • Loading branch information
rnyak authored Dec 16, 2020
1 parent 711efd8 commit fd3f35a
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion nvtabular/ops/lambdaop.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,14 +75,15 @@ class LambdaOp(Operator):
Whether to replace existing columns or create new ones.
"""

def __init__(self, f):
def __init__(self, f, dependency=None):
super().__init__()
if f is None:
raise ValueError("f cannot be None. LambdaOp op applies f to dataframe")
self.f = f
self._param_count = len(signature(self.f).parameters)
if self._param_count not in (1, 2):
raise ValueError("lambda function must accept either one or two parameters")
self.dependency = dependency

@annotate("DFLambda_op", color="darkgreen", domain="nvt_python")
def transform(self, columns, gdf: cudf.DataFrame):
Expand All @@ -96,3 +97,6 @@ def transform(self, columns, gdf: cudf.DataFrame):
# shouldn't ever happen,
raise RuntimeError(f"unhandled lambda param count {self._param_count}")
return new_gdf

def dependencies(self):
return self.dependency

0 comments on commit fd3f35a

Please sign in to comment.