Hi everyone, new to using Hamilton as part of our MLOps stack! I couldn't find a reference for passing on a column from a dataframe to another function. For example, we're loading a table from Postgres initially for user IDs, and then using those user IDs for a subsequent SQL queries to filter other dataframes by. Any help would be greatly appreciated!

Welcome! So, a few simple solutions: Just manually extract it

def original_data() -> pd.DataFrame:
    ... # return some dataframe

def id_column(original_data: pd.DataFrame) -> pd.Series:
    return original_data.user_id # or whatever

def sql_query(id_column: pd.Series, other_stuff: ...) -> pd.DataFrame:
    ...

If you have multiple columns, you might want to consider

@extract_columns

— would look something like this:

@extract_columns("user_id", ...) # one node for original_data, one node (pd.Series) for each extracted column
def original_data() -> pd.DataFrame:
    ... # return some dataframe

def sql_query(user_id: pd.Series, other_stuff: ...) -> pd.DataFrame:
    ...

This is an ergonomic/concise way of doing it — see docs here: https://hamilton.dagworks.io/en/latest/reference/decorators/extract_columns/

Ahh, okay that makes a ton of sense, I’ll take a closer look at the decorators! Thanks so much for the quick reply 🙂

Yeah! Let us know if you need anything else

@Ethan Kim @with_columns might also be useful to look at (useful once you hit trying to reuse column transformations).