We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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**



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `nvtabular/column_group.py`
Content:
```
1 #
2 # Copyright (c) 2021, NVIDIA CORPORATION.
3 #
4 # Licensed under the Apache License, Version 2.0 (the "License");
5 # you may not use this file except in compliance with the License.
6 # You may obtain a copy of the License at
7 #
8 #     http://www.apache.org/licenses/LICENSE-2.0
9 #
10 # Unless required by applicable law or agreed to in writing, software
11 # distributed under the License is distributed on an "AS IS" BASIS,
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 # See the License for the specific language governing permissions and
14 # limitations under the License.
15 #
16 import collections.abc
17 
18 from dask.core import flatten
19 
20 from nvtabular.ops import LambdaOp, Operator
21 
22 
23 class ColumnGroup:
24     """A ColumnGroup is a group of columns that you want to apply the same transformations to.
25     ColumnGroup's can be transformed by shifting operators on to them, which returns a new
26     ColumnGroup with the transformations applied. This lets you define a graph of operations
27     that makes up your workflow
28 
29     Parameters
30     ----------
31     columns: list of (str or tuple of str)
32         The columns to select from the input Dataset. The elements of this list are strings
33         indicating the column names in most cases, but can also be tuples of strings
34         for feature crosses.
35     """
36 
37     def __init__(self, columns):
38         self.parents = []
39         self.children = []
40         self.op = None
41         self.kind = None
42         self.dependencies = None
43 
44         if isinstance(columns, str):
45             columns = [columns]
46 
47         # if any of the values we're passed are a columngroup
48         # we have to ourselves as a childnode in the graph.
49         if any(isinstance(col, ColumnGroup) for col in columns):
50             self.columns = []
51             self.kind = "[...]"
52             for col in columns:
53                 if not isinstance(col, ColumnGroup):
54                     col = ColumnGroup(col)
55                 else:
56                     # we can't handle nesting arbitrarily deep here
57                     # only accept non-nested (str) columns here
58                     if any(not isinstance(c, str) for c in col.columns):
59                         raise ValueError("Can't handle more than 1 level of nested columns")
60 
61                 col.children.append(self)
62                 self.parents.append(col)
63                 self.columns.append(tuple(col.columns))
64 
65         else:
66             self.columns = [_convert_col(col) for col in columns]
67 
68     def __rshift__(self, operator):
69         """Transforms this ColumnGroup by applying an Operator
70 
71         Parameters
72         -----------
73         operators: Operator or callable
74 
75         Returns
76         -------
77         ColumnGroup
78         """
79         if isinstance(operator, type) and issubclass(operator, Operator):
80             # handle case where an operator class is passed
81             operator = operator()
82         elif callable(operator):
83             # implicit lambdaop conversion.
84             operator = LambdaOp(operator)
85 
86         if not isinstance(operator, Operator):
87             raise ValueError(f"Expected operator or callable, got {operator.__class__}")
88 
89         child = ColumnGroup(operator.output_column_names(self.columns))
90         child.parents = [self]
91         self.children.append(child)
92         child.op = operator
93 
94         dependencies = operator.dependencies()
95         if dependencies:
96             child.dependencies = set()
97             if not isinstance(dependencies, collections.abc.Sequence):
98                 dependencies = [dependencies]
99 
100             for dependency in dependencies:
101                 if not isinstance(dependency, ColumnGroup):
102                     dependency = ColumnGroup(dependency)
103                 dependency.children.append(child)
104                 child.parents.append(dependency)
105                 child.dependencies.add(dependency)
106 
107         return child
108 
109     def __add__(self, other):
110         """Adds columns from this ColumnGroup with another to return a new ColumnGroup
111 
112         Parameters
113         -----------
114         other: ColumnGroup or str or list of str
115 
116         Returns
117         -------
118         ColumnGroup
119         """
120         if isinstance(other, str):
121             other = ColumnGroup([other])
122         elif isinstance(other, collections.abc.Sequence):
123             other = ColumnGroup(other)
124 
125         # check if there are any columns with the same name in both column groups
126         overlap = set(self.columns).intersection(other.columns)
127         if overlap:
128             raise ValueError(f"duplicate column names found: {overlap}")
129 
130         child = ColumnGroup(self.columns + other.columns)
131         child.parents = [self, other]
132         child.kind = "+"
133         self.children.append(child)
134         other.children.append(child)
135         return child
136 
137     # handle the "column_name" + ColumnGroup case
138     __radd__ = __add__
139 
140     def __sub__(self, other):
141         """Adds columns from this ColumnGroup with another to return a new ColumnGroup
142 
143         Parameters
144         -----------
145         other: ColumnGroup or str or list of str
146 
147         Returns
148         -------
149         ColumnGroup
150         """
151         if isinstance(other, ColumnGroup):
152             to_remove = set(other.columns)
153         elif isinstance(other, str):
154             to_remove = {other}
155         elif isinstance(other, collections.abc.Sequence):
156             to_remove = set(other)
157         else:
158             raise ValueError(f"Expected ColumnGroup, str, or list of str. Got {other.__class__}")
159         new_columns = [c for c in self.columns if c not in to_remove]
160         child = ColumnGroup(new_columns)
161         child.parents = [self]
162         self.children.append(child)
163         child.kind = f"- {list(to_remove)}"
164         return child
165 
166     def __repr__(self):
167         output = " output" if not self.children else ""
168         return f"<ColumnGroup {self.label}{output}>"
169 
170     @property
171     def flattened_columns(self):
172         return list(flatten(self.columns, container=tuple))
173 
174     @property
175     def input_column_names(self):
176         """Returns the names of columns in the main chain"""
177         dependencies = self.dependencies or set()
178         return [
179             col for parent in self.parents for col in parent.columns if parent not in dependencies
180         ]
181 
182     @property
183     def label(self):
184         if self.op:
185             return str(self.op.__class__.__name__)
186         elif self.kind:
187             return self.kind
188         elif not self.parents:
189             return f"input cols=[{self._cols_repr}]"
190         else:
191             return "??"
192 
193     @property
194     def _cols_repr(self):
195         cols = ", ".join(map(str, self.columns[:3]))
196         if len(self.columns) > 3:
197             cols += "..."
198         return cols
199 
200     @property
201     def graph(self):
202         return _to_graphviz(self)
203 
204 
205 def iter_nodes(nodes):
206     queue = nodes[:]
207     while queue:
208         current = queue.pop()
209         yield current
210         # TODO: deduplicate nodes?
211         for parent in current.parents:
212             queue.append(parent)
213 
214 
215 def _to_graphviz(column_group):
216     """Converts a ColumnGroup to a GraphViz DiGraph object useful for display in notebooks"""
217     from graphviz import Digraph
218 
219     column_group = _merge_add_nodes(column_group)
220     graph = Digraph()
221 
222     # get all the nodes from parents of this columngroup
223     # and add edges between each of them
224     allnodes = list(set(iter_nodes([column_group])))
225     node_ids = {v: str(k) for k, v in enumerate(allnodes)}
226     for node, nodeid in node_ids.items():
227         graph.node(nodeid, node.label)
228         for parent in node.parents:
229             graph.edge(node_ids[parent], nodeid)
230 
231     # add a single 'output' node representing the final state
232     output_node_id = str(len(allnodes))
233     graph.node(output_node_id, f"output cols=[{column_group._cols_repr}]")
234     graph.edge(node_ids[column_group], output_node_id)
235     return graph
236 
237 
238 def _merge_add_nodes(graph):
239     """Merges repeat '+' nodes, leading to nicer looking outputs"""
240     # lets take a copy to avoid mutating the input
241     import copy
242 
243     graph = copy.copy(graph)
244 
245     queue = [graph]
246     while queue:
247         current = queue.pop()
248         if current.kind == "+":
249             changed = True
250             while changed:
251                 changed = False
252                 parents = []
253                 for i, parent in enumerate(current.parents):
254                     if parent.kind == "+" and len(parent.children) == 1:
255                         changed = True
256                         # disconnect parent, point all the grandparents at current instead
257                         parents.extend(parent.parents)
258                         for grandparent in parent.parents:
259                             grandparent.children = [
260                                 current if child == parent else child
261                                 for child in grandparent.children
262                             ]
263                     else:
264                         parents.append(parent)
265                 current.parents = parents
266 
267         queue.extend(current.parents)
268 
269     return graph
270 
271 
272 def _convert_col(col):
273     if isinstance(col, (str, tuple)):
274         return col
275     elif isinstance(col, list):
276         return tuple(col)
277     else:
278         raise ValueError(f"Invalid column value for ColumnGroup: {col}")
279
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     __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 @@
         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}>"

{"golden_diff": "diff --git a/nvtabular/column_group.py b/nvtabular/column_group.py\n--- a/nvtabular/column_group.py\n+++ b/nvtabular/column_group.py\n@@ -138,11 +138,12 @@\n     __radd__ = __add__\n \n     def __sub__(self, other):\n-        \"\"\"Adds columns from this ColumnGroup with another to return a new ColumnGroup\n+        \"\"\"Removes columns from this ColumnGroup with another to return a new ColumnGroup\n \n         Parameters\n         -----------\n         other: ColumnGroup or str or list of str\n+            Columns to remove\n \n         Returns\n         -------\n@@ -163,6 +164,28 @@\n         child.kind = f\"- {list(to_remove)}\"\n         return child\n \n+    def __getitem__(self, columns):\n+        \"\"\"Selects certain columns from this ColumnGroup, and returns a new Columngroup with only\n+        those columns\n+\n+        Parameters\n+        -----------\n+        columns: str or list of str\n+            Columns to select\n+\n+        Returns\n+        -------\n+        ColumnGroup\n+        \"\"\"\n+        if isinstance(columns, str):\n+            columns = [columns]\n+\n+        child = ColumnGroup(columns)\n+        child.parents = [self]\n+        self.children.append(child)\n+        child.kind = str(columns)\n+        return child\n+\n     def __repr__(self):\n         output = \" output\" if not self.children else \"\"\n         return f\"<ColumnGroup {self.label}{output}>\"\n", "issue": "[FEA] Select a subset of a ColumnGroup\n**Is your feature request related to a problem? Please describe.**\r\nNo\r\n\r\n**Describe the solution you'd like**\r\nCurrently, 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'].\r\n\r\n**Describe alternatives you've considered**\r\nAn alternative solution would be to add a select op.\r\n\r\n**Additional context**\r\n\r\n\n", "before_files": [{"content": "#\n# Copyright (c) 2021, NVIDIA CORPORATION.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\nimport collections.abc\n\nfrom dask.core import flatten\n\nfrom nvtabular.ops import LambdaOp, Operator\n\n\nclass ColumnGroup:\n    \"\"\"A ColumnGroup is a group of columns that you want to apply the same transformations to.\n    ColumnGroup's can be transformed by shifting operators on to them, which returns a new\n    ColumnGroup with the transformations applied. This lets you define a graph of operations\n    that makes up your workflow\n\n    Parameters\n    ----------\n    columns: list of (str or tuple of str)\n        The columns to select from the input Dataset. The elements of this list are strings\n        indicating the column names in most cases, but can also be tuples of strings\n        for feature crosses.\n    \"\"\"\n\n    def __init__(self, columns):\n        self.parents = []\n        self.children = []\n        self.op = None\n        self.kind = None\n        self.dependencies = None\n\n        if isinstance(columns, str):\n            columns = [columns]\n\n        # if any of the values we're passed are a columngroup\n        # we have to ourselves as a childnode in the graph.\n        if any(isinstance(col, ColumnGroup) for col in columns):\n            self.columns = []\n            self.kind = \"[...]\"\n            for col in columns:\n                if not isinstance(col, ColumnGroup):\n                    col = ColumnGroup(col)\n                else:\n                    # we can't handle nesting arbitrarily deep here\n                    # only accept non-nested (str) columns here\n                    if any(not isinstance(c, str) for c in col.columns):\n                        raise ValueError(\"Can't handle more than 1 level of nested columns\")\n\n                col.children.append(self)\n                self.parents.append(col)\n                self.columns.append(tuple(col.columns))\n\n        else:\n            self.columns = [_convert_col(col) for col in columns]\n\n    def __rshift__(self, operator):\n        \"\"\"Transforms this ColumnGroup by applying an Operator\n\n        Parameters\n        -----------\n        operators: Operator or callable\n\n        Returns\n        -------\n        ColumnGroup\n        \"\"\"\n        if isinstance(operator, type) and issubclass(operator, Operator):\n            # handle case where an operator class is passed\n            operator = operator()\n        elif callable(operator):\n            # implicit lambdaop conversion.\n            operator = LambdaOp(operator)\n\n        if not isinstance(operator, Operator):\n            raise ValueError(f\"Expected operator or callable, got {operator.__class__}\")\n\n        child = ColumnGroup(operator.output_column_names(self.columns))\n        child.parents = [self]\n        self.children.append(child)\n        child.op = operator\n\n        dependencies = operator.dependencies()\n        if dependencies:\n            child.dependencies = set()\n            if not isinstance(dependencies, collections.abc.Sequence):\n                dependencies = [dependencies]\n\n            for dependency in dependencies:\n                if not isinstance(dependency, ColumnGroup):\n                    dependency = ColumnGroup(dependency)\n                dependency.children.append(child)\n                child.parents.append(dependency)\n                child.dependencies.add(dependency)\n\n        return child\n\n    def __add__(self, other):\n        \"\"\"Adds columns from this ColumnGroup with another to return a new ColumnGroup\n\n        Parameters\n        -----------\n        other: ColumnGroup or str or list of str\n\n        Returns\n        -------\n        ColumnGroup\n        \"\"\"\n        if isinstance(other, str):\n            other = ColumnGroup([other])\n        elif isinstance(other, collections.abc.Sequence):\n            other = ColumnGroup(other)\n\n        # check if there are any columns with the same name in both column groups\n        overlap = set(self.columns).intersection(other.columns)\n        if overlap:\n            raise ValueError(f\"duplicate column names found: {overlap}\")\n\n        child = ColumnGroup(self.columns + other.columns)\n        child.parents = [self, other]\n        child.kind = \"+\"\n        self.children.append(child)\n        other.children.append(child)\n        return child\n\n    # handle the \"column_name\" + ColumnGroup case\n    __radd__ = __add__\n\n    def __sub__(self, other):\n        \"\"\"Adds columns from this ColumnGroup with another to return a new ColumnGroup\n\n        Parameters\n        -----------\n        other: ColumnGroup or str or list of str\n\n        Returns\n        -------\n        ColumnGroup\n        \"\"\"\n        if isinstance(other, ColumnGroup):\n            to_remove = set(other.columns)\n        elif isinstance(other, str):\n            to_remove = {other}\n        elif isinstance(other, collections.abc.Sequence):\n            to_remove = set(other)\n        else:\n            raise ValueError(f\"Expected ColumnGroup, str, or list of str. Got {other.__class__}\")\n        new_columns = [c for c in self.columns if c not in to_remove]\n        child = ColumnGroup(new_columns)\n        child.parents = [self]\n        self.children.append(child)\n        child.kind = f\"- {list(to_remove)}\"\n        return child\n\n    def __repr__(self):\n        output = \" output\" if not self.children else \"\"\n        return f\"<ColumnGroup {self.label}{output}>\"\n\n    @property\n    def flattened_columns(self):\n        return list(flatten(self.columns, container=tuple))\n\n    @property\n    def input_column_names(self):\n        \"\"\"Returns the names of columns in the main chain\"\"\"\n        dependencies = self.dependencies or set()\n        return [\n            col for parent in self.parents for col in parent.columns if parent not in dependencies\n        ]\n\n    @property\n    def label(self):\n        if self.op:\n            return str(self.op.__class__.__name__)\n        elif self.kind:\n            return self.kind\n        elif not self.parents:\n            return f\"input cols=[{self._cols_repr}]\"\n        else:\n            return \"??\"\n\n    @property\n    def _cols_repr(self):\n        cols = \", \".join(map(str, self.columns[:3]))\n        if len(self.columns) > 3:\n            cols += \"...\"\n        return cols\n\n    @property\n    def graph(self):\n        return _to_graphviz(self)\n\n\ndef iter_nodes(nodes):\n    queue = nodes[:]\n    while queue:\n        current = queue.pop()\n        yield current\n        # TODO: deduplicate nodes?\n        for parent in current.parents:\n            queue.append(parent)\n\n\ndef _to_graphviz(column_group):\n    \"\"\"Converts a ColumnGroup to a GraphViz DiGraph object useful for display in notebooks\"\"\"\n    from graphviz import Digraph\n\n    column_group = _merge_add_nodes(column_group)\n    graph = Digraph()\n\n    # get all the nodes from parents of this columngroup\n    # and add edges between each of them\n    allnodes = list(set(iter_nodes([column_group])))\n    node_ids = {v: str(k) for k, v in enumerate(allnodes)}\n    for node, nodeid in node_ids.items():\n        graph.node(nodeid, node.label)\n        for parent in node.parents:\n            graph.edge(node_ids[parent], nodeid)\n\n    # add a single 'output' node representing the final state\n    output_node_id = str(len(allnodes))\n    graph.node(output_node_id, f\"output cols=[{column_group._cols_repr}]\")\n    graph.edge(node_ids[column_group], output_node_id)\n    return graph\n\n\ndef _merge_add_nodes(graph):\n    \"\"\"Merges repeat '+' nodes, leading to nicer looking outputs\"\"\"\n    # lets take a copy to avoid mutating the input\n    import copy\n\n    graph = copy.copy(graph)\n\n    queue = [graph]\n    while queue:\n        current = queue.pop()\n        if current.kind == \"+\":\n            changed = True\n            while changed:\n                changed = False\n                parents = []\n                for i, parent in enumerate(current.parents):\n                    if parent.kind == \"+\" and len(parent.children) == 1:\n                        changed = True\n                        # disconnect parent, point all the grandparents at current instead\n                        parents.extend(parent.parents)\n                        for grandparent in parent.parents:\n                            grandparent.children = [\n                                current if child == parent else child\n                                for child in grandparent.children\n                            ]\n                    else:\n                        parents.append(parent)\n                current.parents = parents\n\n        queue.extend(current.parents)\n\n    return graph\n\n\ndef _convert_col(col):\n    if isinstance(col, (str, tuple)):\n        return col\n    elif isinstance(col, list):\n        return tuple(col)\n    else:\n        raise ValueError(f\"Invalid column value for ColumnGroup: {col}\")\n", "path": "nvtabular/column_group.py"}], "after_files": [{"content": "#\n# Copyright (c) 2021, NVIDIA CORPORATION.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\nimport collections.abc\n\nfrom dask.core import flatten\n\nfrom nvtabular.ops import LambdaOp, Operator\n\n\nclass ColumnGroup:\n    \"\"\"A ColumnGroup is a group of columns that you want to apply the same transformations to.\n    ColumnGroup's can be transformed by shifting operators on to them, which returns a new\n    ColumnGroup with the transformations applied. This lets you define a graph of operations\n    that makes up your workflow\n\n    Parameters\n    ----------\n    columns: list of (str or tuple of str)\n        The columns to select from the input Dataset. The elements of this list are strings\n        indicating the column names in most cases, but can also be tuples of strings\n        for feature crosses.\n    \"\"\"\n\n    def __init__(self, columns):\n        self.parents = []\n        self.children = []\n        self.op = None\n        self.kind = None\n        self.dependencies = None\n\n        if isinstance(columns, str):\n            columns = [columns]\n\n        # if any of the values we're passed are a columngroup\n        # we have to ourselves as a childnode in the graph.\n        if any(isinstance(col, ColumnGroup) for col in columns):\n            self.columns = []\n            self.kind = \"[...]\"\n            for col in columns:\n                if not isinstance(col, ColumnGroup):\n                    col = ColumnGroup(col)\n                else:\n                    # we can't handle nesting arbitrarily deep here\n                    # only accept non-nested (str) columns here\n                    if any(not isinstance(c, str) for c in col.columns):\n                        raise ValueError(\"Can't handle more than 1 level of nested columns\")\n\n                col.children.append(self)\n                self.parents.append(col)\n                self.columns.append(tuple(col.columns))\n\n        else:\n            self.columns = [_convert_col(col) for col in columns]\n\n    def __rshift__(self, operator):\n        \"\"\"Transforms this ColumnGroup by applying an Operator\n\n        Parameters\n        -----------\n        operators: Operator or callable\n\n        Returns\n        -------\n        ColumnGroup\n        \"\"\"\n        if isinstance(operator, type) and issubclass(operator, Operator):\n            # handle case where an operator class is passed\n            operator = operator()\n        elif callable(operator):\n            # implicit lambdaop conversion.\n            operator = LambdaOp(operator)\n\n        if not isinstance(operator, Operator):\n            raise ValueError(f\"Expected operator or callable, got {operator.__class__}\")\n\n        child = ColumnGroup(operator.output_column_names(self.columns))\n        child.parents = [self]\n        self.children.append(child)\n        child.op = operator\n\n        dependencies = operator.dependencies()\n        if dependencies:\n            child.dependencies = set()\n            if not isinstance(dependencies, collections.abc.Sequence):\n                dependencies = [dependencies]\n\n            for dependency in dependencies:\n                if not isinstance(dependency, ColumnGroup):\n                    dependency = ColumnGroup(dependency)\n                dependency.children.append(child)\n                child.parents.append(dependency)\n                child.dependencies.add(dependency)\n\n        return child\n\n    def __add__(self, other):\n        \"\"\"Adds columns from this ColumnGroup with another to return a new ColumnGroup\n\n        Parameters\n        -----------\n        other: ColumnGroup or str or list of str\n\n        Returns\n        -------\n        ColumnGroup\n        \"\"\"\n        if isinstance(other, str):\n            other = ColumnGroup([other])\n        elif isinstance(other, collections.abc.Sequence):\n            other = ColumnGroup(other)\n\n        # check if there are any columns with the same name in both column groups\n        overlap = set(self.columns).intersection(other.columns)\n        if overlap:\n            raise ValueError(f\"duplicate column names found: {overlap}\")\n\n        child = ColumnGroup(self.columns + other.columns)\n        child.parents = [self, other]\n        child.kind = \"+\"\n        self.children.append(child)\n        other.children.append(child)\n        return child\n\n    # handle the \"column_name\" + ColumnGroup case\n    __radd__ = __add__\n\n    def __sub__(self, other):\n        \"\"\"Removes columns from this ColumnGroup with another to return a new ColumnGroup\n\n        Parameters\n        -----------\n        other: ColumnGroup or str or list of str\n            Columns to remove\n\n        Returns\n        -------\n        ColumnGroup\n        \"\"\"\n        if isinstance(other, ColumnGroup):\n            to_remove = set(other.columns)\n        elif isinstance(other, str):\n            to_remove = {other}\n        elif isinstance(other, collections.abc.Sequence):\n            to_remove = set(other)\n        else:\n            raise ValueError(f\"Expected ColumnGroup, str, or list of str. Got {other.__class__}\")\n        new_columns = [c for c in self.columns if c not in to_remove]\n        child = ColumnGroup(new_columns)\n        child.parents = [self]\n        self.children.append(child)\n        child.kind = f\"- {list(to_remove)}\"\n        return child\n\n    def __getitem__(self, columns):\n        \"\"\"Selects certain columns from this ColumnGroup, and returns a new Columngroup with only\n        those columns\n\n        Parameters\n        -----------\n        columns: str or list of str\n            Columns to select\n\n        Returns\n        -------\n        ColumnGroup\n        \"\"\"\n        if isinstance(columns, str):\n            columns = [columns]\n\n        child = ColumnGroup(columns)\n        child.parents = [self]\n        self.children.append(child)\n        child.kind = str(columns)\n        return child\n\n    def __repr__(self):\n        output = \" output\" if not self.children else \"\"\n        return f\"<ColumnGroup {self.label}{output}>\"\n\n    @property\n    def flattened_columns(self):\n        return list(flatten(self.columns, container=tuple))\n\n    @property\n    def input_column_names(self):\n        \"\"\"Returns the names of columns in the main chain\"\"\"\n        dependencies = self.dependencies or set()\n        return [\n            col for parent in self.parents for col in parent.columns if parent not in dependencies\n        ]\n\n    @property\n    def label(self):\n        if self.op:\n            return str(self.op.__class__.__name__)\n        elif self.kind:\n            return self.kind\n        elif not self.parents:\n            return f\"input cols=[{self._cols_repr}]\"\n        else:\n            return \"??\"\n\n    @property\n    def _cols_repr(self):\n        cols = \", \".join(map(str, self.columns[:3]))\n        if len(self.columns) > 3:\n            cols += \"...\"\n        return cols\n\n    @property\n    def graph(self):\n        return _to_graphviz(self)\n\n\ndef iter_nodes(nodes):\n    queue = nodes[:]\n    while queue:\n        current = queue.pop()\n        yield current\n        # TODO: deduplicate nodes?\n        for parent in current.parents:\n            queue.append(parent)\n\n\ndef _to_graphviz(column_group):\n    \"\"\"Converts a ColumnGroup to a GraphViz DiGraph object useful for display in notebooks\"\"\"\n    from graphviz import Digraph\n\n    column_group = _merge_add_nodes(column_group)\n    graph = Digraph()\n\n    # get all the nodes from parents of this columngroup\n    # and add edges between each of them\n    allnodes = list(set(iter_nodes([column_group])))\n    node_ids = {v: str(k) for k, v in enumerate(allnodes)}\n    for node, nodeid in node_ids.items():\n        graph.node(nodeid, node.label)\n        for parent in node.parents:\n            graph.edge(node_ids[parent], nodeid)\n\n    # add a single 'output' node representing the final state\n    output_node_id = str(len(allnodes))\n    graph.node(output_node_id, f\"output cols=[{column_group._cols_repr}]\")\n    graph.edge(node_ids[column_group], output_node_id)\n    return graph\n\n\ndef _merge_add_nodes(graph):\n    \"\"\"Merges repeat '+' nodes, leading to nicer looking outputs\"\"\"\n    # lets take a copy to avoid mutating the input\n    import copy\n\n    graph = copy.copy(graph)\n\n    queue = [graph]\n    while queue:\n        current = queue.pop()\n        if current.kind == \"+\":\n            changed = True\n            while changed:\n                changed = False\n                parents = []\n                for i, parent in enumerate(current.parents):\n                    if parent.kind == \"+\" and len(parent.children) == 1:\n                        changed = True\n                        # disconnect parent, point all the grandparents at current instead\n                        parents.extend(parent.parents)\n                        for grandparent in parent.parents:\n                            grandparent.children = [\n                                current if child == parent else child\n                                for child in grandparent.children\n                            ]\n                    else:\n                        parents.append(parent)\n                current.parents = parents\n\n        queue.extend(current.parents)\n\n    return graph\n\n\ndef _convert_col(col):\n    if isinstance(col, (str, tuple)):\n        return col\n    elif isinstance(col, list):\n        return tuple(col)\n    else:\n        raise ValueError(f\"Invalid column value for ColumnGroup: {col}\")\n", "path": "nvtabular/column_group.py"}]}