load_and_join.py

#data_loaders.py
@config.when(mode="batch")
@extract_columns("budget", "age", "gender", "client_id")
def survey_results__batch(survey_results_table: str, survey_results_db: str) -> pd.DataFrame:
    """Map operation to explode survey results to all fields
    Data comes in JSON, we've grouped it into a series.
    """
    return utils.query_table(table=survey_results_table, db=survey_results_db)
  
  
@config.when(mode="batch")
def client_login_data__batch(client_login_db: str, client_login_table: str) -> pd.DataFrame:
    return utils.query_table(table=client_login_table, db=client_login_db)

# join.py
@config.when(mode="batch")
def last_logged_in__batch(client_id: pd.Series, client_login_data: pd.DataFrame) -> pd.Series:
    return pd.merge(client_id, client_login_data, left_on="client_id", right_index=True)[
        "last_logged_in"
    ]