👋 Hello, team!, I have two functions (

get_dataset1

and

get_dataset2

) that are essentially doing the same and would like to DRY and merge them into one function

get_dataset(dataset:str)

with one parameter (

dataset

):

@extract_fields(
    {
        "dataset1_table1": pl.DataFrame,
        "dataset1_table2": pl.DataFrame,
        "dataset1_table3": pl.DataFrame,
    }
)
def get_dataset1() -> dict:
    ...
    return {
        "dataset1_table1": dataset1_table1,
        "dataset1_table2": dataset1_table2,
        "dataset1_table3": dataset1_table3,
    }

# second function

@extract_fields(
    {
        "dataset2_table1": pl.DataFrame,
        "dataset2_table2": pl.DataFrame,
        "dataset2_table3": pl.DataFrame,
    }
)
def get_dataset2() -> dict:
    ...
    return {
        "dataset2_table1": dataset2_table1,
        "dataset2_table2": dataset2_table2,
        "dataset2_table3": dataset2_table3,
    }

I would like to achieve something like this, one function with

dataset

as a parameter:

@extract_fields(
    {
        "{dataset}_table1": pl.DataFrame,
        "{dataset}_table2": pl.DataFrame,
        "{dataset}_table3": pl.DataFrame,
    }
)
def get_dataset(dataset: str) -> dict:
    # dataset dependent logic
    return {
        "{dataset}_table1": dataset_table1,
        "{dataset}_table2": dataset_table2,
        "{dataset}_table3": dataset_table3,
    }

Is this possible?

Hi @Yaser Martínez Palenzuela! What kind of problem are you working on? What you have in mind is very much doable! Although, we warn about "too much dynamism" because it can hurt observability and debugging. You seem familiar with Hamilton, so I'll give you a high-level rundown: ## Initial code (non-DRY) does the following • You wrote the Python function

get_dataset1()

with the decorator

@extract_fields(dataset1_table1, ...)

• When you build the

Driver

, it loads function definitions to build the nodes and dataflow, resulting in

get_dataset1

,

dataset1_table1

,

dataset1_table2

,

dataset1_table3

. (you know the nodes are generated because you can visualize the DAG) • When you do

Driver.execute()

data flows through the DAG to produce results ## Suggested refactor To work, you would need to • Write the generic function

get_dataset(dataset_name: str)

with a "templated"

@extract_fields

• When you build the

Driver

, "resolve" the

@extract_fields

with the right dataset to be able to generate the graph, which will have

get_dataset

and

dataset1_table1

or

dataset2_table1

(etc.) • Then, just like the first scenario,

Driver.execute()

sends data through the DAG. ##

@resolve

to the rescue The construct for that is hamilton.function_modifiers.resolve or the shorthand hamilton.function_modifiers.resolve_from_config, which allows to add decorators / function modifiers to a function dynamically. It would look something like this

from hamilton.function_modifiers import resolve_from_config, extract_fields

@resolve_from_config(
  decorate_with=lambda dataset: extract_fields({
        "{dataset}_table1": pl.DataFrame,
        "{dataset}_table2": pl.DataFrame,
        "{dataset}_table3": pl.DataFrame,
  })
)
def get_dataset(dataset: str) -> dict:
  ...
  return {
        "{dataset}_table1": dataset_table1,
        "{dataset}_table2": dataset_table2,
        "{dataset}_table3": dataset_table3,
    }

This pattern can be very powerful and is well-liked for feature stores and data platforms. A community member talks about it in this . Feel free to ping me for more info!

wow this is one of the best responses I've ever gotten! Thanks, @Thierry Jean

@Yaser Martínez Palenzuela to clarify a bit more, do you want to run them in the same DAG? or would they be run separately?

hmm, they run on the same DAG. I guess I can only use one specific config at a time right?

are you then joining the outputs? or are they very separate? or?

dataset1_table1

,

dataset1_table2

,

dataset1_table3

get joined in another node as a dataframe

dataset1

and the same for

dataset2

something like:

def dataset1(dataset1_table1, dataset1_table2 , dataset1_table3) -> pl.DataFrame
    ...
    return pl.concat([dataset1_table1, dataset1_table2 , dataset1_table3])

is it always the same three? Are they used by anything else? or? Our advice is to keep things simple, e.g. any reason to not just have a single function for each dataset?

def dataset1(dataset1: str) -> pl.DataFrame
    # load data sets ...
    return pl.concat([dataset1_table1, dataset1_table2 , dataset1_table3])

I mean thats why I have now, it is just that both functions

dataset1

and

dataset2

have like 40 lines of the same code

Sure. You could have a helper function?

def _helper( ... ) -> ...

def dataset1(...) -> ...
  df = _helper(...)
  return df

def dataset2(...)-> ...
  df = _helper(...)
  return df

ok will try something like that thanks

functions prefixed with

_

aren’t made into nodes

good tip

@Yaser Martínez Palenzuela just to mention if you don’t like that, there are other ways — but their use depends on downstream situations and how dynamic things are or are not.

Thank you Stefan, your suggestion works great for me. I'll reach out if I need more dynamis. Thanks for the excellent support