Advanced: Custom DataFrame Handlers

When you need to pass custom dataframe types (such as Polars, DuckDB, or specialized in-house tables) or non-standard file formats between tasks, Flyte's built-in type system needs to know how to serialize and deserialize them. In flyte-sdk, you implement this by subclassing DataFrameEncoder and DataFrameDecoder and registering them with DataFrameTransformerEngine.

Registering Custom Handlers

The following complete example creates a custom encoder and decoder pair for Polars DataFrames using Parquet format and registers them with DataFrameTransformerEngine:

import os
import typing
from pathlib import Path

import polars as pl
from flyteidl.core import literals_pb2, types_pb2

from flyte import storage
from flyte._context import internal_ctx
from flyte._logging import logger
from flyte.io import DataFrame, DataFrameDecoder, DataFrameEncoder, DataFrameTransformerEngine
from flyte.io._dataframe.dataframe import PARQUET


class PolarsToParquetEncodingHandler(DataFrameEncoder):
    def __init__(self):
        # python_type: pl.DataFrame
        # protocol: None (supports all fsspec-compatible storage protocols)
        # supported_format: PARQUET
        super().__init__(pl.DataFrame, None, PARQUET)

    async def encode(
        self,
        dataframe: DataFrame,
        structured_dataset_type: types_pb2.StructuredDatasetType,
    ) -> literals_pb2.StructuredDataset:
        if not dataframe.uri:
            ctx = internal_ctx()
            uri = str(ctx.raw_data.get_random_remote_path())
        else:
            uri = typing.cast(str, dataframe.uri)

        if not storage.is_remote(uri):
            Path(uri).mkdir(parents=True, exist_ok=True)
        path = os.path.join(uri, f"{0:05}")

        df = typing.cast(pl.DataFrame, dataframe.val)
        storage_options = storage.get_configured_fsspec_kwargs(
            storage.split_protocol(path)[0] or "file"
        )
        df.write_parquet(path, storage_options=storage_options)

        structured_dataset_type.format = PARQUET
        return literals_pb2.StructuredDataset(
            uri=uri,
            metadata=literals_pb2.StructuredDatasetMetadata(
                structured_dataset_type=structured_dataset_type
            ),
        )


class ParquetToPolarsDecodingHandler(DataFrameDecoder):
    def __init__(self):
        super().__init__(pl.DataFrame, None, PARQUET)

    async def decode(
        self,
        flyte_value: literals_pb2.StructuredDataset,
        current_task_metadata: literals_pb2.StructuredDatasetMetadata,
    ) -> pl.DataFrame:
        uri = flyte_value.uri
        columns = None
        if (
            current_task_metadata.structured_dataset_type
            and current_task_metadata.structured_dataset_type.columns
        ):
            columns = [c.name for c in current_task_metadata.structured_dataset_type.columns]

        storage_options = storage.get_configured_fsspec_kwargs(
            storage.split_protocol(uri)[0] or "file"
        )
        try:
            return pl.read_parquet(uri, columns=columns, storage_options=storage_options)
        except Exception as exc:
            if exc.__class__.__name__ == "NoCredentialsError":
                logger.debug("S3 source detected, attempting anonymous S3 access")
                storage_options = storage.get_configured_fsspec_kwargs("s3", anonymous=True)
                return pl.read_parquet(uri, columns=columns, storage_options=storage_options)
            raise


# Register the handlers with DataFrameTransformerEngine
DataFrameTransformerEngine.register(
    PolarsToParquetEncodingHandler(), default_format_for_type=True
)
DataFrameTransformerEngine.register(
    ParquetToPolarsDecodingHandler(), default_format_for_type=True
)

Implementing a DataFrameEncoder

DataFrameEncoder is an abstract generic base class (Generic[T]) responsible for converting in-memory DataFrame instances of type T into Flyte IDL literals_pb2.StructuredDataset literals.

Constructor Parameters

When subclassing DataFrameEncoder, call super().__init__(python_type, protocol, supported_format):

• python_type: The concrete Python type handled by this encoder (e.g., pd.DataFrame, pa.Table, pl.DataFrame).
• protocol: A storage prefix string (such as "s3", "gs", or "bq"). Trailing :// separators are automatically stripped. Pass None if the handler supports all storage backends accessible through Flyte's data persistence layer.
• supported_format: A format string identifier (such as PARQUET, CSV, or "" for generic format handlers).

Async Encoding Contract

Subclasses must implement the coroutine method:

async def encode(
    self,
    dataframe: DataFrame,
    structured_dataset_type: types_pb2.StructuredDatasetType,
) -> literals_pb2.StructuredDataset:
    ...

Key requirements inside encode():

1. Destination URI: Check dataframe.uri. If dataframe.uri is not set by the caller, generate a raw data output path by invoking internal_ctx().raw_data.get_random_remote_path().
2. Local Directory Creation: If the target location is local (storage.is_remote(uri) returns False), create directory paths on disk using Path(uri).mkdir(parents=True, exist_ok=True).
3. Data Serialization: Access the underlying Python object via dataframe.val and write it to disk or remote storage using Flyte's filesystem options (e.g. storage.get_underlying_filesystem(path=path) or get_pandas_storage_options(uri=path)).
4. Metadata Return: Assign the serialized format onto structured_dataset_type.format and return a literals_pb2.StructuredDataset containing uri=uri and metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type=structured_dataset_type).

Implementing a DataFrameDecoder

DataFrameDecoder is an abstract generic base class (Generic[DF]) responsible for decoding literals_pb2.StructuredDataset literals into Python dataframe instances or asynchronous generators.

Async Decoding Contract

Subclasses must implement the coroutine method:

async def decode(
    self,
    flyte_value: literals_pb2.StructuredDataset,
    current_task_metadata: literals_pb2.StructuredDatasetMetadata,
) -> Union[DF, typing.AsyncIterator[DF]]:
    ...

Key requirements inside decode():

1. Column Subsetting: Inspect current_task_metadata.structured_dataset_type.columns. If columns are specified by downstream type annotations, extract their names ([c.name for c in current_task_metadata.structured_dataset_type.columns]) and pass them to the reader to only load required columns.
2. Storage and Credential Handling: Obtain storage configurations for the source URI. When accessing public S3 buckets without explicit AWS credentials, catch NoCredentialsError and retry with anonymous=True options.

Built-In Handlers Reference

flyte-sdk ships with built-in handlers in flyte.io._dataframe.basic_dfs:

Pandas Handlers

• PandasToParquetEncodingHandler / ParquetToPandasDecodingHandler: Serializes pd.DataFrame to Parquet files using df.to_parquet() with microsecond timestamp coercion and deserializes with column projection via pd.read_parquet().
• PandasToCSVEncodingHandler / CSVToPandasDecodingHandler: Serializes pd.DataFrame to .csv format via df.to_csv(index=False) and deserializes using pd.read_csv().

PyArrow Handlers

• ArrowToParquetEncodingHandler: Serializes pa.Table to Parquet using pyarrow.parquet.write_table via storage.get_underlying_filesystem(path=path).
• ParquetToArrowDecodingHandler: Deserializes Parquet URIs to pa.Table using pyarrow.parquet.read_table(path, columns=columns).

Registration and UI Rendering

Registration Parameters

Register encoders and decoders using DataFrameTransformerEngine.register():

DataFrameTransformerEngine.register(
    h=handler,
    default_for_type=False,
    override=False,
    default_format_for_type=True,
    default_storage_for_type=False,
)

• h: The DataFrameEncoder or DataFrameDecoder instance.
• default_format_for_type: Marks this handler's format as the default for the Python type when no format is specified in task signatures.
• default_for_type: Sets both format and storage protocol as the default for raw type returns (cannot be used if protocol=None).
• override: If True, overrides existing registrations for the same (python_type, protocol, format) combination instead of raising DuplicateHandlerError.

Registering HTML Renderers for UI

To render custom DataFrame summaries in Flyte UI decks, register a Renderable implementation using register_renderer():

import pandas as pd
import pyarrow as pa

from flyte.io import DataFrameTransformerEngine
from flyte.types._renderer import ArrowRenderer, TopFrameRenderer

# Register UI renderers for dataframe types
DataFrameTransformerEngine.register_renderer(pd.DataFrame, TopFrameRenderer())
DataFrameTransformerEngine.register_renderer(pa.Table, ArrowRenderer())

Troubleshooting

ValueError: Registering SD handler ... with all protocols should never have default specified

• Cause: Passing default_for_type=True to DataFrameTransformerEngine.register() when the handler was initialized with protocol=None.
• Solution: Set default_format_for_type=True instead of default_for_type=True. protocol=None delegates storage protocol selection to the execution context's raw data prefix.

ValueError: Use None instead of empty string for registering handler

• Cause: Initializing a handler with protocol="".
• Solution: Pass protocol=None to indicate support for all storage persistence protocols.

DuplicateHandlerError

• Cause: Registering multiple handlers for the same (python_type, protocol, format) tuple.
• Solution: Pass override=True to DataFrameTransformerEngine.register() if replacing an existing handler registration.

Coroutine Await Errors

• Cause: Defining encode() or decode() as regular synchronous methods (def encode(...)).
• Solution: DataFrameTransformerEngine invokes serialization asynchronously. Always declare both methods as coroutines using async def.