feat(airflow): Add grouping by namespace option to kedro-airflow

kedro-airflow/kedro_airflow/grouping.py

@@ -113,3 +113,37 @@ def dfs(cur_node_name: str, component: int) -> None:
                 group_dependencies[new_name_parent].append(new_name_child)
 
     return group_to_seq, group_dependencies
+
+
+def group_by_namespace(
+    pipeline: Pipeline,
+) -> tuple[dict[str, list[Node]], dict[str, list[str]]]:
+    """
+    Groups nodes based on their namespace using Pipeline's grouped_nodes_by_namespace property.
+    Non-namespaced nodes are assigned to a default namespace.
+    """
+    nodes_by_namespace: dict[str, list[Node]] = {}
+    dependencies_by_namespace: dict[str, list[str]] = {}
+
+    grouped_nodes = pipeline.grouped_nodes_by_namespace
+
+    node_to_namespace = {}
+
+    for group_name, group_info in grouped_nodes.items():
+        ns = group_info["name"] if group_info["type"] == "namespace" else "__default__"
+        nodes_by_namespace.setdefault(ns, [])
+        dependencies_by_namespace.setdefault(ns, [])
+        nodes_by_namespace[ns].extend(group_info["nodes"])
+        for node in group_info["nodes"]:
+            node_to_namespace[node] = ns
+
+    for ns, nodes in nodes_by_namespace.items():
+        dependent_namespaces = set()
+        for node in nodes:
+            for parent in pipeline.node_dependencies.get(node, []):
+                parent_ns = node_to_namespace.get(parent, "__default__")
+                if parent_ns != ns:
+                    dependent_namespaces.add(parent_ns)
+        dependencies_by_namespace[ns] = list(dependent_namespaces)
+
+    return nodes_by_namespace, dependencies_by_namespace

kedro-airflow/kedro_airflow/plugin.py

@@ -121,6 +121,13 @@ def _get_pipeline_config(config_airflow: dict, params: dict, pipeline_name: str)
     help="Group nodes with at least one MemoryDataset as input/output together, "
     "as they do not persist between Airflow operators.",
 )
+@click.option(
+    "-gn",
+    "--group-by-namespace",
+    is_flag=True,
+    default=False,
+    help="Groups nodes based on their namespace using Kedro's grouped_nodes_by_namespace property.",
+)
 @click.option(
     "--tags",
     type=str,
@@ -149,6 +156,7 @@ def create(  # noqa: PLR0913, PLR0912
     target_path,
     jinja_file,
     group_in_memory,
+    group_by_namespace,
     tags,
     params,
     conf_source,
@@ -218,6 +226,8 @@ def create(  # noqa: PLR0913, PLR0912
             # topological sort order obtained from pipeline.nodes, see group_memory_nodes()
             # implementation
             nodes, dependencies = group_memory_nodes(context.catalog, pipeline)
+        elif group_by_namespace:
+            nodes, dependencies = group_by_namespace(pipeline)
         else:
             # To keep the order of nodes and dependencies deterministic - nodes are
             # iterated in the topological sort order obtained from pipeline.nodes and

kedro-airflow/tests/test_node_grouping.py

@@ -7,7 +7,11 @@
 from kedro.pipeline import Pipeline, node
 from kedro.pipeline.modular_pipeline import pipeline as modular_pipeline
 
-from kedro_airflow.grouping import _is_memory_dataset, group_memory_nodes
+from kedro_airflow.grouping import (
+    _is_memory_dataset,
+    group_by_namespace,
+    group_memory_nodes,
+)
 
 
 class TestDataset(AbstractDataset):
@@ -87,6 +91,63 @@ def identity_one_to_one(x):
     )
 
 
+def mock_kedro_pipeline_with_namespaces() -> Pipeline:
+    def identity_one_to_one(x):
+        return x
+
+    return Pipeline(
+        [
+            node(
+                func=identity_one_to_one,
+                inputs="ds1",
+                outputs="ds2",
+                name="f1",
+                namespace="namespace1",
+            ),
+            node(
+                func=lambda x: (x, x),
+                inputs="ds2",
+                outputs=["ds3", "ds4"],
+                name="f2",
+                namespace="namespace1",
+            ),
+            node(
+                func=identity_one_to_one,
+                inputs="ds3",
+                outputs="ds5",
+                name="f3",
+                namespace="namespace2",
+            ),
+            node(
+                func=identity_one_to_one,
+                inputs="ds3",
+                outputs="ds6",
+                name="f4",
+                namespace="namespace2",
+            ),
+            node(
+                func=identity_one_to_one,
+                inputs="ds4",
+                outputs="ds8",
+                name="f5",
+                namespace="namespace3",
+            ),
+            node(
+                func=identity_one_to_one,
+                inputs="ds6",
+                outputs="ds7",
+                name="f6",
+            ),
+            node(
+                func=lambda x, y: x,
+                inputs=["ds3", "ds6"],
+                outputs="ds9",
+                name="f7",
+            ),
+        ]
+    )
+
+
 @pytest.mark.parametrize(
     "all_nodes,memory_nodes,expected_nodes,expected_dependencies",
     [
@@ -166,3 +227,35 @@ def test_is_memory_dataset(
             assert _is_memory_dataset(mock_catalog, node_name)
         else:
             assert not _is_memory_dataset(mock_catalog, node_name)
+
+
+@pytest.mark.parametrize(
+    "pipeline, expected_nodes, expected_dependencies",
+    [
+        (
+            mock_kedro_pipeline_with_namespaces(),
+            {
+                "__default__": ["f6", "f7"],
+                "namespace1": ["namespace1.f1", "namespace1.f2"],
+                "namespace2": ["namespace2.f3", "namespace2.f4"],
+                "namespace3": ["namespace3.f5"],
+            },
+            {
+                "__default__": ["namespace1", "namespace2"],
+                "namespace1": [],
+                "namespace2": ["namespace1"],
+                "namespace3": ["namespace1"],
+            },
+        ),
+    ],
+)
+def test_group_by_namespace(pipeline, expected_nodes, expected_dependencies):
+    """Test grouping of nodes by namespace."""
+    nodes_by_namespace, dependencies_by_namespace = group_by_namespace(pipeline)
+
+    nodes_by_namespace = {
+        ns: [node.name for node in nodes] for ns, nodes in nodes_by_namespace.items()
+    }
+
+    assert nodes_by_namespace == expected_nodes
+    assert dependencies_by_namespace == expected_dependencies