apache-hamilton 1.90.0.dev0__py3-none-any.whl

hamilton/function_modifiers/recursive.py

@@ -0,0 +1,908 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import abc
+import inspect
+import sys
+import typing
+from collections import defaultdict
+from collections.abc import Callable, Collection
+from types import ModuleType
+from typing import Any, TypedDict
+
+_sys_version_info = sys.version_info
+_version_tuple = (_sys_version_info.major, _sys_version_info.minor, _sys_version_info.micro)
+
+if _version_tuple < (3, 11, 0):
+    from typing_extensions import NotRequired
+else:
+    from typing import NotRequired
+
+# Copied this over from function_graph
+# TODO -- determine the best place to put this code
+from hamilton import graph_utils, node
+from hamilton.function_modifiers import base, dependencies
+from hamilton.function_modifiers.base import InvalidDecoratorException, NodeTransformer
+from hamilton.function_modifiers.dependencies import (
+    LiteralDependency,
+    ParametrizedDependency,
+    UpstreamDependency,
+)
+
+
+def assign_namespace(node_name: str, namespace: str) -> str:
+    return f"{namespace}.{node_name}"
+
+
+def derive_type(dependency: dependencies.LiteralDependency):
+    """Quick hack to derive the type of a static dependency.
+    We might want to consider the type provided by the function that needs it.
+    Or we can use the subclass checker/whatnot in function_graph
+    (althoyugh we'll want to move it out)
+
+    :param dependency: Dependency on which
+    :return: The type of the dependency
+    """
+    return type(dependency.value)
+
+
+def create_identity_node(
+    from_: str, typ: type[type], name: str, namespace: tuple[str, ...], tags: dict[str, Any]
+) -> node.Node:
+    """Creates an identity node -- this passes through the exact
+    value returned by the upstream node.
+
+    :param from_: Source node
+    :param typ: Type of the input node
+    :param name: Name of the final node to create
+    :param namespace: Namespace of the node
+    :return: A node that simply copies the source node
+    """
+
+    def identity(**kwargs):
+        return list(kwargs.values())[0]  # Maybe come up with a better way to do this
+
+    return node.Node(
+        name=name,
+        typ=typ,
+        doc_string="",
+        callabl=identity,
+        input_types={from_: typ},
+        namespace=namespace,
+        tags=tags,
+        # TODO -- add tags?
+    )
+
+
+def extract_all_known_types(nodes: Collection[node.Node]) -> dict[str, type[type]]:
+    """Extracts all known types from a set of nodes given the dependencies.
+    We have to do this as we don't know the dependency types at compile-time of
+    upstream nodes. That said, this is only used for guessing dependency types of
+    identity nodes. In which case, we probably want some sort of sentinel "pass-through"
+    dependency type that handles this better. But, for now, we'll derive it from the
+    dependencies we've seen.
+
+    :param nodes: nodes to look through for dependencies
+    :return: A dictionary of all known types.
+    """
+    observed_types = {}
+    for node_ in nodes:
+        for dep_name, (type_, _) in node_.input_types.items():
+            observed_types[dep_name] = type_
+    return observed_types
+
+
+def create_static_node(
+    typ: type, name: str, value: Any, namespace: tuple[str, ...], tags: dict[str, Any]
+) -> node.Node:
+    """Utility function to create a static node -- this helps us bridge nodes together.
+
+    :param typ: Type of the node to create
+    :param name: Name of the node to create
+    :param value: Value that the node's function always returns
+    :param namespace: Namespace of the node
+    :return: The instantiated static node
+    """
+
+    def node_fn(_value=value):
+        return _value
+
+    return node.Node(
+        name=name, typ=typ, callabl=node_fn, input_types={}, namespace=namespace, tags=tags
+    )
+
+
+def _validate_config_inputs(config: dict[str, Any], inputs: dict[str, Any]):
+    """Validates that the inputs specified in the config are valid.
+
+    :param original_config: Original configuration
+    :return: None
+    """
+    # TODO -- implement this
+    shared_keys = set(config.keys()).intersection(set(inputs.keys()))
+    if shared_keys:
+        raise InvalidDecoratorException(
+            f"Config keys {shared_keys} are shared with inputs. This is not allowed."
+            f"Instead, please specify the inputs you need *just* as part of the config. "
+            f"That way, you only write them once! Or, if you don't need them as a config item,"
+            f"just use them in inputs."
+        )
+    for key, value in inputs.items():
+        if not isinstance(value, (UpstreamDependency, LiteralDependency)):
+            raise InvalidDecoratorException(
+                f"Input {key} must be either an UpstreamDependency or a LiteralDependency ,"
+                f" not {type(value)}."
+            )
+
+
+def _resolve_subdag_configuration(
+    configuration: dict[str, Any], fields: dict[str, Any], function_name: str
+) -> dict[str, Any]:
+    """Resolves the configuration for a subdag.
+
+    :param configuration: the Hamilton configuration
+    :param fields: the fields passed to the subdag decorator
+    :return: resolved configuration to use for this subdag.
+    """
+    sources_to_map = {}
+    values_to_include = {}
+    for key, value in fields.items():
+        if isinstance(value, dependencies.ConfigDependency):
+            sources_to_map[key] = value.source
+        elif isinstance(value, dependencies.LiteralDependency):
+            values_to_include[key] = value.value
+        elif isinstance(value, (dependencies.GroupedDependency, dependencies.SingleDependency)):
+            raise InvalidDecoratorException(
+                f"`{value}` is not allowed in the config= part of the subdag decorator. "
+                "Please use `configuration()` or `value()` or literal python values."
+            )
+    plain_configs = {
+        k: v for k, v in fields.items() if k not in sources_to_map and k not in values_to_include
+    }
+    resolved_config = dict(configuration, **plain_configs, **values_to_include)
+
+    # override any values from sources
+    for key, source in sources_to_map.items():
+        try:
+            resolved_config[key] = resolved_config[source]
+        except KeyError as e:
+            raise InvalidDecoratorException(
+                f"Source {source} was not found in the configuration. This is required for the {function_name} subdag."
+            ) from e
+    return resolved_config
+
+
+NON_FINAL_TAGS = {NodeTransformer.NON_FINAL_TAG: True}
+
+
+class subdag(base.NodeCreator):
+    """The `@subdag` decorator enables you to rerun components of your DAG with varying parameters.
+    That is, it enables you to "chain" what you could express with a driver into a single DAG.
+
+    That is, instead of using Hamilton within itself:
+
+    .. code-block:: python
+
+        def feature_engineering(source_path: str) -> pd.DataFrame:
+            '''You could recursively use Hamilton within itself.'''
+            dr = driver.Driver({}, feature_modules)
+            df = dr.execute(["feature_df"], inputs={"path": source_path})
+            return df
+
+    You instead can use the `@subdag` decorator to do the same thing, with the added benefit of visibility into the\
+    whole DAG:
+
+    .. code-block:: python
+
+        @subdag(
+            feature_modules,
+            inputs={"path": source("source_path")},
+            config={}
+        )
+        def feature_engineering(feature_df: pd.DataFrame) -> pd.DataFrame:
+            return feature_df
+
+
+    Note that this is immensely powerful -- if we draw analogies from Hamilton to standard procedural programming \
+    paradigms, we might have the following correspondence:
+
+    - `config.when` + friends -- `if/else` statements
+    - `parameterize`/`extract_columns` -- `for` loop
+    - `does` -- effectively macros
+
+    And so on. `@subdag` takes this one step further:
+
+    - `@subdag` -- subroutine definition
+
+    E.G. take a certain set of nodes, and run them with specified parameters.
+
+    @subdag declares parameters that are outputs of its subdags. Note that, if you want to use outputs of other
+    components of the DAG, you can use the `external_inputs` parameter to declare the parameters that do *not* come
+    from the subDAG.
+
+    Why might you want to use this? Let's take a look at some examples:
+
+    1. You have a feature engineering pipeline that you want to run on multiple datasets. If its exactly the same, \
+    this is perfect. If not, this works perfectly as well, you just have to utilize different functions in each or \
+    the `config.when` + `config` parameter to rerun it.
+    2. You want to train multiple models in the same DAG that share some logic (in features or training) -- this \
+    allows you to reuse and continually add more.
+    3. You want to combine multiple similar DAGs (e.g. one for each business line) into one so you can build a \
+    cross-business line model.
+
+    This basically bridges the gap between the flexibility of non-declarative pipelining frameworks with the \
+    readability/maintainability of declarative ones.
+    """
+
+    def __init__(
+        self,
+        *load_from: ModuleType | Callable,
+        inputs: dict[str, ParametrizedDependency] = None,
+        config: dict[str, Any] = None,
+        namespace: str = None,
+        final_node_name: str = None,
+        external_inputs: list[str] = None,
+    ):
+        """Adds a subDAG to the main DAG.
+
+        :param load_from: The functions that will be used to generate this subDAG.
+        :param inputs: Parameterized dependencies to inject into all sources of this subDAG.
+            This should *not* be an intermediate node in the subDAG.
+        :param config: A configuration dictionary for *just* this subDAG. Note that this passed in
+            value takes precedence over the DAG's config.
+        :param namespace: Namespace with which to prefix nodes. This is optional -- if not included,
+            this will default to the function name.
+        :param final_node_name: Name of the final node in the subDAG. This is optional -- if not included,
+            this will default to the function name.
+        :param external_inputs: Parameters in the function that are not produced by the functions
+            passed to the subdag. This is useful if you want to perform some logic with other inputs
+            in the subdag's processing function. Note that this is currently required to
+            differentiate and clarify the inputs to the subdag.
+
+        """
+        self.subdag_functions = subdag.collect_functions(load_from)
+        self.inputs = inputs if inputs is not None else {}
+        self.config = config if config is not None else {}
+        self.external_inputs = external_inputs if external_inputs is not None else []
+        _validate_config_inputs(self.config, self.inputs)
+        self.namespace = namespace
+        self.final_node_name = final_node_name
+
+    @staticmethod
+    def collect_functions(
+        load_from: Collection[ModuleType] | Collection[Callable],
+    ) -> list[Callable]:
+        """Utility function to collect functions from a list of callables/modules.
+
+        :param load_from: A list of callables or modules to load from
+        :return: a list of callables to use to create a DAG.
+        """
+        if len(load_from) == 0:
+            raise ValueError(f"No functions were passed to {subdag.__name__}(load_from=...)")
+        out = []
+        for item in load_from:
+            if isinstance(item, Callable):
+                out.append(item)
+            out.extend(
+                [function for _, function in graph_utils.find_functions(function_module=item)]
+            )
+        return out
+
+    @staticmethod
+    def collect_nodes(config: dict[str, Any], subdag_functions: list[Callable]) -> list[node.Node]:
+        nodes = []
+        for fn in subdag_functions:
+            for node_ in base.resolve_nodes(fn, config):
+                nodes.append(node_.copy_with(tags={**node_.tags, **NON_FINAL_TAGS}))
+        return nodes
+
+    def _create_additional_static_nodes(
+        self, nodes: Collection[node.Node], namespace: str
+    ) -> Collection[node.Node]:
+        # These already have the namespace on them
+        # This allows us to inject values into the replayed subdag
+        node_types = extract_all_known_types(nodes)
+        out = []
+        for key, value in self.inputs.items():
+            # TODO -- fix type derivation. Currently we don't use the specified type as we don't
+            #  really know what it should be...
+            new_node_name = assign_namespace(key, namespace)
+            if value.get_dependency_type() == dependencies.ParametrizedDependencySource.LITERAL:
+                out.append(
+                    create_static_node(
+                        typ=derive_type(value),
+                        name=key,
+                        value=value.value,
+                        namespace=(namespace,),
+                        tags=NON_FINAL_TAGS,
+                    )
+                )
+            elif value.get_dependency_type() == dependencies.ParametrizedDependencySource.UPSTREAM:
+                if new_node_name not in node_types:
+                    continue
+                out.append(
+                    create_identity_node(
+                        from_=value.source,
+                        typ=node_types[new_node_name],
+                        name=key,
+                        namespace=(namespace,),
+                        tags=NON_FINAL_TAGS,
+                    )
+                )
+        for key, value in self.config.items():
+            out.append(
+                create_static_node(
+                    typ=type(value),
+                    name=key,
+                    value=value,
+                    namespace=(namespace,),
+                    tags=NON_FINAL_TAGS,
+                )
+            )
+        return out
+
+    @staticmethod
+    def add_namespace(
+        nodes: list[node.Node],
+        namespace: str,
+        inputs: dict[str, Any] = None,
+        config: dict[str, Any] = None,
+    ) -> list[node.Node]:
+        """Utility function to add a namespace to nodes.
+
+        :param nodes:
+        :return:
+        """
+        inputs = inputs if inputs is not None else {}
+        config = config if config is not None else {}
+        new_nodes = []
+        new_name_map = {}
+        # First pass we validate + collect names so we can alter dependencies
+        for node_ in nodes:
+            new_name = assign_namespace(node_.name, namespace)
+            new_name_map[node_.name] = new_name
+        for dep, _value in inputs.items():
+            # We create nodes for both namespace assignment and source assignment
+            # Why? Cause we need unique parameter names, and with source() some can share params
+            new_name_map[dep] = assign_namespace(dep, namespace)
+
+        for dep, _value in config.items():
+            new_name_map[dep] = assign_namespace(dep, namespace)
+
+        # Reassign sources
+        for node_ in nodes:
+            # This is not perfect -- we might get strangeness if its dynamically generated
+            # that said, it should work
+            is_async = inspect.iscoroutinefunction(node_.callable)
+            new_name = new_name_map[node_.name]
+            kwarg_mapping = {
+                (new_name_map[key] if key in new_name_map else key): key
+                for key in node_.input_types
+            }
+
+            # Map of argument in function to source, can't be the other way
+            # around as sources can potentially serve multiple destinations (with the source()) decorator
+            def fn(
+                _callabl=node_.callable,
+                _kwarg_mapping=dict(kwarg_mapping),  # noqa: B006
+                _new_name=new_name,
+                _new_name_map=dict(new_name_map),  # noqa: B006
+                **kwargs,
+            ):
+                new_kwargs = {_kwarg_mapping[kwarg]: value for kwarg, value in kwargs.items()}
+                return _callabl(**new_kwargs)
+
+            async def async_fn(
+                _callabl=node_.callable,
+                _kwarg_mapping=dict(kwarg_mapping),  # noqa: B006
+                _new_name=new_name,
+                _new_name_map=dict(new_name_map),  # noqa: B006
+                **kwargs,
+            ):
+                new_kwargs = {_kwarg_mapping[kwarg]: value for kwarg, value in kwargs.items()}
+                return await _callabl(**new_kwargs)
+
+            new_input_types = {
+                dep: node_.input_types[original_dep] for dep, original_dep in kwarg_mapping.items()
+            }
+            fn_to_use = async_fn if is_async else fn
+
+            new_nodes.append(
+                node_.copy_with(input_types=new_input_types, name=new_name, callabl=fn_to_use)
+            )
+        return new_nodes
+
+    def add_final_node(self, fn: Callable, node_name: str, namespace: str):
+        """
+
+        :param fn:
+        :return:
+        """
+        is_async = inspect.iscoroutinefunction(fn)  # determine if its async
+        node_ = node.Node.from_fn(fn)
+        namespaced_input_map = {
+            (assign_namespace(key, namespace) if key not in self.external_inputs else key): key
+            for key in node_.input_types
+        }
+
+        new_input_types = {
+            (assign_namespace(key, namespace) if key not in self.external_inputs else key): value
+            for key, value in node_.input_types.items()
+        }
+
+        def new_function(**kwargs):
+            kwargs_without_namespace = {
+                namespaced_input_map[key]: value for key, value in kwargs.items()
+            }
+            # Have to translate it back to use the kwargs the fn is expecting
+            return fn(**kwargs_without_namespace)
+
+        async def async_function(**kwargs):
+            return await new_function(**kwargs)
+
+        fn_to_use = async_function if is_async else new_function
+
+        return node_.copy_with(name=node_name, input_types=new_input_types, callabl=fn_to_use)
+
+    def _derive_namespace(self, fn: Callable) -> str:
+        """Utility function to derive a namespace from a function.
+
+        :param fn: Function we're decorating.
+        :return: The function we're outputting.
+        """
+        return fn.__name__ if self.namespace is None else self.namespace
+
+    def _derive_name(self, fn: Callable) -> str:
+        """Utility function to derive a name from a function.
+        The user will be able to likely pass this in as an override, but
+        we have not exposed it yet.
+
+        :param fn: Function we're decorating.
+        :return: The function we're outputting.
+        """
+        return fn.__name__ if self.final_node_name is None else self.final_node_name
+
+    def generate_nodes(self, fn: Callable, configuration: dict[str, Any]) -> Collection[node.Node]:
+        # Resolve all nodes from passed in functions
+        # if self.config has configuration() or value() in it, we need to resolve it
+        resolved_config = _resolve_subdag_configuration(configuration, self.config, fn.__name__)
+        # resolved_config = dict(configuration, **self.config)
+        nodes = self.collect_nodes(config=resolved_config, subdag_functions=self.subdag_functions)
+        # Derive the namespace under which all these nodes will live
+        namespace = self._derive_namespace(fn)
+        final_node_name = self._derive_name(fn)
+        # Rename them all to have the right namespace
+        nodes = self.add_namespace(nodes, namespace, self.inputs, self.config)
+        # Create any static input nodes we need to translate
+        nodes += self._create_additional_static_nodes(nodes, namespace)
+        # Add the final node that does the translation
+        nodes += [self.add_final_node(fn, final_node_name, namespace)]
+        return nodes
+
+    def _validate_parameterization(self):
+        invalid_values = []
+        for _key, value in self.inputs.items():
+            if not isinstance(value, dependencies.ParametrizedDependency):
+                invalid_values.append(value)
+        if invalid_values:
+            raise ValueError(
+                f"Parameterization using the following values is not permitted -- "
+                f"must be either source() or value(): {invalid_values}"
+            )
+
+    def validate(self, fn):
+        """Validates everything we can before we create the subdag.
+
+        :param fn: Function that this decorates
+        :raises InvalidDecoratorException: if this is not a valid decorator
+        """
+
+        self._validate_parameterization()
+
+    def required_config(self) -> list[str] | None:
+        """Currently we do not filter for subdag as we do not *statically* know what configuration
+        is required. This is because we need to parse the function so that we can figure it out,
+        and that is not available at the time that we call required_config. We need to think about
+        the best way to do this, but its likely that we'll want to allow required_config to consume
+        the function itself, and pass it in when its called with that.
+
+        That said, we don't have sufficient justification to do that yet, so we're just going to
+        return None for now, meaning that it has access to all configuration variables.
+
+        :return:
+        """
+        return None
+
+
+class SubdagParams(TypedDict):
+    inputs: NotRequired[dict[str, ParametrizedDependency]]
+    config: NotRequired[dict[str, Any]]
+    external_inputs: NotRequired[list[str]]
+
+
+class parameterized_subdag(base.NodeCreator):
+    """parameterized subdag is when you want to create multiple subdags at one time.
+    Why might you want to do this?
+
+    1. You have multiple data sets you want to run the same feature engineering pipeline on.
+    2. You want to run some sort of optimization routine with a variety of results
+    3. You want to run some sort of pipeline over slightly different configuration (E.G. region/business line)
+
+    Note that this really is just syntactic sugar for creating multiple subdags, just as `@parameterize
+    is syntactic sugar for creating multiple nodes from a function. That said, it is common that you
+    won't know what you want until compile time (E.G. when you have the config available), so this
+    decorator along with the `@resolve` decorator is a good way to make that feasible. Note that
+    we are getting into *advanced* Hamilton here -- we don't recommend starting with this. In fact,
+    we generally recommend repeating subdags multiple times if you don't have too many. That said,
+    that can get cumbersome if you have a lot, so this decorator is a good way to help with that.
+
+    Let's take a look at an example:
+
+    .. code-block:: python
+
+        @parameterized_subdag(
+            feature_modules,
+            from_datasource_1={"inputs" : {"data" : value("datasource_1.csv")}},
+            from_datasource_2={"inputs" : {"data" : value("datasource_2.csv")}},
+            from_datasource_3={
+                "inputs" : {"data" : value("datasource_3.csv")},
+                "config" : {"filter" : "only_even_client_ids"}
+            }
+        )
+        def feature_engineering(feature_df: pd.DataFrame) -> pd.DataFrame:
+            return feature_df
+
+    This is (obviously) contrived, but what it does is create three subdags, each with a different
+    data source. The third one also applies a configuration to that subdags. Note that we can also
+    pass in inputs/config to the decorator itself, which will be applied to all subdags.
+
+    This is effectively the same as the example above.
+
+    .. code-block:: python
+
+        @parameterized_subdag(
+            feature_modules,
+            inputs={"data" : value("datasource_1.csv")},
+            from_datasource_1={},
+            from_datasource_2={
+                    "inputs" : {"data" : value("datasource_2.csv")}
+            },
+            from_datasource_3={
+                    "inputs" : {"data" : value("datasource_3.csv")},
+                    "config" : {"filter" : "only_even_client_ids"},
+            }
+        )
+
+    Again, think about whether this feature is really the one you want -- often times, verbose,
+    static DAGs are far more readable than very concise, highly parameterized DAGs.
+    """
+
+    def __init__(
+        self,
+        *load_from: ModuleType | Callable,
+        inputs: dict[
+            str, dependencies.ParametrizedDependency | dependencies.LiteralDependency
+        ] = None,
+        config: dict[str, Any] = None,
+        external_inputs: list[str] = None,
+        **parameterization: SubdagParams,
+    ):
+        """Initializes a parameterized_subdag decorator.
+
+        :param load_from: Modules to load from
+        :param inputs: Inputs for each subdag generated by the decorated function
+        :param config: Config for each subdag generated by the decorated function
+        :param external_inputs: External inputs to all parameterized subdags. Note that
+            if you pass in any external inputs from local subdags, it overrides this (does not merge).
+        :param parameterization: Parameterizations for each subdag generated.
+            Note that this *overrides* any inputs/config passed to the decorator itself.
+
+            Furthermore, note the following:
+
+            1. The parameterizations passed to the constructor are \\*\\*kwargs, so you are not
+            allowed to name these `load_from`, `inputs`, or `config`. That's a good thing, as these
+            are not good names for variables anyway.
+
+            2. Any empty items (not included) will default to an empty dict (or an empty list in
+            the case of parameterization)
+        """
+        self.load_from = load_from
+        self.inputs = inputs if inputs is not None else {}
+        self.config = config if config is not None else {}
+        self.parameterization = parameterization
+        self.external_inputs = external_inputs if external_inputs is not None else []
+
+    def _gather_subdag_generators(self) -> list[subdag]:
+        subdag_generators = []
+        for key, parameterization in self.parameterization.items():
+            subdag_generators.append(
+                subdag(
+                    *self.load_from,
+                    inputs={**self.inputs, **parameterization.get("inputs", {})},
+                    config={**self.config, **parameterization.get("config", {})},
+                    external_inputs=parameterization.get("external_inputs", self.external_inputs),
+                    namespace=key,
+                    final_node_name=key,
+                )
+            )
+        return subdag_generators
+
+    def generate_nodes(self, fn: Callable, config: dict[str, Any]) -> list[node.Node]:
+        generated_nodes = []
+        for subdag_generator in self._gather_subdag_generators():
+            generated_nodes.extend(subdag_generator.generate_nodes(fn, config))
+        return generated_nodes
+
+    def validate(self, fn: Callable):
+        for subdag_generator in self._gather_subdag_generators():
+            subdag_generator.validate(fn)
+
+    def required_config(self) -> list[str] | None:
+        """See note for subdag.required_config -- this is the same pattern.
+
+        :return: Any required config items.
+        """
+        return None
+
+
+def prune_nodes(nodes: list[node.Node], select: list[str] | None = None) -> list[node.Node]:
+    """Prunes the nodes to only include those upstream from the select columns.
+    Conducts a depth-first search using the nodes `input_types` field.
+
+    If select is None, we just assume all nodes should be included.
+
+    :param nodes: Full set of nodes
+    :param select: Columns to select
+    :return:  Pruned set of nodes
+    """
+    if select is None:
+        return nodes
+
+    node_name_map = {node_.name: node_ for node_ in nodes}
+    seen_nodes = set(select)
+    stack = list({node_name_map[col] for col in select if col in node_name_map})
+    output = []
+    while len(stack) > 0:
+        node_ = stack.pop()
+        output.append(node_)
+        for dep in node_.input_types:
+            if dep not in seen_nodes and dep in node_name_map:
+                dep_node = node_name_map[dep]
+                stack.append(dep_node)
+                seen_nodes.add(dep)
+
+    if not set(select) <= set([node_.name for node_ in output]):
+        raise ValueError(
+            "At least one of the selected nodes is not not in the DAG. "
+            f"You selected: {select}, but we only found nodes: {nodes}."
+        )
+    return output
+
+
+def _default_inject_parameter(fn: Callable, target_dataframe: str = None) -> str:
+    if target_dataframe is not None:
+        inject_parameter = target_dataframe
+    else:
+        # If we don't have a specified dataframe we assume it's the first argument
+        function_parameters = list(inspect.signature(fn).parameters.values())
+        if function_parameters:
+            inject_parameter = function_parameters[0].name
+        else:
+            raise ValueError(
+                f"Function {fn.__qualname__} has no parameters, but was "
+                f"decorated with with_columns. with_columns requires the first "
+                f"parameter to be a dataframe or using the on_input argument."
+            )
+    return inject_parameter
+
+
+class with_columns_base(base.NodeInjector, abc.ABC):
+    """Factory for with_columns operation on a dataframe. This is used when you want to extract some
+    columns out of the dataframe, perform operations on them and then append to the original dataframe.
+
+    This is an internal class that is meant to be extended by each individual dataframe library implementing
+    the following abstract methods:
+
+    - get_initial_nodes
+    - get_subdag_nodes
+    - chain_subdag_nodes
+    - validate
+    """
+
+    # TODO: if we rename the column nodes into something smarter this can be avoided and
+    # can also modify columns in place
+    @staticmethod
+    def contains_duplicates(nodes_: list[node.Node]) -> bool:
+        """Ensures that we don't run into name clashing of columns and group operations.
+
+        In the case when we extract columns for the user, because ``columns_to_pass`` was used, we want
+        to safeguard against nameclashing with functions that are passed into ``with_columns`` - i.e.
+        there are no functions that have the same name as the columns. This effectively means that
+        using ``columns_to_pass`` will only append new columns to the dataframe and for changing
+        existing columns ``pass_dataframe_as`` or ``on_input`` needs to be used.
+        """
+        node_counter = defaultdict(int)
+        for node_ in nodes_:
+            node_counter[node_.name] += 1
+            if node_counter[node_.name] > 1:
+                return True
+        return False
+
+    @staticmethod
+    def validate_dataframe(
+        fn: Callable, inject_parameter: str, params: dict[str, type[type]], required_type: type
+    ) -> None:
+        input_types = typing.get_type_hints(fn)
+        if inject_parameter not in params:
+            raise InvalidDecoratorException(
+                f"Function: {fn.__name__} does not have the parameter {inject_parameter} as a dependency. "
+                f"@with_columns requires the parameter names to match the function parameters. "
+                f"If you wish do not wish to use the first argument, please use ``pass_dataframe_as`` or ``on_input`` option. "
+                f"It might not be compatible with some other decorators."
+            )
+
+        if isinstance(input_types[inject_parameter], required_type):
+            raise InvalidDecoratorException(
+                "The selected dataframe parameter is not the correct dataframe type. "
+                f"You selected a parameter of type {input_types[inject_parameter]}, but we expect to get {required_type}"
+            )
+
+    def __init__(
+        self,
+        *load_from: Callable | ModuleType,
+        columns_to_pass: list[str] = None,
+        pass_dataframe_as: str = None,
+        on_input: str = None,
+        select: list[str] = None,
+        namespace: str = None,
+        config_required: list[str] = None,
+        dataframe_type: type = None,
+    ):
+        """Instantiates a ``@with_columns`` decorator.
+
+        :param load_from: The functions or modules that will be used to generate the group of map operations.
+        :param columns_to_pass: The initial schema of the dataframe. This is used to determine which
+            upstream inputs should be taken from the dataframe, and which shouldn't. Note that, if this is
+            left empty (and external_inputs is as well), we will assume that all dependencies come
+            from the dataframe. This cannot be used in conjunction with pass_dataframe_as.
+        :param pass_dataframe_as: The name of the dataframe that we're modifying, as known to the subdag.
+            If you pass this in, you are responsible for extracting columns out. If not provided, you have
+            to pass columns_to_pass in, and we will extract the columns out for you.
+        :param on_input: the dataframe parameter that we are applying with_columns on. By default we
+            will assume the first parameter is the corresponding dataframe.
+        :param select: The end nodes that represent columns to be appended to the original dataframe
+            via with_columns. Existing columns will be overridden.
+        :param namespace: The namespace of the nodes, so they don't clash with the global namespace
+            and so this can be reused. If its left out, there will be no namespace (in which case you'll want
+            to be careful about repeating it/reusing the nodes in other parts of the DAG.)
+        :param config_required: the list of config keys that are required to resolve any functions. Pass in None\
+            if you want the functions/modules to have access to all possible config.
+        """
+
+        self.subdag_functions = subdag.collect_functions(load_from)
+        self.select = select
+
+        # This is here to restrict to using either pass_dataframe_as or on_input or columns_to_pass
+        # TODO: decouple columns_to_pass, pass_dataframe_as and on_input
+        # For spark, we always perform with_columns on first parameter and use pass_dataframe_as;
+        # for pandas/polars, we can select which dataframe with on_input, but columns_to_pass, will always only work on first parameter
+        # We can decouple it so that on_input selects the target dataframe parameter that will inject into the next node
+        # pass_dataframe_as selects the original dataframe we want to extract columns from
+        # columns_to_pass is optinal helper that can be toggled on/off so no need to raise this error.
+        if (
+            int(pass_dataframe_as is None) + int(columns_to_pass is None) + int(on_input is None)
+            == 1
+        ):
+            raise ValueError(
+                "You must specify only one of ``columns_to_pass``, ``pass_dataframe_as``, and ``on_input``. "
+                "This is because specifying ``pass_dataframe_as`` or ``on_input`` injects into "
+                "the set of columns, allowing you to perform your own extraction"
+                "from the dataframe. We then execute all columns in the subdag"
+                "in order, passing in that initial dataframe. If you want"
+                "to reference columns in your code, you'll have to specify "
+                "the set of initial columns, and allow the subdag decorator "
+                "to inject the dataframe through. The initial columns tell "
+                "us which parameters to take from that dataframe, so we can"
+                "feed the right data into the right columns."
+            )
+
+        self.initial_schema = columns_to_pass
+        self.dataframe_subdag_param = pass_dataframe_as
+        self.target_dataframe = on_input
+        self.namespace = namespace
+        self.config_required = config_required
+
+        if dataframe_type is None:
+            raise InvalidDecoratorException(
+                "Please provide the dataframe type for this specific library."
+            )
+
+        self.dataframe_type = dataframe_type
+
+    def required_config(self) -> list[str]:
+        return self.config_required
+
+    @abc.abstractmethod
+    def get_initial_nodes(
+        self, fn: Callable, params: dict[str, type[type]]
+    ) -> tuple[str, Collection[node.Node]]:
+        """Preparation stage where columns get extracted into nodes. In case `pass_dataframe_as` or `on_input` is
+        used, this should return an empty list (no column nodes) since the users will extract it
+        themselves.
+
+        :param fn: the function we are decorating. By using the inspect library you can get information.
+            about what arguments it has / find out the dataframe argument.
+        :param params: Dictionary of all the type names one wants to inject.
+        :return: name of the dataframe parameter and list of nodes representing the extracted columns (can be empty).
+        """
+        pass
+
+    @abc.abstractmethod
+    def get_subdag_nodes(self, fn: Callable, config: dict[str, Any]) -> Collection[node.Node]:
+        """Creates subdag from the passed in module / functions.
+
+        :param config: Configuration with which the DAG was constructed.
+        :return: the subdag as a list of nodes.
+        """
+        pass
+
+    @abc.abstractmethod
+    def chain_subdag_nodes(
+        self, fn: Callable, inject_parameter: str, generated_nodes: Collection[node.Node]
+    ) -> node.Node:
+        """Combines the origanl dataframe with selected columns. This should produce a
+        dataframe output that is injected into the decorated function with new columns
+        appended and existing columns overriden.
+
+        :param inject_parameter: the name of the original dataframe that.
+        :return: the new dataframe with the columns appended / overwritten.
+        """
+        pass
+
+    def inject_nodes(
+        self, params: dict[str, type[type]], config: dict[str, Any], fn: Callable
+    ) -> tuple[list[node.Node], dict[str, str]]:
+        namespace = fn.__name__ if self.namespace is None else self.namespace
+
+        inject_parameter, initial_nodes = self.get_initial_nodes(fn=fn, params=params)
+        subdag_nodes = self.get_subdag_nodes(fn=fn, config=config)
+        generated_nodes = initial_nodes + subdag_nodes
+        # TODO: for now we restrict that if user wants to change columns that already exist, he needs to
+        # pass the dataframe and extract them himself. If we add namespace to initial nodes and rewire the
+        # initial node names with the ongoing ones that have a column argument, we can also allow in place
+        # changes when using columns_to_pass
+        if with_columns_base.contains_duplicates(generated_nodes):
+            raise ValueError(
+                "You can only specify columns once. You used `columns_to_pass` and we "
+                "extract the columns for you. In this case they cannot be overwritten -- only new columns get "
+                "appended. If you want to modify in-place columns pass in a dataframe and "
+                "extract + modify the columns and afterwards select them."
+            )
+
+        pruned_nodes = prune_nodes(nodes=generated_nodes, select=self.select)
+        if len(pruned_nodes) == 0:
+            raise ValueError(
+                f"No nodes found upstream from select columns: {self.select} for function: "
+                f"{fn.__qualname__}"
+            )
+
+        # Node combining columns and dataframe might need info about prior nodes
+        output_nodes, current_param = self.chain_subdag_nodes(
+            fn=fn, inject_parameter=inject_parameter, generated_nodes=pruned_nodes
+        )
+        output_nodes = subdag.add_namespace(output_nodes, namespace)
+        return output_nodes, {inject_parameter: assign_namespace(current_param, namespace)}