Extending DataFrame Support

Register a dataframe format

To add a dataframe library or file format to flyte-sdk, implement an asynchronous DataFrameEncoder and DataFrameDecoder, then register each instance with DataFrameTransformerEngine rather than with the core TypeEngine directly. The public extension surface is re-exported from flyte.io:

from flyte.io import DataFrame
from flyte.io import DataFrameDecoder
from flyte.io import DataFrameEncoder
from flyte.io import DataFrameTransformerEngine

DataFrameTransformerEngine is the higher-level transformer for all dataframe types. Its registration path also registers the concrete dataframe type with TypeEngine, so a handler participates in normal Flyte type inference without a separate TypeEngine.register_additional_type call.

Follow the encoder and decoder contracts

An encoder is selected when a Python dataframe value is converted to a Flyte literal. Its constructor identifies the concrete Python type, storage protocol, and format; its encode method receives the user-facing DataFrame wrapper and must return a Flyte StructuredDataset protobuf:

class DataFrameEncoder(ABC, Generic[T]):
    def __init__(
        self,
        python_type: Type[T],
        protocol: Optional[str] = None,
        supported_format: Optional[str] = None,
    ):
        self._python_type = python_type
        self._protocol = protocol.replace("://", "") if protocol else None
        self._supported_format = supported_format or ""

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

A decoder is used in the reverse direction. It receives the incoming StructuredDataset and the metadata calculated for the currently running task. It may return one dataframe or an async iterator:

class DataFrameDecoder(ABC, Generic[DF]):
    def __init__(
        self,
        python_type: Type[DF],
        protocol: Optional[str] = None,
        supported_format: Optional[str] = None,
        additional_protocols: Optional[List[str]] = None,
    ):
        self._python_type = python_type
        self._protocol = protocol.replace("://", "") if protocol else None
        self._supported_format = supported_format or ""

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

The additional_protocols constructor argument is accepted by DataFrameDecoder, but the constructor stores only the Python type, normalized protocol, and format. Protocols such as s3:// are stored as s3; passing None means that the handler is protocol-independent. An omitted format becomes the empty generic format (GENERIC_FORMAT). Use an explicit format for a format-specific handler.

The encoder input is always the wrapper, even when task code starts with a plain dataframe value. The wrapper exposes the underlying value as dataframe.val and an optional destination as dataframe.uri. The encoder must return literals_pb2.StructuredDataset, not a Python dataframe and not a complete literals_pb2.Literal. The engine creates the enclosing Literal after the handler returns.

Use the built-in Pandas CSV path as a model

The built-in CSV pair in io/_dataframe/basic_dfs.py shows the complete storage and protobuf pattern:

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)
        )

This handler uses None for the protocol, obtains a random remote destination when the wrapper has no URI, creates a local directory only for a non-remote URI, and writes the artifact at <uri>/.csv. A custom format should make its artifact layout equally explicit. Remote access goes through the storage configuration rather than through unconditional local filesystem calls.

The corresponding decoder reads the task's requested columns from current_task_metadata:

class CSVToPandasDecodingHandler(DataFrameDecoder):
    def __init__(self):
        super().__init__(pd.DataFrame, None, CSV)

    async def decode(
        self,
        proto_value: literals_pb2.StructuredDataset,
        current_task_metadata: literals_pb2.StructuredDatasetMetadata,
    ) -> "pd.DataFrame":
        uri = proto_value.uri
        columns = None
        kwargs = get_pandas_storage_options(uri=uri)
        path = os.path.join(uri, ".csv")
        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_csv(path, usecols=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_csv(path, usecols=columns, storage_options=kwargs)
            else:
                raise

The decoder receives metadata for the current task, not just the metadata embedded in the incoming literal. The built-in handler turns requested dataset columns into Pandas usecols. It also retries an S3 read anonymously only when the exception class name is NoCredentialsError; other exceptions are re-raised.

Register both sides of a format

The built-in CSV registration function is the minimal format-specific pair-registration pattern:

@functools.lru_cache(maxsize=None)
def register_csv_handlers():
    from .basic_dfs import CSVToPandasDecodingHandler, PandasToCSVEncodingHandler

    DataFrameTransformerEngine.register(PandasToCSVEncodingHandler(), default_format_for_type=True)
    DataFrameTransformerEngine.register(CSVToPandasDecodingHandler(), default_format_for_type=True)

For a custom dataframe class, instantiate your encoder and decoder with the same python_type, protocol policy, and format identifier, then call DataFrameTransformerEngine.register for each. The handler class itself is responsible for the library-specific read/write operation; registration supplies the lookup key used by the engine.

The Pandas Parquet registration demonstrates the same approach for another explicit format and adds an optional renderer:

@functools.lru_cache(maxsize=None)
def register_pandas_handlers():
    import pandas as pd

    from flyte.types._renderer import TopFrameRenderer

    from .basic_dfs import PandasToParquetEncodingHandler, ParquetToPandasDecodingHandler

    DataFrameTransformerEngine.register(PandasToParquetEncodingHandler(), default_format_for_type=True)
    DataFrameTransformerEngine.register(ParquetToPandasDecodingHandler(), default_format_for_type=True)
    DataFrameTransformerEngine.register_renderer(pd.DataFrame, TopFrameRenderer())

register_renderer is optional and affects HTML rendering; it does not select an encoder or decoder.

Choose registration flags deliberately

register accepts these flags:

• default_format_for_type=True records the handler's non-generic format as the default for its Python dataframe type. It does not select a storage protocol. The built-in protocol-independent Pandas and Arrow handlers use this option.
• default_storage_for_type=True records the handler's protocol as the default storage protocol.
• default_for_type=True records both format and protocol defaults. It is invalid when the handler's protocol is None, because a protocol-independent handler must derive the protocol from the URI or active raw-data context.
• override=True permits replacing an existing handler and replacing an existing default.

A protocol-independent handler is routed into the special fsspec bucket. Do not pass an empty string as the protocol: the registration code raises ValueError and explicitly requires None instead. A handler with protocol="s3://" is normalized to s3 before lookup.

Only one default format and one default protocol are retained for a Python dataframe type. Without override, a later conflicting default is not adopted. A duplicate handler key—(python_type, protocol, supported_format)—raises DuplicateHandlerError, a subclass of ValueError, from register_for_protocol. The lazy built-in registration path catches that exception; custom registration should avoid accidental duplicate global registrations or use override=True intentionally.

Understand handler resolution

The engine stores handlers in class-level ENCODERS and DECODERS maps, keyed by Python type, protocol, and format. get_encoder and get_decoder first attempt an exact match. If that fails, _finder considers the fsspec bucket, generic format, configured default format, and—in the generic-format case—a single available handler. If no candidate is found, it raises:

ValueError: Failed to find a handler for <type>, protocol [<protocol>], fmt ['<format>']

An empty supported format is therefore a fallback meaning “any format.” For a file-format implementation, use an explicit identifier such as the built-in CSV or PARQUET; otherwise the engine can select a generic handler and then force the output metadata format to that handler's empty supported_format.

The protocol comes from a registered default when one exists. Otherwise _protocol_from_type_or_prefix derives it from the dataframe URI or internal_ctx().raw_data.path. This makes protocol-independent handlers suitable for the storage protocols supported by Flyte's storage layer, provided their read/write implementation uses the corresponding storage APIs.

Trace the conversion lifecycle

For a wrapped dataframe value, the conversion path is:

1. DataFrameTransformerEngine.to_literal identifies the wrapped value's concrete type and resolves its protocol and format.
2. encode calls get_encoder and awaits the handler's encode method.
3. The engine supplies missing metadata, sets the metadata format to handler.supported_format, wraps the returned StructuredDataset in a Literal, modifies uploaded flyte:// URIs, and records the literal on the wrapper.

The engine-side finalization is visible here:

async def encode(
    self,
    sd: DataFrame,
    df_type: Type,
    protocol: str,
    format: str,
    structured_literal_type: types_pb2.StructuredDatasetType,
) -> literals_pb2.Literal:
    handler: DataFrameEncoder
    handler = self.get_encoder(df_type, protocol, format)

    sd_model = await handler.encode(sd, structured_literal_type)
    if sd_model.metadata is None:
        sd_model.metadata = literals_pb2.StructuredDatasetMetadata(structured_dataset_type=structured_literal_type)
    if sd_model.metadata and sd_model.metadata.structured_dataset_type is None:
        sd_model.metadata.structured_dataset_type = structured_literal_type
    sd_model.metadata.structured_dataset_type.format = handler.supported_format
    lit = literals_pb2.Literal(scalar=literals_pb2.Scalar(structured_dataset=sd_model))

    modify_literal_uris(lit)
    sd._literal_sd = sd_model
    sd._already_uploaded = True
    return lit

On the read path, DataFrame.open(dataframe_type) lazy-loads dataframe handlers and selects the local Python type. DataFrame.all() eventually calls open_as; open_as obtains the URI protocol, resolves a decoder from the incoming format, and awaits decode. A structured-dataset annotation remains a lazy DataFrame wrapper until opened, while a concrete dataframe type is opened immediately by to_python_value.

DataFrame.iter() uses iter_as, which is stricter than open_as: it directly indexes the decoder map rather than using the general fallback finder, and the decoder result must be an AsyncGeneratorType. A normal coroutine returning one dataframe raises ValueError in this path. Provide an exact decoder registration and an async-generator result if iteration is part of the format's support.

Handle optional dependencies and operational failures

The built-in registration functions are cached with functools.lru_cache and import Pandas or PyArrow inside the registration function. The core type engine calls flyte.io.lazy_import_dataframe_handler as part of lazy transformer loading, so optional dataframe libraries are not eagerly imported merely to use Flyte's type system. Follow the same delayed-import pattern when your handler depends on an optional library.

When a wrapper is created around an in-memory dataframe, pass the underlying value as DataFrame(val=...). When a wrapper represents an existing dataset, provide DataFrame(uri=..., file_format=...); to_literal requires a URI when val is absent and uploads a local URI through storage before creating the literal. The engine rejects a plain dataframe value when the declared type is the wrapper DataFrame; wrap the value instead.

Decoders should use current_task_metadata.structured_dataset_type.columns for projection, as the built-in Pandas and Arrow decoders do. Encoders should return metadata compatible with the supplied StructuredDatasetType; the engine fills absent type metadata but always overwrites its format with the selected handler's supported_format.

There are no discovered repository tests for custom handlers. Downstream implementations should add round-trip coverage for their encoder and decoder, duplicate-registration and override behavior, protocol and format selection, local and remote URIs, requested-column projection, and missing-credentials behavior.