Handling Tabular Data with DataFrame

The DataFrame wrapper and its lifecycle

When a task receives tabular data, the value is not automatically a Pandas object. flyte.io.DataFrame is a Flyte wrapper around a structured-dataset value: it may contain an in-memory dataframe (_val), a URI, a format, metadata, and the internal literals_pb2.StructuredDataset literal. It is also deliberately distinct from the protobuf literals.StructuredDataset; the latter is the serialized model used to exchange the dataset between tasks.

Import the public type from flyte.io:

from flyte.io import DataFrame

For a task input, select the dataframe library and then materialize the data. The wrapper's documented Pandas pattern is open(pd.DataFrame).all(); because all is asynchronous, an async consumer awaits it:

import pandas as pd

from flyte.io import DataFrame


async def consume(df: DataFrame) -> None:
    pandas_df = await df.open(pd.DataFrame).all()
    print(pandas_df)

open stores the requested Python dataframe type and returns the same wrapper, which makes chaining possible. It also calls lazy_import_dataframe_handler(). This matters when, as the DataFrame.open docstring demonstrates, Pandas is imported inside the task rather than when Flyte deserializes the task input. Calling all() without first calling open() raises:

ValueError: No dataframe type set. Use open() to set the local dataframe type you want to use.

Once a type has been selected, all() calls the dataframe transformer to decode the structured-dataset literal. If the wrapper has a URI but no in-memory value, all() first invokes _set_literal() so that the URI and format are represented as a Flyte literal. metadata and literal expose the wrapper's current metadata and structured-dataset literal when those values are available.

Opening a URI-backed dataset

Construct a wrapper with a URI when the dataset already exists, and retain its storage format explicitly:

import pandas as pd

from flyte.io import DataFrame


async def read_existing_dataset() -> pd.DataFrame:
    df = DataFrame(
        uri="s3://my-s3-bucket/s3_flyte_dir/df.parquet",
        file_format="parquet",
    )
    return await df.open(pd.DataFrame).all()

This is the URI-backed Parquet pattern included in the DataFrame source documentation. The URI identifies the structured dataset location; it is not necessarily the path of one physical file. The built-in Pandas Parquet decoder selects a decoder using the URI protocol and the literal's format, then calls pd.read_parquet.

For a streaming-capable custom decoder, use iter() instead:

async def consume_in_chunks(df: DataFrame) -> None:
    async for chunk in await df.open(pd.DataFrame).iter():
        print(chunk)

iter() delegates to DataFrameTransformerEngine.iter_as. Unlike open_as, iter_as requires the decoder result to be an actual async generator. The built-in Pandas CSV and Parquet decoders return one dataframe directly, so they support all() and do not provide dataframe iteration through iter().

Producing and exchanging data

A task can return a wrapper containing a dataframe value:

import pandas as pd

from flyte.io import DataFrame


async def produce() -> DataFrame:
    return DataFrame(val=pd.DataFrame({"name": ["Ada"], "score": [1]}))

DataFrameTransformerEngine.to_literal() handles three wrapper cases:

1. Existing literal: a wrapper passed through from an earlier task can reuse its _literal_sd. If the wrapper has both an existing literal and an in-memory value, conversion raises ValueError.
2. URI-only wrapper: when val is None, the URI must be truthy. A local URI is uploaded with storage.put() before the structured-dataset literal is built; a remote URI is retained. The wrapper's format is used when the expected task format is still the generic format.
3. In-memory value: the engine takes type(dataframe.val), determines a storage protocol, finds an encoder, and calls encode().

For a plain dataframe value rather than a DataFrame wrapper, the engine creates a wrapper internally and sends it through the same encoder path. Conversely, when a task expects DataFrame, to_python_value() preserves the structured-dataset literal in a lazy wrapper. When the expected type is a concrete dataframe type such as pd.DataFrame, it opens and decodes the value immediately instead.

After an encoder returns a protobuf StructuredDataset, the engine fills missing type metadata, sets the metadata format to the selected handler's supported_format, adjusts uploaded Flyte URIs, and stores the literal back on the wrapper. This stored literal is what makes a wrapper suitable for later task-output conversion or passthrough.

DataFrame also supplies serialization hooks for dataclass and Pydantic-related flows. _serialize() and serialize_dataframe() represent the value as a dictionary containing uri and file_format; _deserialize() and deserialize_dataframe() reconstruct it through DataFrameTransformerEngine. A dictionary without a URI is rejected with ValueError.

Formats and physical layout

The built-in handlers are in io/_dataframe/basic_dfs.py. Pandas CSV and Parquet handlers use None as their protocol, making them fsspec-capable handlers, and advertise csv and parquet formats respectively.

CSV is stored beneath the dataset URI in a .csv child:

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

    async def encode(self, dataframe, structured_dataset_type):
        if not dataframe.uri:
            from flyte._context import internal_ctx

            uri = internal_ctx().raw_data.get_random_remote_path()
        else:
            uri = dataframe.uri

        path = os.path.join(uri, ".csv")
        dataframe.val.to_csv(path, index=False)
        structured_dataset_type.format = CSV
        return literals_pb2.StructuredDataset(
            uri=uri,
            metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type),
        )

The complete project handler additionally creates local directories and supplies Pandas storage options. Its decoder reads the matching <uri>/.csv path with pd.read_csv.

Pandas and PyArrow Parquet encoders write a partition-like 00000 child below the dataset URI. The Pandas handler calls pd.to_parquet with coerce_timestamps="us" and allow_truncated_timestamps=False; its decoder calls pd.read_parquet on the dataset URI. The Arrow handlers use pyarrow.parquet and the same 00000 layout. Thus, do not treat a DataFrame URI as necessarily pointing directly to file.csv or one Parquet file.

The Pandas storage helper returns configured S3 fsspec options for S3 URLs, an empty dictionary for other fsspec URLs, and None for local paths because Pandas does not accept storage_options for non-fsspec paths. The CSV and Parquet decoders retry with anonymous S3 options only when the caught exception's class name is exactly NoCredentialsError; other exceptions are re-raised.

Columns and schema metadata

DataFrame.columns() returns an empty mapping by default, and column_names() derives names from that mapping:

from flyte.io import DataFrame


class Scores(DataFrame):
    @classmethod
    def columns(cls):
        return {"name": str, "score": float}


assert Scores.column_names() == ["name", "score"]

The transformer converts column mappings into StructuredDatasetType.DatasetColumn values. It recursively flattens nested dictionaries and dataclass descriptions into dot-separated names. Column types are converted to Flyte literal types; unsupported types cause _get_dataset_column_literal_type() to raise AssertionError.

When a task input supplies column metadata, to_python_value() builds updated metadata for the currently running task. If the input describes columns, those requested columns replace the incoming column list. If it does not, incoming literal columns are retained. The built-in Pandas decoders then pass the resulting names to usecols for CSV or columns for Parquet:

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]
return pd.read_parquet(uri, columns=columns, storage_options=kwargs)

The same projection logic is used by the CSV decoder with pd.read_csv(path, usecols=columns, ...). The annotation-processing helper accepts only one format, one column map, and one Arrow schema in an annotated type; supplying multiple instances raises ValueError.

Adding a dataframe library or format

Use DataFrameEncoder for Python-to-literal conversion and DataFrameDecoder for literal-to-Python conversion. Both constructors record a Python dataframe type, an optional protocol with :// removed, and a format. An omitted format becomes the generic empty format, which represents a handler that can handle any format.

An encoder receives the Flyte DataFrame wrapper and a types_pb2.StructuredDatasetType, and returns a protobuf literals_pb2.StructuredDataset:

from flyte.io import DataFrameEncoder, DataFrameTransformerEngine


class MyEncoder(DataFrameEncoder):
    def __init__(self, dataframe_type):
        super().__init__(dataframe_type, None, "my-format")

    async def encode(self, dataframe, structured_dataset_type):
        # Write dataframe.val and return a literals_pb2.StructuredDataset.
        raise NotImplementedError


DataFrameTransformerEngine.register(MyEncoder(MyDataFrameType))

The abstract method contract requires a concrete implementation; the placeholder above is therefore an extension shape, not a complete handler. A decoder implements decode(flyte_value, current_task_metadata) and may return either one dataframe or an async iterator. The task metadata argument is where a decoder receives the effective column projection.

Register the corresponding concrete encoder and decoder with DataFrameTransformerEngine.register(). Its options are significant:

• default_for_type=True sets both the default format and storage protocol for a concrete handler. It cannot be used with a handler whose protocol is None, because that protocol means all fsspec-capable protocols.
• default_format_for_type=True sets only the default format.
• default_storage_for_type=True sets only the default protocol.
• override=True permits replacing an existing registration or default.

Protocols are normalized by removing ://; registering an empty-string protocol is rejected, and None should be used for the all-protocol case. Duplicate (Python type, protocol, format) registrations raise DuplicateHandlerError unless overridden. Registration also calls TypeEngine.register_additional_type(..., override=True), making the concrete dataframe type visible to Flyte's type engine.

Handler lookup first attempts an exact dataframe-type/protocol/format match. DataFrameTransformerEngine._finder() then considers fsspec handlers, protocol-specific generic or default formats, and a single available handler when the requested format is generic. If no candidate matches, it raises ValueError.

Optional dependencies and lazy registration

Importing the dataframe core does not eagerly import Pandas or PyArrow. lazy_import_dataframe_handler() checks which optional modules are already imported, then registers only the applicable handlers. If Pandas is imported, it registers CSV and Parquet Pandas handlers and a Pandas renderer. If PyArrow is imported, it registers Arrow Parquet handlers and an Arrow renderer. BigQuery and Snowflake handlers are attempted only when their corresponding modules are imported, and missing connector packages are logged and skipped.

The built-in registration functions are cached with functools.lru_cache, and duplicate registration is caught during lazy loading. DataFrame.open() calls this lazy-registration function explicitly, while Flyte's type-engine lazy transformer loading also invokes it. Install the dataframe library you intend to use, import it in the execution environment, and use that concrete type in open(); Pandas Parquet operations additionally require a compatible Parquet engine.

DataFrame versus File[DataFrame]

File[DataFrame] is a separate file-oriented API in io/_file.py. It gives you a file handle and leaves parsing to your task code:

async with file.open() as f:
    df = pd.read_csv(f)

The DataFrame wrapper instead exchanges a structured-dataset literal and selects a registered encoder or decoder based on dataframe type, URI protocol, and format. Choose File[DataFrame] when you want direct file I/O and explicit Pandas parsing; choose DataFrame when you want Flyte's structured-dataset conversion, format metadata, handler dispatch, and column metadata propagation.