Enable the usage of generator functions in nodes (#2161)

* Modify the node and run_node to enable generator nodes

Signed-off-by: Ivan Danov <idanov@users.noreply.github.com>

* Add tests to cover all types of generator functions

Signed-off-by: Ivan Danov <idanov@users.noreply.github.com>

* Fail on running a generator node with async load/save

Signed-off-by: Ivan Danov <idanov@users.noreply.github.com>

* Lint my code changes

Signed-off-by: Ivan Danov <idanov@users.noreply.github.com>

* Add changelog to RELEASE.md

Signed-off-by: Ivan Danov <idanov@users.noreply.github.com>

* Simplify the usage of spy and clarify with a comment

Signed-off-by: Ivan Danov <idanov@users.noreply.github.com>

* Improve error messaging

* Improve readability slightly in certain places

* Correct the expected error message in node run tests

Signed-off-by: Ivan Danov <idanov@users.noreply.github.com>

* Revert the eager evaluation of the result in _from_dict

Generators cannot refer to themselves in their definition, or they will fail when used.

Signed-off-by: Ivan Danov <idanov@users.noreply.github.com>

* Add a comment for the eager evaluation

Signed-off-by: Ivan Danov <idanov@users.noreply.github.com>

Signed-off-by: Ivan Danov <idanov@users.noreply.github.com>

Makefile

@@ -19,6 +19,9 @@ test:
 test-no-spark:
  pytest tests --no-cov --ignore tests/extras/datasets/spark --numprocesses 4 --dist loadfile
 
+test-no-datasets:
+ pytest tests --no-cov --ignore tests/extras/datasets/ --numprocesses 4 --dist loadfile
+
 e2e-tests:
  behave
 

RELEASE.md

@@ -14,6 +14,9 @@
 * Added new `OmegaConfLoader` which uses `OmegaConf` for loading and merging configuration.
 * Added the `--conf-source` option to `kedro run`, allowing users to specify a source for project configuration for the run.
 * Added `omegaconf` syntax as option for `--params`. Keys and values can now be separated by colons or equals signs.
+* Added support for generator functions as nodes, i.e. using `yield` instead of return.
+ * Enable chunk-wise processing in nodes with generator functions.
+ * Save node outputs after every `yield` before proceeding with next chunk.
 
 ## Bug fixes and other changes
 * Fix bug where `micropkg` manifest section in `pyproject.toml` isn't recognised as allowed configuration.

dependency/requirements.txt

@@ -10,6 +10,7 @@ importlib-metadata>=3.6; python_version >= '3.8'
 importlib_metadata>=3.6, <5.0; python_version < '3.8' # The "selectable" entry points were introduced in `importlib_metadata` 3.6 and Python 3.10. Bandit on Python 3.7 relies on a library with `importlib_metadata` < 5.0
 importlib_resources>=1.3 # The `files()` API was introduced in `importlib_resources` 1.3 and Python 3.9.
 jmespath>=0.9.5, <1.0
+more_itertools~=9.0
 omegaconf~=2.3
 pip-tools~=6.12
 pluggy~=1.0.0

kedro/pipeline/node.py

@@ -9,6 +9,8 @@
 from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Union
 from warnings import warn
 
+from more_itertools import spy, unzip
+
 
 class Node:
  """``Node`` is an auxiliary class facilitating the operations required to
@@ -397,38 +399,57 @@ def _run_with_dict(self, inputs: Dict[str, Any], node_inputs: Dict[str, str]):
 
  def _outputs_to_dictionary(self, outputs):
  def _from_dict():
- if set(self._outputs.keys()) != set(outputs.keys()):
+ result, iterator = outputs, None
+ # generator functions are lazy and we need a peek into their first output
+ if inspect.isgenerator(outputs):
+ (result,), iterator = spy(outputs)
+
+ keys = list(self._outputs.keys())
+ names = list(self._outputs.values())
+ if not isinstance(result, dict):
+ raise ValueError(
+ f"Failed to save outputs of node {self}.\n"
+ f"The node output is a dictionary, whereas the "
+ f"function output is {type(result)}."
+ )
+ if set(keys) != set(result.keys()):
  raise ValueError(
  f"Failed to save outputs of node {str(self)}.\n"
- f"The node's output keys {set(outputs.keys())} do not match with "
- f"the returned output's keys {set(self._outputs.keys())}."
+ f"The node's output keys {set(result.keys())} "
+ f"do not match with the returned output's keys {set(keys)}."
  )
- return {name: outputs[key] for key, name in self._outputs.items()}
+ if iterator:
+ exploded = map(lambda x: tuple(x[k] for k in keys), iterator)
+ result = unzip(exploded)
+ else:
+ # evaluate this eagerly so we can reuse variable name
+ result = tuple(result[k] for k in keys)
+ return dict(zip(names, result))
 
  def _from_list():
- if not isinstance(outputs, (list, tuple)):
+ result, iterator = outputs, None
+ # generator functions are lazy and we need a peek into their first output
+ if inspect.isgenerator(outputs):
+ (result,), iterator = spy(outputs)
+
+ if not isinstance(result, (list, tuple)):
  raise ValueError(
  f"Failed to save outputs of node {str(self)}.\n"
  f"The node definition contains a list of "
  f"outputs {self._outputs}, whereas the node function "
- f"returned a '{type(outputs).__name__}'."
+ f"returned a '{type(result).__name__}'."
  )
- if len(outputs) != len(self._outputs):
+ if len(result) != len(self._outputs):
  raise ValueError(
  f"Failed to save outputs of node {str(self)}.\n"
- f"The node function returned {len(outputs)} output(s), "
+ f"The node function returned {len(result)} output(s), "
  f"whereas the node definition contains {len(self._outputs)} "
  f"output(s)."
  )
 
- return dict(zip(self._outputs, outputs))
-
- if isinstance(self._outputs, dict) and not isinstance(outputs, dict):
- raise ValueError(
- f"Failed to save outputs of node {self}.\n"
- f"The node output is a dictionary, whereas the "
- f"function output is not."
- )
+ if iterator:
+ result = unzip(iterator)
+ return dict(zip(self._outputs, result))
 
  if self._outputs is None:
  return {}

kedro/runner/runner.py

@@ -2,6 +2,8 @@
 implementations.
 """
 
+import inspect
+import itertools as it
 import logging
 from abc import ABC, abstractmethod
 from collections import deque
@@ -12,8 +14,9 @@
  as_completed,
  wait,
 )
-from typing import Any, Dict, Iterable, List, Set
+from typing import Any, Dict, Iterable, Iterator, List, Set
 
+from more_itertools import interleave
 from pluggy import PluginManager
 
 from kedro.framework.hooks.manager import _NullPluginManager
@@ -294,10 +297,22 @@ def run_node(
  asynchronously with threads. Defaults to False.
  session_id: The session id of the pipeline run.
 
+ Raises:
+ ValueError: Raised if is_async is set to True for nodes wrapping
+ generator functions.
+
  Returns:
  The node argument.
 
  """
+ if is_async and inspect.isgeneratorfunction(node.func):
+ raise ValueError(
+ f"Async data loading and saving does not work with "
+ f"nodes wrapping generator functions. Please make "
+ f"sure you don't use `yield` anywhere "
+ f"in node {str(node)}."
+ )
+
  if is_async:
  node = _run_node_async(node, catalog, hook_manager, session_id)
  else:
@@ -399,7 +414,21 @@ def _run_node_sequential(
  node, catalog, inputs, is_async, hook_manager, session_id=session_id
  )
 
- for name, data in outputs.items():
+ items: Iterable = outputs.items()
+ # if all outputs are iterators, then the node is a generator node
+ if all(isinstance(d, Iterator) for d in outputs.values()):
+ # Python dictionaries are ordered so we are sure
+ # the keys and the chunk streams are in the same order
+ # [a, b, c]
+ keys = list(outputs.keys())
+ # [Iterator[chunk_a], Iterator[chunk_b], Iterator[chunk_c]]
+ streams = list(outputs.values())
+ # zip an endless cycle of the keys
+ # with an interleaved iterator of the streams
+ # [(a, chunk_a), (b, chunk_b), ...] until all outputs complete
+ items = zip(it.cycle(keys), interleave(*streams))
+
+ for name, data in items:
  hook_manager.hook.before_dataset_saved(dataset_name=name, data=data)
  catalog.save(name, data)
  hook_manager.hook.after_dataset_saved(dataset_name=name, data=data)

tests/pipeline/test_node_run.py

@@ -109,8 +109,8 @@ def test_run_dict_diff_size(self, mocked_dataset):
 
 class TestNodeRunInvalidOutput:
  def test_miss_matching_output_types(self, mocked_dataset):
- pattern = r"The node output is a dictionary, whereas the function "
- pattern += r"output is not\."
+ pattern = "The node output is a dictionary, whereas the function "
+ pattern += "output is <class 'kedro.io.lambda_dataset.LambdaDataSet'>."
  with pytest.raises(ValueError, match=pattern):
  node(one_in_one_out, "ds1", dict(a="ds")).run(dict(ds1=mocked_dataset))
 

tests/runner/test_run_node.py

@@ -0,0 +1,89 @@
+import pytest
+
+from kedro.framework.hooks.manager import _NullPluginManager
+from kedro.pipeline import node
+from kedro.runner import run_node
+
+
+def generate_one():
+ yield from range(10)
+
+
+def generate_tuple():
+ for i in range(10):
+ yield i, i * i
+
+
+def generate_list():
+ for i in range(10):
+ yield [i, i * i]
+
+
+def generate_dict():
+ for i in range(10):
+ yield dict(idx=i, square=i * i)
+
+
+class TestRunGeneratorNode:
+ def test_generator_fail_async(self, mocker, catalog):
+ fake_dataset = mocker.Mock()
+ catalog.add("result", fake_dataset)
+ n = node(generate_one, inputs=None, outputs="result")
+
+ with pytest.raises(Exception, match="nodes wrapping generator functions"):
+ run_node(n, catalog, _NullPluginManager(), is_async=True)
+
+ def test_generator_node_one(self, mocker, catalog):
+ fake_dataset = mocker.Mock()
+ catalog.add("result", fake_dataset)
+ n = node(generate_one, inputs=None, outputs="result")
+ run_node(n, catalog, _NullPluginManager())
+
+ expected = [((i,),) for i in range(10)]
+ assert 10 == fake_dataset.save.call_count
+ assert fake_dataset.save.call_args_list == expected
+
+ def test_generator_node_tuple(self, mocker, catalog):
+ left = mocker.Mock()
+ right = mocker.Mock()
+ catalog.add("left", left)
+ catalog.add("right", right)
+ n = node(generate_tuple, inputs=None, outputs=["left", "right"])
+ run_node(n, catalog, _NullPluginManager())
+
+ expected_left = [((i,),) for i in range(10)]
+ expected_right = [((i * i,),) for i in range(10)]
+ assert 10 == left.save.call_count
+ assert left.save.call_args_list == expected_left
+ assert 10 == right.save.call_count
+ assert right.save.call_args_list == expected_right
+
+ def test_generator_node_list(self, mocker, catalog):
+ left = mocker.Mock()
+ right = mocker.Mock()
+ catalog.add("left", left)
+ catalog.add("right", right)
+ n = node(generate_list, inputs=None, outputs=["left", "right"])
+ run_node(n, catalog, _NullPluginManager())
+
+ expected_left = [((i,),) for i in range(10)]
+ expected_right = [((i * i,),) for i in range(10)]
+ assert 10 == left.save.call_count
+ assert left.save.call_args_list == expected_left
+ assert 10 == right.save.call_count
+ assert right.save.call_args_list == expected_right
+
+ def test_generator_node_dict(self, mocker, catalog):
+ left = mocker.Mock()
+ right = mocker.Mock()
+ catalog.add("left", left)
+ catalog.add("right", right)
+ n = node(generate_dict, inputs=None, outputs=dict(idx="left", square="right"))
+ run_node(n, catalog, _NullPluginManager())
+
+ expected_left = [((i,),) for i in range(10)]
+ expected_right = [((i * i,),) for i in range(10)]
+ assert 10 == left.save.call_count
+ assert left.save.call_args_list == expected_left
+ assert 10 == right.save.call_count
+ assert right.save.call_args_list == expected_right