#666 Implement user-facing `HashFunction` Pull request opened by cosmicBboy This pull request introduces advanced examples and core support for content-based caching of DataFrames in Flyte using custom hash functions. It adds three new example scripts demonstrating how to enable cache key computation for pandas, Polars, and Flyte DataFrames, and updates the core

flyte.io.DataFrame

implementation to support hash-based caching via a flexible API. Key highlights: • New examples show how to use custom hash functions for caching with pandas, Polars, and Flyte DataFrames, including both annotation-based and explicit API usage. • Core changes to

flyte.io.DataFrame

and related modules add support for passing custom hash functions and storing hash values, enabling cache-aware workflows. --- ### New Examples for Content-Based Caching • Added `examples/advanced/hash_pandas_dataframe.py`: Demonstrates using

HashFunction

with

typing.Annotated

for pandas DataFrames, showing cache hits by comparing task outputs. • Added `examples/advanced/hash_polars_dataframe.py`: Demonstrates similar content-based caching for Polars DataFrames using custom hash functions and type annotations. • Added `examples/advanced/hash_flyte_dataframe.py`: Shows how to use a custom hash function with

flyte.io.DataFrame

via the new

hash_method

parameter in

from_local

, enabling cache key computation for Flyte DataFrames. ### Core API and Implementation Enhancements • Updated `flyte.io.DataFrame`: • Added

hash

field and internal

_hash_method

attribute to store and propagate hash values for cache key computation. [1] [2] • Modified

from_local

and

from_local_sync

methods to accept a

hash_method

parameter, allowing users to specify a

HashMethod

or precomputed hash string for content-based caching. [1] [2] [3] [4] • Added a module-level

_dataframe_hash_cache

to maintain hashes for raw DataFrames, supporting cache-aware behavior even when type annotations are not present. • Internal imports and typing improvements: • Imported

HashMethod

in

dataframe.py

and adjusted type signatures accordingly. • Added

HashFunction

to the public API in

flyte/io/__init__.py

. • Minor typing and import improvements for maintainability. --- ### Usage #### With

flyte.io.DataFrame

@env.task async def produce_dataframe() -> DataFrame: """ Producer task that creates a DataFrame with a custom hash method. The hash is computed from the DataFrame content using hash_pandas_dataframe, which allows downstream tasks to use content-based caching. """ df = pd.DataFrame(SAMPLE_DATA) # Create a HashFunction from our custom hash function hash_method = HashFunction.from_fn(hash_pandas_dataframe) # Use from_local to create a DataFrame with the custom hash return await DataFrame.from_local(df, hash_method=hash_method) @env.task(cache=Cache(behavior="override", version_override="v1")) async def consume_dataframe(input_df: DataFrame) -> str: """ Consumer task that processes the DataFrame. This task is cached based on the hash of the input DataFrame. If called with a DataFrame that has the same hash, the cached result will be returned. """ ... #### With third-party dataframe type e.g.

pd.DataFrame

from flyte.io import HashFunction HashedPandasDataFrame = Annotated[pd.DataFrame, HashFunction.from_fn(hash_pandas_dataframe)] @env.task async def produce_dataframe() -> HashedPandasDataFrame: """ Producer task that creates a pandas DataFrame. The return type annotation includes HashFunction, which tells Flyte to compute a content-based hash from the DataFrame using our custom hash function. """ df = pd.DataFrame(SAMPLE_DATA) df["action_name"] = flyte.ctx().action.run_name return df --- These changes make it much easier to implement cache-aware, content-based workflows with DataFrames in Flyte, supporting both explicit and annotation-based hash specification for pandas, Polars, and Flyte DataFrames. flyteorg/flyte-sdk