Rework executor transform methods to accept a Graph (#158)

* Rework executor `transform` methods to accept a `Graph`

* Allow `Node` as a potential value of `transform`'s `graph` arg (for now)

Author: karlhigley

merlin/dag/executors.py

@@ -25,7 +25,7 @@
     global_dask_client,
     set_client_deprecated,
 )
-from merlin.dag import ColumnSelector, Node
+from merlin.dag import ColumnSelector, Graph, Node
 from merlin.io.worker import clean_worker_cache
 
 LOG = logging.getLogger("merlin")
@@ -39,7 +39,7 @@ class LocalExecutor:
     def transform(
         self,
         transformable,
-        nodes,
+        graph,
         output_dtypes=None,
         additional_columns=None,
         capture_dtypes=False,
@@ -48,6 +48,21 @@ def transform(
         Transforms a single dataframe (possibly a partition of a Dask Dataframe)
         by applying the operators from a collection of Nodes
         """
+        nodes = []
+        if isinstance(graph, Graph):
+            nodes.append(graph.output_node)
+        elif isinstance(graph, Node):
+            nodes.append(graph)
+        elif isinstance(graph, list):
+            nodes = graph
+        else:
+            raise TypeError(
+                f"LocalExecutor detected unsupported type of input for graph: {type(graph)}."
+                " `graph` argument must be either a `Graph` object (preferred)"
+                " or a list of `Node` objects (deprecated, but supported for backward "
+                " compatibility.)"
+            )
+
         output_data = None
 
         for node in nodes:
@@ -220,12 +235,26 @@ def __getstate__(self):
         return {k: v for k, v in self.__dict__.items() if k != "client"}
 
     def transform(
-        self, ddf, nodes, output_dtypes=None, additional_columns=None, capture_dtypes=False
+        self, ddf, graph, output_dtypes=None, additional_columns=None, capture_dtypes=False
     ):
         """
         Transforms all partitions of a Dask Dataframe by applying the operators
         from a collection of Nodes
         """
+        nodes = []
+        if isinstance(graph, Graph):
+            nodes.append(graph.output_node)
+        elif isinstance(graph, Node):
+            nodes.append(graph)
+        elif isinstance(graph, list):
+            nodes = graph
+        else:
+            raise TypeError(
+                f"DaskExecutor detected unsupported type of input for graph: {type(graph)}."
+                " `graph` argument must be either a `Graph` object (preferred)"
+                " or a list of `Node` objects (deprecated, but supported for backward"
+                " compatibility.)"
+            )
 
         self._clear_worker_cache()