extract_columns

# functions.py - declare and link your transformations as functions....
import pandas as pd
from hamilton.htypes import Parallelizable, Collect
from hamilton.function_modifiers import extract_columns

def a(input: pd.Series) -> Parallelizable[int]:
    for x in range(0, 10):
        yield x

def b(a: int) -> int:
    return a * 2

def c(a: int, b: int) -> int:
    return a * 3 + b * 2

def d(c: int) -> int:
    return c ** 3

@extract_columns("q", "r")
def e(d: Collect[int]) -> pd.DataFrame:
    return pd.DataFrame({"q": pd.Series(list(d)), "r": pd.Series(list(d))})

# And run them!
import functions
from hamilton import driver, base
from hamilton.execution import executors
# create a driver
dr = ( 
  driver.Builder()
    .with_modules(functions)
    .enable_dynamic_execution(allow_experimental_mode=True)
    .with_local_executor(executors.SynchronousLocalTaskExecutor())
    .with_remote_executor(executors.SynchronousLocalTaskExecutor())
    .with_adapters(base.PandasDataFrameResult())
    .build()
)
# execute it
result = dr.execute(
  ['q', 'r'], 
  inputs={'input': pd.Series([1, 2, 3, 4, 5])}
)
print(result)
# visualize it
gr = dr.display_all_functions(orient="TB", show_legend=False)

Driver

Traceback (most recent call last):
    .build()
     ^^^^^^^
  File "C:\codebase\ice-irm2\irm-venv\Lib\site-packages\hamilton\driver.py", line 1791, in build
    return Driver(
           ^^^^^^^
  File "C:\codebase\ice-irm2\irm-venv\Lib\site-packages\hamilton\driver.py", line 397, in __init__
    raise e
  File "C:\codebase\ice-irm2\irm-venv\Lib\site-packages\hamilton\driver.py", line 384, in __init__
    self.graph = graph.FunctionGraph.from_modules(*modules, config=config, adapter=adapter)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\codebase\ice-irm2\irm-venv\Lib\site-packages\hamilton\graph.py", line 676, in from_modules
    nodes = create_function_graph(*modules, config=config, adapter=adapter)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\codebase\ice-irm2\irm-venv\Lib\site-packages\hamilton\graph.py", line 178, in create_function_graph
    nodes = update_dependencies(nodes, adapter, reset_dependencies=False)  # no dependencies
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\codebase\ice-irm2\irm-venv\Lib\site-packages\hamilton\graph.py", line 139, in update_dependencies
    add_dependency(n, node_name, nodes, param_name, param_type, adapter)
  File "C:\codebase\ice-irm2\irm-venv\Lib\site-packages\hamilton\graph.py", line 61, in add_dependency
    types_do_match = types_match(param_type, required_node.type)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\codebase\ice-irm2\irm-venv\Lib\site-packages\hamilton\htypes.py", line 130, in types_match
    elif custom_subclass_check(required_node_type, param_type):
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\codebase\ice-irm2\irm-venv\Lib\site-packages\hamilton\htypes.py", line 69, in custom_subclass_check
    (requested_type_arg,) = _get_args(requested_type)
    ^^^^^^^^^^^^^^^^^^^^^
ValueError: too many values to unpack (expected 1)

Collect

@extract_columns

pd.DataFrame

@extract_columns("q", "r")
def e(d: Collect[int]) -> pd.DataFrame:
    return pd.DataFrame({"q": pd.Series(list(d)), "r": pd.Series(list(d))})

@config.when(data_product="api")
def commodity_df(positions_validated: pd.DataFrame) -> Parallelizable[Any, pd.DataFrame]:
    """Extract commodity data from formatted positions."""
    for cmdty, df in positions_validated.groupby("commodity"):
        yield cmdty, df


@config.when(data_product="api")
def formatted_positions__api(commodity_df: pd.DataFrame) -> dict: ...

Parallelizable[Any, pd.DataFrame]