Handling Tabular Data with `DataFrame`

The DataFrame type in flyte-sdk provides a high-level abstraction for structured tabular data. Unlike a standard in-memory dataframe, a flyte.io.DataFrame acts as a reference to data stored in remote locations (like S3, GCS, or BigQuery) in specific formats (like Parquet or CSV). This allows flyte-sdk to handle massive datasets efficiently by deferring data loading until it is explicitly requested within a task.

Accessing Data in Tasks

When a task receives a DataFrame as an input, the data is not automatically loaded into memory. To work with the data, you must specify which concrete dataframe library you want to use (e.g., pandas or PyArrow) and then load the content.

If you attempt to access the data without opening it, you will encounter a ValueError stating that no dataframe type has been set. Use the .open() and .all() pattern to retrieve the data:

from flytekit import task
from flyte.io import DataFrame
import pandas as pd

@task
def process_data(df: DataFrame) -> int:
    # Specify the desired local type and load all data
    pandas_df = df.open(pd.DataFrame).all()
    
    # Now you can use standard pandas operations
    return len(pandas_df)

Internally, df.open(pd.DataFrame) triggers lazy_import_dataframe_handler() in io/_dataframe/dataframe.py to ensure the appropriate handlers for pandas are registered. The .all() method then uses the DataFrameTransformerEngine to find a registered DataFrameDecoder that matches the storage protocol (e.g., s3) and the file format (e.g., parquet).

Producing and Returning Dataframes

You can return tabular data from a task in two ways: by returning a concrete dataframe object directly or by returning a flyte.io.DataFrame wrapper.

Returning Concrete Types

If you return a standard pd.DataFrame, flyte-sdk automatically wraps it and uses the DataFrameTransformerEngine to select an encoder. By default, pandas dataframes are often encoded as Parquet files in the workflow's configured raw data prefix.

@task
def create_data() -> pd.DataFrame:
    return pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})

Returning the DataFrame Wrapper

You can also explicitly return a flyte.io.DataFrame object. This is useful when you want to point to an existing file in remote storage without loading it into the current task's memory, or when you need to specify a specific format like CSV.

@task
def link_external_data() -> DataFrame:
    # Point to an existing Parquet file in S3
    return DataFrame(
        uri="s3://my-bucket/data/table.parquet",
        file_format="parquet"
    )

When this task completes, the DataFrameTransformerEngine.to_literal method handles the conversion. If the val attribute is empty but a uri is provided, flyte-sdk creates a StructuredDataset literal pointing to that location.

How the Engine Works

The DataFrameTransformerEngine (found in io/_dataframe/dataframe.py) is a meta-transformer that manages a registry of encoders and decoders. It decouples the Flyte type system from specific dataframe libraries.

Encoders and Decoders

• DataFrameEncoder: Responsible for taking a Python object (like pd.DataFrame) and writing it to storage. For example, the PandasToCSVEncodingHandler in io/_dataframe/basic_dfs.py uses df.to_csv() to write data to a URI provided by the Flyte context.
• DataFrameDecoder: Responsible for reading data from a URI and reconstructing the Python object.

The engine selects the correct handler using the _finder method, which looks for a match based on:

1. The Python type (e.g., pandas.DataFrame)
2. The storage protocol (e.g., s3, gs, file)
3. The file format (e.g., parquet, csv)

Streaming with Iterators

For very large datasets that do not fit in memory, the DataFrame class supports an asynchronous iterator interface. If the underlying decoder supports it, you can stream data in chunks:

@task
async def stream_process(df: DataFrame):
    # Use .iter() instead of .all()
    async for chunk in await df.open(pd.DataFrame).iter():
        print(f"Processing chunk of size {len(chunk)}")

This invokes DataFrameTransformerEngine.iter_as, which calls the decode method of the registered decoder. The decoder must return an AsyncIterator for this to work.

Custom Handlers

You can extend flyte-sdk to support new dataframe libraries or storage formats by implementing the DataFrameEncoder and DataFrameDecoder interfaces and registering them with the engine:

from flyte.io import DataFrameEncoder, DataFrameTransformerEngine

class MyCustomEncoder(DataFrameEncoder):
    def __init__(self):
        super().__init__(MyCustomType, protocol="s3", supported_format="custom_fmt")

    async def encode(self, dataframe, structured_dataset_type):
        # Implementation to write MyCustomType to S3
        ...

# Register the handler
DataFrameTransformerEngine.register(MyCustomEncoder())