Repository: NVIDIA-Merlin/NVTabular

Issue:
[FEA] Select a subset of a ColumnGroup
**Is your feature request related to a problem? Please describe.**
No

**Describe the solution you'd like**
Currently, the best way I know of to select a ColumnGroup generated by your code (say as the result of a LambdaOp) is set-subtractive, I.E., you have a ColumnGroup with generated features A,B,C, and just want B,C, the way to get this is group - ['A']. A better way to do this in some cases would be set-selective, IE, group['B','C'].

**Describe alternatives you've considered**
An alternative solution would be to add a select op.

**Additional context**




Discussion:

Resolving patch:
diff --git a/nvtabular/column_group.py b/nvtabular/column_group.py
--- a/nvtabular/column_group.py
+++ b/nvtabular/column_group.py
@@ -138,11 +138,12 @@ def __add__(self, other):
     __radd__ = __add__
 
     def __sub__(self, other):
-        """Adds columns from this ColumnGroup with another to return a new ColumnGroup
+        """Removes columns from this ColumnGroup with another to return a new ColumnGroup
 
         Parameters
         -----------
         other: ColumnGroup or str or list of str
+            Columns to remove
 
         Returns
         -------
@@ -163,6 +164,28 @@ def __sub__(self, other):
         child.kind = f"- {list(to_remove)}"
         return child
 
+    def __getitem__(self, columns):
+        """Selects certain columns from this ColumnGroup, and returns a new Columngroup with only
+        those columns
+
+        Parameters
+        -----------
+        columns: str or list of str
+            Columns to select
+
+        Returns
+        -------
+        ColumnGroup
+        """
+        if isinstance(columns, str):
+            columns = [columns]
+
+        child = ColumnGroup(columns)
+        child.parents = [self]
+        self.children.append(child)
+        child.kind = str(columns)
+        return child
+
     def __repr__(self):
         output = " output" if not self.children else ""
         return f"<ColumnGroup {self.label}{output}>"


Test patch:
diff --git a/tests/unit/test_column_group.py b/tests/unit/test_column_group.py
--- a/tests/unit/test_column_group.py
+++ b/tests/unit/test_column_group.py
@@ -1,11 +1,31 @@
 import cudf
 import pytest
+from cudf.tests.utils import assert_eq
 
 from nvtabular import ColumnGroup, Dataset, Workflow
 from nvtabular.ops import Categorify, Rename
 
 
-def test_nested_column_group(tmpdir):
+def test_column_group_select():
+    df = cudf.DataFrame({"a": [1, 4, 9, 16, 25], "b": [0, 1, 2, 3, 4], "c": [25, 16, 9, 4, 1]})
+
+    input_features = ColumnGroup(["a", "b", "c"])
+    sqrt_features = input_features[["a", "c"]] >> cudf.sqrt
+    plus_one_features = input_features["b"] >> (lambda col: col + 1)
+    features = sqrt_features + plus_one_features
+
+    workflow = Workflow(features)
+    df_out = workflow.fit_transform(Dataset(df)).to_ddf().compute(scheduler="synchronous")
+
+    expected = cudf.DataFrame()
+    expected["a"] = cudf.sqrt(df["a"])
+    expected["c"] = cudf.sqrt(df["c"])
+    expected["b"] = df["b"] + 1
+
+    assert_eq(expected, df_out)
+
+
+def test_nested_column_group():
     df = cudf.DataFrame(
         {
             "geo": ["US>CA", "US>NY", "CA>BC", "CA>ON"],