Add Rename to the set of core DAG ops that work in all DAGs (#312)

* Add a utility for testing equality between `TensorTables`

* Add `Rename` to the set of core DAG ops that work in all DAGs

This should work with either dataframes or `TensorTables`, which requires minor changes to TT to add a definition of equality and make it possible to change column names.

Adjust param names and types to fix linter issue

* Add docstrings to TensorTable methods

* Adjust the implementation of `TensorTable.columns`

* Mark assignment in `Rename` to be ignored by `mypy` linter

* Appease the linter's line length checks

* Update `Rename` op tests to use `assert_transformable_equal`

Author: karlhigley

merlin/dag/ops/rename.py

@@ -0,0 +1,83 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from merlin.core.protocols import Transformable
+from merlin.dag import ColumnSelector
+from merlin.dag.base_operator import BaseOperator
+
+
+class Rename(BaseOperator):
+    """This operation renames columns by one of several methods:
+
+        - using a user defined lambda function to transform column names
+        - appending a postfix string to every column name
+        - renaming a single column to a single fixed string
+
+    Example usage::
+
+        # Rename columns after LogOp
+        cont_features = cont_names >> nvt.ops.LogOp() >> Rename(postfix='_log')
+        processor = nvt.Workflow(cont_features)
+
+    Parameters
+    ----------
+    f : callable, optional
+        Function that takes a column name and returns a new column name
+    postfix : str, optional
+        If set each column name in the output will have this string appended to it
+    name : str, optional
+        If set, a single input column will be renamed to this string
+    """
+
+    def __init__(self, f=None, postfix=None, name=None):
+        if not f and postfix is None and name is None:
+            raise ValueError("must specify name, f, or postfix, for Rename op")
+
+        self.f = f
+        self.postfix = postfix
+        self.name = name
+        super().__init__()
+
+    def transform(
+        self, col_selector: ColumnSelector, transformable: Transformable
+    ) -> Transformable:
+        transformable = transformable[col_selector.names]
+        transformable.columns = list(  # type: ignore[assignment]
+            self.column_mapping(col_selector).keys()
+        )
+        return transformable
+
+    transform.__doc__ = BaseOperator.transform.__doc__
+
+    def column_mapping(self, col_selector):
+        column_mapping = {}
+        for col_name in col_selector.names:
+            if self.f:
+                new_col_name = self.f(col_name)
+            elif self.postfix:
+                new_col_name = col_name + self.postfix
+            elif self.name:
+                if len(col_selector.names) == 1:
+                    new_col_name = self.name
+                else:
+                    raise RuntimeError("Single column name provided for renaming multiple columns")
+            else:
+                raise RuntimeError(
+                    "The Rename op requires one of f, postfix, or name to be provided"
+                )
+
+            column_mapping[new_col_name] = [col_name]
+
+        return column_mapping

merlin/table/tensor_table.py

@@ -209,6 +209,13 @@ def columns(self) -> List[str]:
         """
         return list(self.keys())
 
+    @columns.setter
+    def columns(self, col_names):
+        renamed_columns = {}
+        for col, col_name in zip(self.columns, col_names):
+            renamed_columns[col_name] = self._columns[col]
+        self._columns = renamed_columns
+
     @property
     def column_type(self):
         return type(list(self.values())[0])
@@ -241,10 +248,26 @@ def to_dict(self):
         return result
 
     def cpu(self):
+        """
+        Move this TensorTable and its columns to CPU
+
+        Returns
+        -------
+        TensorTable
+            A new TensorTable containing the same columns but on CPU
+        """
         columns = {col_name: col_values.cpu() for col_name, col_values in self.items()}
         return TensorTable(columns)
 
     def gpu(self):
+        """
+        Move this TensorTable and its columns to GPU
+
+        Returns
+        -------
+        TensorTable
+            A new TensorTable containing the same columns but on GPU
+        """
         columns = {col_name: col_values.gpu() for col_name, col_values in self.items()}
         return TensorTable(columns)
 

merlin/testing/__init__.py

@@ -0,0 +1,18 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# flake8: noqa
+from merlin.testing.assert_equals import assert_transformable_equal

merlin/testing/assert_equals.py

@@ -0,0 +1,51 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from merlin.core.compat import pandas as pd
+from merlin.dispatch.lazy import lazy_singledispatch
+from merlin.table import TensorTable
+
+
+def assert_table_equal(left: TensorTable, right: TensorTable):
+    pd.testing.assert_frame_equal(left.cpu().to_df(), right.cpu().to_df())
+
+
+@lazy_singledispatch
+def assert_transformable_equal(left, right):
+    raise NotImplementedError
+
+
+@assert_transformable_equal.register(TensorTable)
+def _assert_equal_table(left, right):
+    assert_table_equal(left, right)
+
+
+@assert_transformable_equal.register_lazy("cudf")
+def _register_assert_equal_df_cudf():
+    import cudf
+
+    @assert_transformable_equal.register(cudf.DataFrame)
+    def _assert_equal_df_cudf(left, right):
+        cudf.testing.assert_frame_equal(left, right)
+
+
+@assert_transformable_equal.register_lazy("pandas")
+def _register_assert_equal_pandas():
+    import pandas
+
+    @assert_transformable_equal.register(pandas.DataFrame)
+    def _assert_equal_pandas(left, right):
+        pandas.testing.assert_frame_equal(left, right)

tests/unit/dag/ops/test_rename.py

@@ -0,0 +1,62 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import numpy as np
+import pandas as pd
+import pytest
+
+from merlin.core.compat import cudf
+from merlin.dag import ColumnSelector
+from merlin.dag.ops.rename import Rename
+from merlin.table import TensorTable
+from merlin.testing import assert_transformable_equal
+
+transformables = [pd.DataFrame, TensorTable]
+if cudf:
+    transformables.append(cudf.DataFrame)
+
+
+@pytest.mark.parametrize("transformable", transformables)
+def test_rename(transformable):
+    df = transformable({"x": np.array([1, 2, 3, 4, 5]), "y": np.array([6, 7, 8, 9, 10])})
+
+    selector = ColumnSelector(["x", "y"])
+
+    op = Rename(f=lambda name: name.upper())
+    transformed = op.transform(selector, df)
+    expected = transformable({"X": np.array([1, 2, 3, 4, 5]), "Y": np.array([6, 7, 8, 9, 10])})
+    assert_transformable_equal(transformed, expected)
+
+    op = Rename(postfix="_lower")
+    transformed = op.transform(selector, df)
+    expected = transformable(
+        {
+            "x_lower": np.array([1, 2, 3, 4, 5]),
+            "y_lower": np.array([6, 7, 8, 9, 10]),
+        }
+    )
+    assert_transformable_equal(transformed, expected)
+
+    selector = ColumnSelector(["x"])
+
+    op = Rename(name="z")
+    transformed = op.transform(selector, df)
+    expected = transformable({"z": np.array([1, 2, 3, 4, 5])})
+    assert_transformable_equal(transformed, expected)
+
+    op = Rename(f=lambda name: name.upper())
+    transformed = op.transform(selector, df)
+    expected = transformable({"X": np.array([1, 2, 3, 4, 5])})
+    assert_transformable_equal(transformed, expected)