Skip to main content

Data I/O

Choose the right data reference​

Use flyte.io.File for one object, flyte.io.Dir for a directory tree, and flyte.io.DataFrame for a structured dataset that can be encoded by a registered dataframe handler. The following task writes bytes directly to task storage and returns a File reference; it does not first create a local file:

@env.task
async def create_new_remote_file(content: str) -> File:
    f = File.new_remote()
    async with f.open("wb") as fh:
        await fh.write(content.encode("utf-8"))
    print(f"Created new remote file at: {f.path}")
    return f


@env.task
async def read_file_async(f: File) -> str:
    async with f.open("rb") as fh:
        contents = bytes(await fh.read())  # read() returns a memoryview
        text_content = contents.decode("utf-8")
        return text_content

File and Dir are Pydantic BaseModel/SerializableType references. Their literal transformers serialize references as Flyte blobs; ordinary uploads, reads, and downloads happen through their explicit methods (or through a lazy uploader created by from_local). The path field may therefore be local or remote, and name defaults to the basename of that path.

Read and write files​

Stream a remote file​

Use the async API in an async task and the _sync API in a synchronous task. File.open() is an async context manager; File.open_sync() is a normal context manager.

@env.task
def read_file_sync(f: File) -> str:
    with f.open_sync("rb") as fh:
        content = fh.read()
        return content.decode("utf-8")


@env.task
async def write_file() -> File:
    f = File.new_remote()
    async with f.open("wb") as fh:
        await fh.write(b"Hello, World!")
    return f

open() accepts mode, block_size, cache_type, cache_options, compression, and additional fsspec open arguments. open_sync() exposes the same options. Async reads commonly return a memoryview, so convert it with bytes(...) before decoding or passing it to an API that expects bytes.

File.new_remote(file_name=None, hash_method=None) obtains a generated raw-data path from the active task context. File.named_remote(name) instead derives a stable path from the active task/run context and is intended for retry-stable outputs. It raises ValueError when there is no active task execution context. File.from_existing_remote(remote_path, file_cache_key=None) only creates a reference to an existing object:

@env.task
async def process_existing_file() -> str:
    file = File.from_existing_remote("s3://my-bucket/data.csv")
    async with file.open("rb") as f:
        content = await f.read()
        return bytes(content).decode("utf-8")

Upload and download local files​

Use File.from_local() in an async task and File.from_local_sync() in a synchronous task. Both require the local file to exist; a missing path raises ValueError. An optional remote_destination selects the destination, while hash_method can supply a cache-key computation or a precomputed string.

@env.task
async def upload_file() -> File:
    with open("/tmp/data.csv", "w") as f:
        f.write("col1,col2\n1,2\n3,4\n")
    return await File.from_local("/tmp/data.csv")


@env.task
def upload_file_sync() -> File:
    with open("/tmp/data.csv", "w") as f:
        f.write("col1,col2\n1,2\n3,4\n")
    return File.from_local_sync("/tmp/data.csv")

Outside a raw-data task context, from_local() and from_local_sync() install a lazy uploader. The local path can remain unchanged until dataframe/file literal conversion invokes that callback; in the main local run mode the callback can deliberately return the local path without uploading. Within a task with raw-data storage, the methods upload immediately or use the configured local storage optimization.

Download to a generated local path or provide a destination:

@env.task
async def download_file(f: File) -> str:
    local_path = await f.download()
    with open(local_path, "r") as fh:
        return fh.read()


@env.task
def download_file_sync(f: File) -> str:
    local_path = f.download_sync("/tmp/myfile.csv")
    with open(local_path, "r") as fh:
        return fh.read()

exists() and exists_sync() query the underlying storage filesystem. For a local source, downloading to the same path returns or copies through the local filesystem; for a remote source, the storage layer performs the download.

Attach a custom cache key​

HashFunction adapts a callable to Flyte's HashMethod interface:

def my_hash(data: bytes) -> str:
    return hashlib.md5(data).hexdigest()


hash_fn = HashFunction.from_fn(my_hash)
hash_fn.update(b"hello")
assert hash_fn.result() == "5d41402abc4b2a76b9719d911017c592"

update(data) calls the supplied function and replaces the previously stored result; it is not an incremental accumulator. File.from_local*, dataframe uploads, and hashing writers use HashMethod implementations to populate literal hashes.

Upload and traverse directories​

Dir.from_local() recursively uploads a local directory and Dir.from_existing_remote() creates a reference without uploading. The async lifecycle is:

@env.task
async def upload_directory_from_local() -> Dir:
    local_path = await create_test_local_directory()
    uploaded_dir = await Dir.from_local(local_path)
    print(f"Uploaded local directory {local_path} to remote: {uploaded_dir.path}")
    return uploaded_dir


@env.task
async def list_files_in_directory(d: Dir) -> list[File]:
    files = await d.list_files()
    print(f"Found {len(files)} files in directory {d.path}:")
    for file in files:
        print(f"  - {file.name}: {file.path}")
    return files


@env.task
async def walk_directory_async(d: Dir) -> list[File]:
    all_files = []
    async for file in d.walk(recursive=True):
        print(f"  Found file: {file.name} at {file.path}")
        all_files.append(file)
    return all_files

list_files() and list_files_sync() are non-recursive. walk() defaults to recursive traversal and supports recursive and max_depth; walk(recursive=False) internally limits the walk to the top level. Each yielded item is a File reference and can be opened without downloading the whole directory. The synchronous equivalents are walk_sync(), list_files_sync(), get_file_sync(), download_sync(), and exists_sync().

walk_sync() has a file_pattern argument in its public signature, but the current implementation does not apply that pattern when yielding files. Do not rely on file_pattern="*.txt" to filter results. Directory lookup also uses filesystem-specific separators: async get_file() joins with the filesystem separator, while get_file_sync() uses os.path.join.

Upload concurrency can be supplied to the async method:

uploaded_dir = await Dir.from_local(
    "/tmp/data_dir/",
    remote_destination="s3://my-bucket/data/",
    batch_size=16,
)

When batch_size is omitted, the storage upload default is controlled by FLYTE_IO_BATCH_SIZE (default 32). Dir.from_local_sync() has the same upload/reference behavior but does not expose batch_size.

Download a complete tree when local processing needs all of its files:

@env.task
async def download_directory(d: Dir) -> str:
    local_dir = await d.download("/tmp/my_data/")
    return local_dir


@env.task
def download_directory_sync(d: Dir) -> str:
    local_dir = d.download_sync("/tmp/my_data/")
    return local_dir

Dir.new_remote(dir_name=None, hash=None) allocates a remote directory reference for output written directly into task storage. Dir.empty() returns an EmptyDir sentinel whose path is ::flyte-empty-dir::. After serialization and deserialization, the result is a plain Dir, so test result.is_empty rather than isinstance(result, EmptyDir). This sentinel is the documented alternative to Optional[Dir], which cannot round-trip through the dataclass transformer in this flow.

Pass structured dataframes​

flyte.io.DataFrame is a lazy wrapper around a registered native dataframe type. DataFrame.wrap_df() (and the deprecated from_df() used by existing examples) wraps an in-memory value; wrapping alone does not encode or upload it. Native dataframe values can also be used directly when their plugin transformer is registered.

@env.task
def process_df(df: pl.DataFrame) -> pl.DataFrame:
    return df


@env.task
def process_fdf_to_df(df: flyte.io.DataFrame) -> pl.DataFrame:
    return df


@env.task
def process_df_to_fdf(df: pl.DataFrame) -> flyte.io.DataFrame:
    return flyte.io.DataFrame.from_df(df)


@env.task
def process_lf(lf: pl.LazyFrame) -> pl.LazyFrame:
    return lf


@env.task
def process_fdf_to_lf(df: flyte.io.DataFrame) -> pl.LazyFrame:
    return df

When the declared type is flyte.io.DataFrame, returning an unwrapped native dataframe causes DataFrameTransformerEngine.to_literal() to raise TypeTransformerFailedError; wrap it first. Conversely, when the declared type is a native registered type such as pl.DataFrame, the engine invokes that type's encoder. DataFrame.from_existing_remote() creates a structured-dataset reference, and DataFrame.from_local()/from_local_sync() attach metadata and upload immediately in a task context or install a lazy uploader outside one.

A wrapper can remain lazy on the receiving side. DataFrame.open(dataframe_type) selects the native target type, all() and all_sync() materialize it, and iter() uses an async-iterator decoder when one is registered. If a wrapper is passed through unchanged, the engine preserves its original structured-dataset literal rather than invoking another encoder.

Column and format annotations become StructuredDatasetType metadata. For example, the Spark integration uses kwtypes to request a column schema:

columns = kwtypes(name=str, age=int)


@spark_env.task
async def sum_of_all_ages(sd: Annotated[pyspark.sql.DataFrame, columns]) -> int:
    total_age = sd.groupBy().sum("age").collect()[0][0]
    print("Total age sum:", total_age)
    return total_age

Before decoding, DataFrameTransformerEngine.to_python_value() uses consumer-side columns when present; otherwise it carries forward incoming literal columns. Decoders receive the resulting metadata and are expected to honor it. A wrapper with a URI is encoded as a structured dataset reference; a local wrapper sent to a remote run needs DataFrame.from_local(...) or an equivalent usable task storage context.

Add a dataframe handler​

Implement DataFrameEncoder for Python-to-Flyte conversion and DataFrameDecoder for Flyte-to-Python conversion, then register both with DataFrameTransformerEngine. The encoder constructor takes the native dataframe type, an optional protocol, and an optional format. Its encode() receives a Flyte DataFrame wrapper and a protobuf StructuredDatasetType, and returns a protobuf StructuredDataset. The decoder's decode() receives a protobuf StructuredDataset and current task metadata and returns a native dataframe or an async iterator.

Registration is the dispatch contract:

  • protocol=None registers an all-fsspec handler under the generic fsspec entry. It cannot be registered with default_for_type=True.
  • protocol="s3://" and protocol="s3" normalize to s3; / normalizes to file.
  • supported_format=None becomes the generic empty format ("").
  • default_for_type=True sets both default format and storage; default_format_for_type=True and default_storage_for_type=True set only the corresponding default.
  • Duplicate (dataframe type, protocol, format) registrations require override=True; an empty protocol string is rejected in favor of None.

Handler lookup first tries an exact protocol and format, then protocol-specific generic/default formats, then generic fsspec handlers and their generic/default formats. open_as() calls a decoder that returns one dataframe. iter_as() is stricter: the decoder result must be an actual async generator or the engine raises ValueError.

Use the bundled dataframe integrations​

Polars and Parquet​

The Polars plugin registers generic-protocol Parquet handlers for both eager and lazy frames. Eager pl.DataFrame values are written with write_parquet() to <uri>/00000; pl.LazyFrame values use sink_parquet() and write to <uri>/00000.parquet. The corresponding decoders use read_parquet() and scan_parquet(). The lazy decoder applies requested columns with select(*columns) without eagerly materializing the frame.

The different filenames are observable in the round-trip tests: the eager test asserts 00000 exists and 00000.parquet does not, while the lazy test restores a pl.LazyFrame. Partitioning is not implemented in these handlers. Storage credentials are converted to Polars' flat string-key options for supported object stores; when an S3 read raises an exception whose class name is NoCredentialsError, the handlers retry with anonymous S3 options.

Spark and Parquet​

SparkToParquetEncoder obtains a Spark session, sets mapreduce.fileoutputcommitter.marksuccessfuljobs to false, and writes the wrapped raw Spark dataframe with overwrite mode. It chooses the wrapper URI first, then task or internal raw-data paths. ParquetToSparkDecoder reads with spark.read.parquet(path) and calls .select(*columns) when current task metadata declares columns; with no declared columns it returns the full read.

Snowflake and pandas​

The Snowflake adapter registers pd.DataFrame for protocol snowflake and the generic format. Its write URI is positional:

snowflake://user/account/warehouse/database/schema/table

The encoder parses those components, obtains a connector connection, and calls write_pandas(). The read path uses the same prefix with a query ID in the final component, calls cursor.get_results_from_sfqid(query_id), and returns fetch_pandas_all(). The read handler does not use current task columns to subset the result.

A task configures the connection fields and injects key-pair material as an environment variable like this:

sf_config = SnowflakeConfig(
    user="KEVIN",
    account="PWGJLTH-XKB21544",
    database="FLYTE",
    schema="PUBLIC",
    warehouse="COMPUTE_WH",
)

env = flyte.TaskEnvironment(
    name="snowflake_example_env",
    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-snowflake"),
    secrets=[flyte.Secret(key="snowflake", as_env_var="SNOWFLAKE_PRIVATE_KEY")],
)


@env.task
async def downstream_task(df: pd.DataFrame) -> float:
    return df["AGE"].mean().item()

The adapter reads SNOWFLAKE_PRIVATE_KEY and optional SNOWFLAKE_PRIVATE_KEY_PASSPHRASE; the PEM key is converted to DER for connector key-pair authentication. Other authentication options belong in Snowflake connector configuration.

Troubleshoot data I/O​

  • The file read cannot be awaited. Use async with file.open(...) and await fh.read() in async tasks. Use with file.open_sync(...) and fh.read() in sync tasks.
  • A local file was not uploaded yet. Outside raw-data context, from_local* may return a reference with a lazy uploader. Conversion to a literal invokes it; local run mode may intentionally keep the local path.
  • A native dataframe is rejected for flyte.io.DataFrame. Wrap it with DataFrame.from_df(...) (or the current wrap_df() API). Native values are valid when the task annotation is the registered native dataframe type.
  • No dataframe handler matches. Check the native type, normalized protocol, and format. Use protocol=None for an all-fsspec handler, not protocol="", and register duplicates with override=True.
  • Polars cannot find an object. Verify the eager/lazy filename convention and object-store credentials. The anonymous retry is limited to S3 NoCredentialsError handling.
  • Snowflake parsing fails. Verify the positional seven-component write URI and the query-ID final component used for reads; an absent or empty URI raises ValueError.
  • A directory pattern did not filter files. Dir.walk_sync(file_pattern=...) currently accepts but ignores the pattern. Filter the yielded File objects in task code instead.
  • An optional directory does not deserialize. Return Dir.empty() or EmptyDir() and branch on is_empty after the value crosses the Flyte boundary.