Skip to main content

Interoperable Tabular Data with DataFrames

The DataFrame reference and its materialization boundary​

A task may need to pass tabular data between pandas, PyArrow, and other dataframe libraries without choosing one library at every task boundary. Flyte exposes flyte.io.DataFrame for that purpose: it is a Pydantic/SerializableType wrapper whose public state is a uri, a format, and an optional hash. Its private state can hold the in-memory dataframe, StructuredDataset metadata and literal, the requested dataframe type, and an optional lazy uploader.

Use wrap_df() when you already have a dataframe in memory and want a Flyte reference around it:

import pandas as pd
from flyte.io._dataframe.dataframe import DataFrame

pd_df = pd.DataFrame({"employee_id": [1001, 1002], "salary": [55000, 75000]})
fdf = DataFrame.wrap_df(pd_df)

wrap_df(val=None, uri=None) stores val in the wrapper; it does not invoke an encoder. from_df() has the same behavior but is deprecated, so new code should use wrap_df(). The public import is flyte.io.DataFrame, which is re-exported from flyte.io; the _dataframe modules are implementation and registration modules.

For a reference to an existing dataset, provide its URI and, when useful, its format:

from flyte.io._dataframe.dataframe import DataFrame

fdf = DataFrame.from_existing_remote(
    "s3://bucket/data.parquet",
    format="parquet",
)

from_existing_remote() marks the wrapper as already uploaded. It does not read the dataset. A wrapper with neither an in-memory value nor a URI cannot be serialized: the transformer requires a URI when it has no dataframe value to encode.

Local data and optional hashes​

Use from_local() or from_local_sync() when you are preparing local data with upload semantics, for example in a script or notebook:

import pandas as pd
from flyte.io._dataframe.dataframe import DataFrame

local_df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
local_reference = DataFrame.from_local_sync(local_df)

Despite the method name, from_local() is not unconditionally an immediate upload. When no raw-data context is available, DataFrame.from_local() and from_local_sync() create a wrapper and install a lazy uploader. In local run mode that uploader returns the original dataframe without uploading; in remote mode it uploads/converts the dataframe through Flyte's local data upload machinery. When a raw-data context is available, the wrapper is associated with the configured raw-data destination.

You can attach a cache hash either as a string or with HashFunction. The advanced example uses the dataframe content as the hash input:

import pandas as pd
from flyte.io._dataframe.dataframe import DataFrame
from flyte.io import HashFunction


def hash_pandas_dataframe(df: pd.DataFrame) -> str:
    return str(pd.util.hash_pandas_object(df).sum())


async def produce_dataframe() -> DataFrame:
    df = pd.DataFrame({"value": [10, 20]})
    hash_method = HashFunction.from_fn(hash_pandas_dataframe)
    return await DataFrame.from_local(df, hash_method=hash_method)

DataFrame.from_local() calls hash_method.update(df) and stores hash_method.result() in hash. During literal conversion, DataFrameTransformerEngine carries that value into the Flyte literal. The repository's consume_dataframe example then opens the reference as pandas and uses input_df.hash while processing it.

Select a dataframe type, then materialize it​

A DataFrame reference is intentionally not an in-memory pandas or Arrow object. Select the target library with open() and cross the materialization boundary with all():

import pandas as pd
from flyte.io._dataframe.dataframe import DataFrame


async def read_dataframe(df: DataFrame) -> pd.DataFrame:
    return await df.open(pd.DataFrame).all()

open(dataframe_type) lazy-loads the dataframe handlers, records dataframe_type, and returns the same wrapper so the call can be chained. all() raises ValueError("No dataframe type set. Use open() ...") if open() was not called. If the wrapper has a URI but no raw value, all() first resolves its StructuredDataset literal and then asks the dataset transformer to open it.

The internal path is:

DataFrame.open(pd.DataFrame)
    -> records pd.DataFrame and lazy-loads handlers
DataFrame.all()
    -> resolves a StructuredDataset literal when needed
DataFrameTransformerEngine.open_as()
    -> selects a decoder by dataframe type, URI protocol, and format
ParquetToPandasDecodingHandler or CSVToPandasDecodingHandler
    -> reads the dataset and returns pd.DataFrame

The same reference can be materialized as an Arrow table when PyArrow is available:

import pyarrow as pa
from flyte.io._dataframe.dataframe import DataFrame


async def read_arrow(df: DataFrame) -> pa.Table:
    return await df.open(pa.Table).all()

all_sync() runs all() through asyncio.run() for synchronous callers. iter() also requires open(), but it is not a general replacement for all(): it requires the selected decoder to return an async generator. The built-in pandas and Arrow decoders return one materialized dataframe/table, so using iter() with those handlers fails with the transformer's async-iterator check.

A raw dataframe task signature materializes during input conversion rather than exposing the wrapper. Conversely, when a task signature expects DataFrame, returning a raw pd.DataFrame is rejected by DataFrameTransformerEngine.to_literal(); wrap it with DataFrame.wrap_df() (or use a raw pandas signature).

Formats, schema metadata, and interoperability​

Flyte stores the dataframe as a StructuredDataset literal: the URI identifies the stored data and StructuredDataset metadata carries its format and columns. PARQUET, CSV, GENERIC_FORMAT, and GENERIC_PROTOCOL are the format/protocol constants used by the dataframe implementation. The built-in default format for pandas and Arrow is Parquet. CSV is available when the expected format is CSV, such as this task signature:

from typing import Annotated
import pandas as pd
from flyte.io._dataframe.dataframe import DataFrame


async def create_csv_dataframe() -> Annotated[DataFrame, "csv"]:
    return DataFrame.wrap_df(
        pd.DataFrame({"employee_id": [1001, 1002], "salary": [55000, 75000]})
    )

The expected task type is authoritative when formats are combined. In to_literal(), the transformer starts from the expected StructuredDataset type. If the expected format is generic, a non-generic format on a URI-backed wrapper can be retained; it does not blindly replace an explicit format from the task annotation.

Column metadata can be supplied with an OrderedDict annotation. The transformer converts the mapping into StructuredDataset columns, recursively flattening nested mappings and converting supported scalar/container types. Lists and dictionaries are handled recursively; unsupported column types raise an assertion such as type ... is currently not supported by DataFrame. DataFrame.columns() defaults to {}, and column_names() returns the keys from that mapping.

On input conversion, DataFrameTransformerEngine.to_python_value() combines the incoming literal metadata with the current task's requested columns. If the current task declares columns, those columns become the decoder's current_task_metadata; if it declares no columns, incoming columns are retained when present. The built-in Parquet decoders pass the resulting names as columns to pandas or PyArrow. This is the projection behavior tested by the dataframe transformer: a consumer can request one column from a two-column incoming dataset.

A CSV round trip exercises both sides of the abstraction:

import typing
import pandas as pd
from flyte.io._dataframe.dataframe import DataFrame
from flyte.types import TypeEngine
from flyte.io._dataframe.dataframe import DataFrameTransformerEngine


async def csv_round_trip(sample_dataframe: pd.DataFrame) -> pd.DataFrame:
    literal_type = TypeEngine.to_literal_type(typing.Annotated[pd.DataFrame, "csv"])
    wrapped = DataFrame.wrap_df(sample_dataframe)
    engine = DataFrameTransformerEngine()

    literal = await engine.to_literal(
        wrapped,
        python_type=pd.DataFrame,
        expected=literal_type,
    )
    lazy_value = await engine.to_python_value(literal, expected_python_type=DataFrame)
    return await lazy_value.open(pd.DataFrame).all()

By default, guess_python_type() returns the meta DataFrame for StructuredDataset literals. If the run context enables preserve_original_types=True and the literal contains a registered type tag, the transformer can restore the original registered dataframe type instead. This setting does not change the explicit open() target selected by a DataFrame consumer.

Built-in pandas and PyArrow handlers​

The built-in handlers are in flyte.io._dataframe.basic_dfs. They are registered lazily by flyte.io._dataframe.__init__: registration checks whether pandas or PyArrow has already been imported, then registers the corresponding handlers and renderer. Installing a package alone is not the registration trigger.

The registration functions in src/flyte/io/_dataframe/__init__.py register PandasToParquetEncodingHandler and ParquetToPandasDecodingHandler for pandas, PandasToCSVEncodingHandler and CSVToPandasDecodingHandler for CSV, and ArrowToParquetEncodingHandler and ParquetToArrowDecodingHandler for PyArrow. They also register TopFrameRenderer for pandas and ArrowRenderer for pa.Table.

Python valueEncoderStored layoutDecoder
pd.DataFramePandasToParquetEncodingHandler<uri>/00000ParquetToPandasDecodingHandler
pd.DataFrame with CSV formatPandasToCSVEncodingHandler<uri>/data.csvCSVToPandasDecodingHandler
pa.TableArrowToParquetEncodingHandler<uri>/00000ParquetToArrowDecodingHandler

Pandas Parquet​

PandasToParquetEncodingHandler is constructed for pd.DataFrame, all storage protocols (protocol=None), and PARQUET. It uses the wrapper URI when present or obtains a random raw-data path from internal_ctx().raw_data, creates a local directory when necessary, and calls:

df.to_parquet(
    path,
    coerce_timestamps="us",
    allow_truncated_timestamps=False,
    storage_options=get_pandas_storage_options(uri=path),
)

The file path is <uri>/00000, while the returned StructuredDataset URI is the directory URI. ParquetToPandasDecodingHandler reads that URI with pd.read_parquet(), passing requested columns and pandas storage options. A caught exception whose class name is exactly NoCredentialsError causes a retry with anonymous S3 options; other exceptions are re-raised.

Pandas CSV​

PandasToCSVEncodingHandler writes <uri>/data.csv, uses index=False, and sets the StructuredDataset format to CSV. CSVToPandasDecodingHandler accepts either a URI ending in .csv or a directory URI, in which case it appends data.csv. Its normal path streams the object through storage.get_stream() into BytesIO before calling pd.read_csv(buf). The source comment identifies this buffering pattern as a workaround for an obstore early-EOF issue with CSV. If the caught exception class name is exactly NoCredentialsError, it retries with pd.read_csv() and anonymous S3 storage options.

The CSV decoder computes requested column names, but its normal buffered pd.read_csv(buf) call does not pass them as usecols; the names are passed only in the anonymous-credentials fallback. This differs from the Parquet decoders and is relevant when diagnosing an unexpected full CSV read.

PyArrow Parquet​

ArrowToParquetEncodingHandler is constructed for pa.Table, all storage protocols, and PARQUET. It writes <uri>/00000 using a Flyte storage filesystem and passes a protocol-stripped path to pyarrow.parquet.write_table(). ParquetToArrowDecodingHandler passes requested column names to pyarrow.parquet.read_table().

For local paths, the Arrow decoder creates the parent directory and strips the protocol before reading. For remote paths it uses the URI and Flyte's storage filesystem behavior. On an exception whose class name is exactly NoCredentialsError, it obtains an anonymous filesystem and retries; if no filesystem can be obtained, that handler returns None.

Handler lookup and extension points​

DataFrameTransformerEngine is the higher-level TypeTransformer for all dataframe types. It keeps class-level ENCODERS, DECODERS, DEFAULT_FORMATS, and DEFAULT_PROTOCOLS maps and registers dataframe types with the main TypeEngine when a handler is registered.

Lookup first tries an exact dataframe-type, protocol, and format match. It can then use an fsspec registration, a generic-format handler, a configured default format, or the single registered handler when the requested format is generic. If no candidate matches, it raises ValueError identifying the dataframe type, protocol, and format.

DataFrameEncoder and DataFrameDecoder are the extension interfaces. An encoder is initialized with a Python dataframe type, optional protocol, and supported format, then implements async encode(DataFrame, StructuredDatasetType) -> StructuredDataset. A decoder implements async decode(StructuredDataset, StructuredDatasetMetadata) and may return either one dataframe or an async iterator. Set protocol=None to support all storage protocols; an empty supported format means a generic fallback. Registering a None-protocol handler with default_for_type=True is invalid because that option requires a specific default storage protocol. Duplicate registrations require override=True.

External dataframe integrations use these interfaces rather than registering independently with the core type engine. The DataFrame docstring names flyteplugins-polars and flyteplugins-spark as examples of ecosystem plugins; each supplies handlers for its dataframe type through the same transformer.

Operational checks​

  • open() omitted: Call df.open(pd.DataFrame) or df.open(pa.Table) before all() or iter(); both operations require a selected dataframe type.
  • Pandas or PyArrow not recognized: Import the library in the task environment. Handler loading is lazy and only registers pandas or Arrow handlers when the corresponding package is already imported.
  • Raw value supplied for a meta-DataFrame signature: Wrap it with DataFrame.wrap_df(); to_literal() rejects a non-DataFrame value when the declared type is DataFrame.
  • Local value outside a task: Use DataFrame.from_local()/from_local_sync() or provide a destination URI. If there is no raw-data context and no URI, the transformer cannot infer storage for a local object by itself.
  • Unexpected URI layout: Pandas and Arrow Parquet use a dataset directory containing 00000; pandas CSV uses a directory containing data.csv. Pass the dataset directory URI when using the generated Parquet outputs.
  • S3 access failures: Credentialed reads and writes use Flyte storage/fsspec configuration. Anonymous retry is narrow and occurs only when the exception class name is exactly NoCredentialsError; other storage errors remain errors.
  • Column annotation failure: Restrict annotated columns to the scalar and nested list/dictionary types supported by the dataframe transformer. Unsupported types fail while the column mapping is converted to StructuredDataset metadata.
  • Original type not returned: StructuredDataset deserialization returns the meta DataFrame by default. Enable preserve_original_types=True when the run context should use a registered original dataframe type and the literal contains a matching type tag.