Hi Team 👋 In our workflow it happens quite often that we have a situation where: a node takes a

PartitionedDataset

as input and

PartitionDataset

as output. So something like this:

def my_node(inputs: dict[str, Callable[[], Any]]) -> dict[str, Any]:
    results = {}
    for key, value in inputs.items():
        response = my_function(value())
        results[key] = response
    return results

Ideally, I would want: • the internal

for

loop to run in parallel • and not all the data to be in memory (because it won't fit) I've noticed that @Guillaume Tauzin mentioned a similar situation in a threat last month. My questions: • does the

DaskRunner

help in that regard? • @Guillaume Tauzin did you find a solution ??? Many thanks! Anna-Lea

Hi @Anna-Lea, I think you're referring to this thread where Guillaume asked about parallel execution. Like Elena said, there's no built in way from the Kedro side to parallelise processing within a node. That's logic you'd have to implement yourself. I personally don't have experience with the

DaskRunner

. These docs are pretty old seem to have some issues too, so not sure you can rely on that right now without a re-write.

Hi @Merel Thanks for the answer. I was a bit afraid of that response, but it's OK. I will be a bit creative.

Let me know if there's anything I can help with!

Hi @Anna-Lea, My question was a more general one and as @Elena Khaustova suggested, I was planning to open a discussion that summarizes the situation and try and propose some alternatives when I have some bandwidth. The general issue is a complex one with a lot of edge cases (nodes with 1-to-1 partitions like your use case, but could be 1-to-many, many-to-1, or 1-to-0, 0-to-1, with any other datasets as inputs and/or outputs) and a few other considerations (partitions obtained at pipeline creation time or node runtime). I think there are some solutions for your use case and it IMO it all depends at which level you wish to have parallelization. Moving to a parallel runner (

ParallelRunner

,

DaskRunner

, ...) only helps running nodes that can be ran in parallel according to the pipeline DAG. This means that in your case, using the

DaskRunner

would not help just by itself to parallelize a single node across partitions. Let me give you a quick recap of what I think your options are with some pros and cons: Parallelization at the pipeline level Here you would try and define one node per partition (or a chunk of partitions). Typically you would load your partition dataset in pipeline.py, get the partition names and iterate over them to create nodes. You will also need to create datasets for each input and output partitions. It could look like something like this: `pipeline.py`:

import os
from kedro.pipeline import pipeline, node
from .nodes import process_single_partition

def create_pipeline(**kwargs):
    # 1. Obtain partition keys (statically or by inspecting the filesystem)
    catalog_ds: PartitionedDataset = catalog.load("my_partitioned_input")
    partition_ids = catalog_ds.load().keys()

    # 2. Create a node for each partition
    nodes = []
    # You could also process partitions per chunk by modifying the loop
    for p_id in partition_ids:
        # Note: You generally need to define `input_shard_{p_id}` in the catalog
        out_pids = ... # Define the output partition key
        nodes.append(
            node(
                func=process_single_partition,
                inputs=f"input_shard_{p_id}",
                outputs=[f"output_shard_{out_p_id}", "{p_id}_is_done"],
                name=f"process_{p_id}"
            )
        )
    
    return pipeline(nodes)

with a

catalog.yml

that could look something like this (not tested):

my_partitioned_input:
  type: PartitionedDataset
  path: data/01_raw/inputs
  dataset: pandas.CSVDataset

my_partitioned_output:
  type: PartitionedDataset
  path: data/02_intermediate/outputs
  dataset: pandas.CSVDataset

"input_shard_{p_id}":
  type: pandas.CSVDataset
  filepath: {pid}


"output_shard_{p_id}":
  type: pandas.CSVDataset
  filepath: {pid}

Pros: • Further delegate parallelization to your runner (e.g.

ParallelRunner

, or

DaskRunner

) or your orchestrator (

Airflow

,

Dagster

, etc...) • High visibility: You can see exactly which partition failed in the logs/Viz. • You actually control how much data is processed per node Cons: • Complex to set up: You must generate catalog entries for every partition dynamically. • A bit hacky: I feel users are not really supposed to load partitions like that in

pipeline.py

. I am not actually sure the setup I showed even works • Need to define dummy

"{pid}_is_done"

datasets and have an extra node gathering them all to make sure they all finish before the next node is called • Use statically defined partitions or partitions obtained at pipeline creation time. The last point could be important as in the standard, sequential case, partitions are collected at node runtime, which means it would collect partitions potentially generated by previous nodes. Parallelization at the Node level You can handle the parallelism inside the node by yourself. One way I like is to use something like

joblib

to parallelize the for loop. `nodes.py`:

from joblib import Parallel, delayed
from typing import Callable, Any, Dict

def _process_single(key: str, loader: Callable[[], Any]) -> tuple[str, Any]:
    # The data is loaded only inside the worker process
    data = loader() 
    result = my_function(data)
    return key, result

def my_node(inputs: Dict[str, Callable[[], Any]]) -> Dict[str, Any]:
    # n_jobs=-1 uses all available cores, control this to control memory consumption as well
    results_list = Parallel(n_jobs=-1)(
        delayed(_process_single)(key, value) 
        for key, value in inputs.items()
    )
    return dict(results_list)

Pros: • Simple to implement: No changes to

pipeline.py

or

catalog.yml

. • Partitions are collected the native Kedro way as you would expect. Cons: • While inputs are lazy,

joblib

will return all results to the main process before the node finishes. If your output dataset for each partition is large, holding

results

in memory before Kedro writes it out might crash the worker. • No DAG visibility: Kedro sees this as one giant node; if one partition fails, the whole node fails. You can also use

Dask

inside as a

joblib

backend if you wish to: https://joblib.readthedocs.io/en/latest/auto_examples/parallel/distributed_backend_simple.html Parallelization at both the Pipeline and Node levels (Chunking) This is a hybrid approach. In

pipeline.py

, you split your partitions into "chunks" (e.g., 10 groups of 50 partitions). You create one node per chunk. Inside each node, you use the

joblib

approach from Option 2 to process that chunk's partitions in parallel. Pros: • Best of both worlds: Keeps the DAG size manageable while keeping memory usage lower than Option 2 (you only hold one chunk of results in memory at a time). Cons: • Highest complexity to implement (requires both dynamic pipeline logic and internal node parallelism). Hope this helps!

Hi @Guillaume Tauzin; Thanks for this amazing detailed answer!! I'll need some time to implement this, but you have given me a lot of food for thoughts and trials One point though about the

joblib

option, as I'm already using

dask.distributed

in my pipelines, I'm wondering if both can cohabit on the same machine ... To be tested probably Cheers

You're welcome! I was wondering the same and my guess is that it would work. To be tested indeed! If you do, please let me know here. I am curious to find out :)

Hi @Guillaume Tauzin A quick follow up on our discussion. We opted for the second approach for now using

joblib

. It comes with this feature

return_as="generator"

which together with a

yield

makes it sort of possible to limit a bit the memory load. Direct example:

from joblib import Parallel, delayed

def dispatched_node(inputs: dict[str, Callable[[], Any]]) -> Generator:
    result_generator = Parallel(return_as="generator")(delayed(my_function)(entry()) for entry in inputs.values())
    yield from result_generator

Together with dask and some chunking it's working (at least in my use case). Thanks again for your suggestions.

Hi @Anna-Lea Thanks a lot for letting me know! Love the generator trick and happy it all worked :)