diff --git a/nvtabular/ops/groupby.py b/nvtabular/ops/groupby.py index 93748e4f19..da343a5608 100644 --- a/nvtabular/ops/groupby.py +++ b/nvtabular/ops/groupby.py @@ -20,7 +20,7 @@ from merlin.core.dispatch import DataFrameType, annotate from merlin.dtypes.shape import DefaultShapes from merlin.schema import Schema -from nvtabular.ops.operator import ColumnSelector, DataFormats, Operator +from nvtabular.ops.operator import ColumnSelector, Operator class Groupby(Operator): @@ -111,10 +111,6 @@ def __init__( self.name_sep = name_sep super().__init__() - @property - def supported_formats(self): - return DataFormats.PANDAS_DATAFRAME | DataFormats.CUDF_DATAFRAME - @annotate("Groupby_op", color="darkgreen", domain="nvt_python") def transform(self, col_selector: ColumnSelector, df: DataFrameType) -> DataFrameType: # Sort if necessary diff --git a/nvtabular/ops/operator.py b/nvtabular/ops/operator.py index 41b3621643..a0aa99ab6a 100644 --- a/nvtabular/ops/operator.py +++ b/nvtabular/ops/operator.py @@ -19,4 +19,9 @@ DataFormats, ) -Operator = BaseOperator + +# Avoid TENSOR_TABLE by default (for now) +class Operator(BaseOperator): + @property + def supported_formats(self): + return DataFormats.PANDAS_DATAFRAME | DataFormats.CUDF_DATAFRAME