Schema management for dataset-types

[elijahbenizzy]

Changes

How I tested this

Notes

Checklist

• PR has an informative and human-readable title (this will be pulled into the release notes)
• Changes are limited to a single goal (no scope creep)
• Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
• Any change in functionality is tested
• New functions are documented (with a description, list of inputs, and expected output)
• Placeholder code is flagged / future TODOs are captured in comments
• Project documentation has been updated if adding/changing functionality.

[sweep-ai]

Apply Sweep Rules to your PR?

• Apply: All new business logic should have corresponding unit tests.
• Apply: Refactor large functions to be more modular.
• Apply: Add docstrings to all functions and file headers.

[elijahbenizzy]

POC taking @roelbertens's implementation and making a little more general.

Design decisions to enable this:

1. This is going to take the form of a decorator
2. The data, for now, is not going to be accessible
3. We store the data under tags in the node as a serialized format. Under hamilton.internal.schema_output=...

Further questions:

1. How do we allow for API change?
   • Keep with this API
   • Have the parameter not be **kwargs -- be able to switch it later
   • Change this to dataset_schema or something like that, use schema later
   • Keep in experimental/call it schema_beta
2. Should schemas allow order?
3. Should we be able to build schemas for specific libraries (spark, pydantic, etc...)?
   • Probably yes -- we should have a schema class
   • This schema class will have the ability to specify a "flavor", which will be generic
   • It will allow for column names/types
   • We should be able to convert between them/check compatibility with the right plugin
4. How to decorate input types?
   • Type-annotation

Proposal for more general schema

class SchemaFlavor(enum.Enum):
    common = "common" # common set of basic ones
    spark = "spark"
    pandas = "pandas"
    pydantic = "pydantic"


@dataclasses.dataclass
class Schema:
    # fields are limited to those supported by "flavor"
    # if ordereddict then it has order, otherwise it doesn't
    fields: Union[Dict[str, Type], OrderedDict[str, Type]] 
    flavor: SchemaFlavor = "common"
    allow_extra: bool = True # whether or not extra are allowed

The user has three entry-points:

1. Pass in a schema (specified above) to a @schema decorator
2. Pass in a Dict[str, Type] to shortcut the schema decorator
3. Add a type-annotation:

from hamilton.htypes import dataframe
SchemaIn=Schema(...)
SchemaOut=Schema(...)

# replace with spark, polars, etc...
def foo(input: dataframe[pd.DataFrame, SchemaIn]) -> dataFrame[pd.DataFrame, SchemaOut]:
    ...

Then for validation we can:

1. Add a lifecycle method that does automated validation at runtime of inputs, outputs
2. Add a graph method that validates types between edges

TODO remaining:

• Add some basic unit tests for schema
• Add/fix some tests for viz
• Ensure that the node modifiers work together (E.G. we have some colors that we need to use with the different node modifiers)
• Allow for disabling at the DAG level
• Break type-registry on spark into separate commits