Handling Tabular Data with DataFrame

In flyte-sdk, tabular data is managed through the DataFrame class, which provides a high-level abstraction for structured datasets. Unlike a standard pandas or pyarrow dataframe, the flytekit.io.DataFrame is a pointer to data stored in a remote location (like S3 or GCS) along with metadata about its format and schema.

Working with Tabular Files

The most common way to handle tabular data in flyte-sdk is by using the File type parameterized with DataFrame. This indicates to the system that the file contains structured data, typically in Parquet or CSV format.

Reading Data in a Task

When a task receives a File[DataFrame], you use standard data libraries to load the content. Because flyte-sdk handles the underlying storage abstraction, you can open the file and pass the handle directly to readers like pandas.read_csv.

import pandas as pd
from flyte.io import File, DataFrame
from flyte import env

@env.task
async def process_table(file: File[DataFrame]):
    # Open the remote file for reading in binary mode
    async with file.open("rb") as f:
        # Load into a pandas DataFrame for manipulation
        df = pd.read_csv(f)
        print(f"Loaded {len(df)} rows")

Writing Data from a Task

To return tabular data, you can create a new remote file reference and write your local dataframe to it. flyte-sdk ensures the data is streamed to the appropriate remote storage.

import pandas as pd
from flyte.io import File, DataFrame
from flyte import env

@env.task
async def generate_table() -> File[DataFrame]:
    # Create some local data
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    
    # Create a reference for a new remote file
    file = File.new_remote()
    
    # Write the dataframe directly to the remote storage
    async with file.open("wb") as f:
        df.to_csv(f, index=False)
        
    return file

The DataFrame Abstraction

The io._dataframe.dataframe.DataFrame class is a dataclass that wraps the location (uri) and the file_format of the data. It is designed to be serializable and compatible with the Flyte type engine.

Internal Mechanism

Internally, flyte-sdk uses the DataFrameTransformerEngine to convert between Python objects and Flyte literals. This engine manages a registry of encoders and decoders:

• DataFrameEncoder: Converts a Python dataframe (like pandas.DataFrame) into a Flyte StructuredDataset literal.
• DataFrameDecoder: Converts a Flyte StructuredDataset literal back into a Python dataframe.

When you return a pandas.DataFrame from a task, the DataFrameTransformerEngine looks up the appropriate encoder based on the type and the configured default format (e.g., Parquet).

Custom Handlers

You can extend flyte-sdk to support new dataframe libraries or storage formats by implementing and registering custom handlers.

from flyte.io import DataFrameEncoder, DataFrameTransformerEngine
from flyteidl.core import literals_pb2, types_pb2

class MyCustomEncoder(DataFrameEncoder):
    def __init__(self):
        super().__init__(python_type=MyCustomDF, supported_format="custom-fmt")

    async def encode(
        self,
        dataframe: DataFrame,
        structured_dataset_type: types_pb2.StructuredDatasetType,
    ) -> literals_pb2.StructuredDataset:
        # Implementation for writing MyCustomDF to storage
        ...

# Register the encoder with the engine
DataFrameTransformerEngine.register(MyCustomEncoder())

Important Considerations

• Binary Mode: When using file.open(), always use binary modes like "rb" or "wb". Remote storage drivers in flyte-sdk require binary access for consistency across different providers.
• Lazy Loading: The DataFrame.open() method is provided to load handlers lazily. This is useful when a specific library (like pandas) is only imported inside the task body to keep the container image lean.
• Type Hints: While File[DataFrame] is the standard way to handle files, you can also use DataFrame directly in type hints. In this case, flyte-sdk will automatically invoke the registered encoders/decoders to handle the data transfer.
• Column Subsetting: The DataFrameTransformerEngine supports column subsetting. If a task signature specifies a subset of columns via Annotated, the engine's to_python_value method ensures that decoders receive the requested schema, allowing them to optimize the download by only fetching required columns if the format (like Parquet) supports it.