Extending Support for Custom DataFrames

Dataframe handling in flyte-sdk decouples in-memory tabular representations from serialized physical storage formats and cloud storage protocols. Instead of registering every dataframe format directly with Flyte's core TypeEngine, flyte-sdk routes dataframe serialization and deserialization through a meta-transformer: DataFrameTransformerEngine.

This architecture enables developers to add support for custom in-memory dataframe libraries (such as Polars, Vaex, or internal matrix wrappers) or specialized file formats (such as Feather, Avro, or custom binary encodings) by implementing and registering custom DataFrameEncoder and DataFrameDecoder classes.

Architecture Overview

The tabular data pipeline in flyte-sdk is structured into three distinct layers:

1. User Interface (flyte.io.DataFrame): The high-level container class that wraps either an in-memory dataframe object (val), a storage URI (uri), or metadata (StructuredDatasetMetadata). Tasks can receive or return DataFrame instances or raw in-memory objects like pd.DataFrame and pa.Table.
2. Meta-Transformer Engine (DataFrameTransformerEngine): A TypeTransformer[DataFrame] subclass that manages registries of encoders and decoders indexed across three dimensions: (python_type, protocol, format). It inspects task type signatures, handles column projection and subsetting, and delegates serialization/deserialization to registered handlers.
3. Encoding and Decoding Handlers (DataFrameEncoder and DataFrameDecoder): Format- and protocol-specific worker classes that write in-memory dataframe structures to storage and read them back.

When a Flyte task returns a dataframe, DataFrameTransformerEngine.to_literal resolves the correct DataFrameEncoder based on the in-memory object's Python type, the destination URI storage protocol (e.g., s3, gs, file), and the target file format (such as parquet or csv). Conversely, when a task accepts a dataframe input, DataFrameTransformerEngine.to_python_value identifies the appropriate DataFrameDecoder to materialize the object.

                   +----------------------------------+
                   |       User Task Signature        |
                   | (e.g. pd.DataFrame, DataFrame)   |
                   +----------------------------------+
                                    |
                                    v
                   +----------------------------------+
                   |   DataFrameTransformerEngine     |
                   |   - Resolves (Type, Proto, Fmt)  |
                   |   - Manages Column Projections   |
                   +----------------------------------+
                               /      \
            encode()          /        \          decode()
                             v          v
  +-----------------------------+    +-----------------------------+
  |      DataFrameEncoder       |    |      DataFrameDecoder       |
  | (e.g. PandasToParquet... )  |    | (e.g. ParquetToPandas... )  |
  +-----------------------------+    +-----------------------------+
                 |                                  ^
                 v                                  |
  +----------------------------------------------------------------+
  |              Physical Storage (S3, GCS, Local)                 |
  |              Flyte IDL StructuredDataset Literal               |
  +----------------------------------------------------------------+

Implementing a Custom DataFrameEncoder

To convert a custom dataframe type or file format into a Flyte literal, subclass DataFrameEncoder defined in flyte.io._dataframe.dataframe. The encoder must implement the asynchronous encode method.

Constructor Arguments

When initializing a DataFrameEncoder, provide:

• python_type: The concrete Python class the encoder serializes (for example, pd.DataFrame, pa.Table, or a custom class).
• protocol: An optional storage scheme prefix (e.g., "s3", "gs", "file"). If set to None, the encoder is registered under "fsspec", allowing it to handle any storage filesystem supported by flyte-sdk's persistence layer. Do not pass an empty string "" (which raises ValueError).
• supported_format: A string identifier for the storage serialization format (e.g., "parquet", "csv", "feather"). Supplying an empty string "" designates the handler as a generic fallback encoder for that type.

Implementing encode

The encode method receives two parameters:

1. dataframe: A flyte.io.DataFrame wrapper object whose val property contains the in-memory dataframe instance and whose uri property optionally contains a user-specified destination path.
2. structured_dataset_type: A protobuf types_pb2.StructuredDatasetType instance populated with interface information from the task signature.

If dataframe.uri is not set, obtain a designated remote storage location using internal_ctx().raw_data.get_random_remote_path().

Example: Custom CSV Encoder

The built-in PandasToCSVEncodingHandler in flyte.io._dataframe.basic_dfs demonstrates how to write a dataframe to CSV format using flyte-sdk's storage options:

import os
import typing
from pathlib import Path
import pandas as pd
from flyteidl.core import literals_pb2, types_pb2
from flyte.io._dataframe.basic_dfs import CSV, get_pandas_storage_options
from flyte.io._dataframe.dataframe import DataFrame, DataFrameEncoder
from flyte import storage

class PandasToCSVEncodingHandler(DataFrameEncoder):
    def __init__(self):
        super().__init__(pd.DataFrame, None, CSV)

    async def encode(
        self,
        dataframe: DataFrame,
        structured_dataset_type: types_pb2.StructuredDatasetType,
    ) -> literals_pb2.StructuredDataset:
        if not dataframe.uri:
            from flyte._context import internal_ctx

            ctx = internal_ctx()
            uri = 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, ".csv")
        df = typing.cast(pd.DataFrame, dataframe.val)
        df.to_csv(
            path,
            index=False,
            storage_options=get_pandas_storage_options(uri=path),
        )
        structured_dataset_type.format = CSV
        return literals_pb2.StructuredDataset(
            uri=uri,
            metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type),
        )

Implementing a Custom DataFrameDecoder

To deserialize stored tabular data back into an in-memory dataframe or an asynchronous iterator, subclass DataFrameDecoder and implement the asynchronous decode method.

Constructor Arguments

DataFrameDecoder takes:

• python_type: The target Python dataframe type produced by this decoder.
• protocol: The storage scheme prefix ("s3", "gs", "file"), or None for all fsspec-supported protocols.
• supported_format: The format string handled by this decoder (e.g., "parquet", "csv").

Column Subsetting and Task Metadata

When executing a task with typed column schemas, DataFrameTransformerEngine populates current_task_metadata.structured_dataset_type.columns. Decoders must inspect this metadata and read only the requested subset of columns from storage to avoid unnecessary I/O overhead.

Example: Custom Parquet Decoder

The built-in ParquetToPandasDecodingHandler in flyte.io._dataframe.basic_dfs reads parquet files using column filtering and handles credential fallbacks:

import logging
import pandas as pd
from flyteidl.core import literals_pb2
from flyte.io._dataframe.basic_dfs import PARQUET, get_pandas_storage_options
from flyte.io._dataframe.dataframe import DataFrameDecoder

logger = logging.getLogger(__name__)

class ParquetToPandasDecodingHandler(DataFrameDecoder):
    def __init__(self):
        super().__init__(pd.DataFrame, None, PARQUET)

    async def decode(
        self,
        flyte_value: literals_pb2.StructuredDataset,
        current_task_metadata: literals_pb2.StructuredDatasetMetadata,
    ) -> pd.DataFrame:
        uri = flyte_value.uri
        columns = None
        kwargs = get_pandas_storage_options(uri=uri)

        # Extract column subset requested by the downstream task signature
        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]

        try:
            return pd.read_parquet(uri, columns=columns, storage_options=kwargs)
        except Exception as exc:
            if exc.__class__.__name__ == "NoCredentialsError":
                logger.debug("S3 source detected, attempting anonymous S3 access")
                kwargs = get_pandas_storage_options(uri=uri, anonymous=True)
                return pd.read_parquet(uri, columns=columns, storage_options=kwargs)
            else:
                raise

Handler Registration and Engine Dispatch

Once concrete encoders and decoders are defined, register them with DataFrameTransformerEngine.register.

Registration Options

DataFrameTransformerEngine.register(h, ...) accepts several configuration parameters:

• h: The DataFrameEncoder or DataFrameDecoder instance.
• default_for_type: If True, sets this handler's format and protocol as the default when a task returns a raw dataframe instance instead of a DataFrame wrapper. Note: Do not pass default_for_type=True if handler.protocol is None, as protocol=None targets all protocols. A ValueError is raised if this is attempted.
• default_format_for_type: Sets this handler's format as the default format for the Python type without locking the storage protocol.
• default_storage_for_type: Sets this handler's protocol as the default storage protocol for the Python type.
• override: If True, replaces any existing handler registered for the same (python_type, protocol, format) combination and overrides existing defaults.

Handling DuplicateHandlerError

Attempting to register a handler for a (python_type, protocol, supported_format) triple that already exists raises flyte.io._dataframe.dataframe.DuplicateHandlerError when override=False.

flyte-sdk uses functools.lru_cache and catches DuplicateHandlerError during registration routines (such as in flyte.io._dataframe.lazy_import_dataframe_handler) to ensure safe, idempotent initialization across multiple task invocations:

import functools
from flyte.io._dataframe.basic_dfs import (
    PandasToParquetEncodingHandler,
    ParquetToPandasDecodingHandler,
)
from flyte.io._dataframe.dataframe import (
    DataFrameTransformerEngine,
    DuplicateHandlerError,
)

@functools.lru_cache(maxsize=None)
def register_custom_handlers():
    try:
        DataFrameTransformerEngine.register(
            PandasToParquetEncodingHandler(),
            default_format_for_type=True,
        )
        DataFrameTransformerEngine.register(
            ParquetToPandasDecodingHandler(),
            default_format_for_type=True,
        )
    except DuplicateHandlerError:
        pass

Handler Resolution Order

When resolving an encoder or decoder during serialization (to_literal) or deserialization (open_as), DataFrameTransformerEngine._finder checks registrations in the following priority order:

1. Exact Match: handler_map[df_type][protocol][format]
2. Generic fsspec Handler: Matches handler_map[df_type]["fsspec"][format] or the generic format "" under "fsspec".
3. Protocol Generic Handler: Matches handler_map[df_type][protocol][""] or the type's default format under that protocol.
4. Single Registered Handler Fallback: If exactly one handler is registered under "fsspec" or the target protocol, that handler is selected.

If no matching handler satisfies the lookup criteria, DataFrameTransformerEngine raises a ValueError.

UI Rendering with Renderable

flyte-sdk allows registering visual HTML renderers for custom dataframe types so task inputs and outputs display formatted previews in the Flyte Deck UI.

Subclass flyte.types._renderer.Renderable and register the renderer with DataFrameTransformerEngine.register_renderer:

import pandas as pd
from flyte.io._dataframe.dataframe import DataFrameTransformerEngine
from flyte.types._renderer import Renderable

class CustomFrameRenderer(Renderable):
    def to_html(self, df: pd.DataFrame) -> str:
        # Generate an HTML representation (e.g. rendering top 10 rows)
        return df.head(10).to_html()

# Register the renderer for pd.DataFrame
DataFrameTransformerEngine.register_renderer(pd.DataFrame, CustomFrameRenderer())

When DataFrameTransformerEngine.to_html is called, it checks DataFrameTransformerEngine.Renderers for the dataframe type and returns the generated HTML string.

Best Practices and Implementation Constraints

1. Protocol Specification: Always use None rather than "" when initializing a handler intended to work across all storage backends. Setting protocol="" raises a ValueError("Use None instead of empty string for registering handler...").
2. Default Flags: When protocol=None, never pass default_for_type=True to DataFrameTransformerEngine.register(). Use default_format_for_type=True instead.
3. Column Projection: Always check current_task_metadata.structured_dataset_type.columns in decode(). Implementing column selection in custom decoders ensures workflows only transfer and load required columns into memory.
4. Direct Type Returns vs. Wrapper Returns: If a task returns an unwrapped dataframe (e.g., return my_df), DataFrameTransformerEngine automatically wraps it using the registered DEFAULT_FORMATS and DEFAULT_PROTOCOLS. If a task's return type annotation is DataFrame but the function body returns an unwrapped instance without matching types, to_literal raises a TypeTransformerFailedError.
5. Lazy Loading: If your custom dataframe library has heavy import times, invoke registration inside a lazy loader guarded by lazy_import_dataframe_handler() and wrap calls in try ... except DuplicateHandlerError to prevent multi-import registration conflicts.