Hey folks, can someone recommend what is the best practice in my case? The story is simple, I'm building an ML pipeline. During the data preparation I have a one hot encoding stage. It looks something like this

@save_to.pickle(path=source("encoders_path"))
def one_hot_categorical_encoding(
    team_nodes: pl.DataFrame,
    data_config: DataConfig,
) -> dict[str, dict[str, OneHotEncoder]]:

This is during the train time. During test/validation however I need to load the encoders from the file, but from reading docs I can't seem to figure out what is best way to achieve that. One options I was thinking about is to have

@load_from.pickle(path=source("encoders_path"))
def load_categorical_encoders()

and then during execution dynamically chose which one to run depending on it being training or validation. Another option I see it to have a stage that encompasses both actions loading and saving, but decides what to do according to an input parameter. I do not particularly like it, because it would be implemented outside of scope of materialisers. What would be the best way to proceed here? Thanks!

Hey, one easy solution that comes to mind is using

@config.when

and building the driver based on the case before you execute:

# module
@save_to.pickle(path=source("encoders_path"))
@config.when(stage="train")
def one_hot_categorical_encoding__train(
    ...)->dict[str, dict[str, OneHotEncoder]]:
    return ...

@load_from.pickle(path=source("encoders_path"))
@config.when(stage="test")
def one_hot_categorical_encoding__test(...)->...:
    return ...

# execution script

case = "train"
dr.Builder().with_modules(your_module).with_config({"stage":case}.build()

dr.execute(...)

In the above you then have the node

one_hot_categorical_encoding

in both case and can use it.

Ah, good idea. Will look into it, thanks!

Thanks @Jernj Frank! Yep, exactly the right way to do this IMO. Another way to do this is with a pytest fixture/overriding. E.G. you have a fixture that loads up the encoders for testing, and then you use overrides in the driver execute to pass that in and bypass the original data loading. I like config quite a bit though here — very easy to read and less to pass in.