Skip to main content

Data Management & Type System

The Data Management & Type System provides robust and extensible capabilities for handling diverse data types, managing local and remote storage, and ensuring seamless serialization and deserialization within the platform. This system is designed to abstract away the complexities of data persistence and type conversion, allowing developers to focus on their application logic.

File and Directory Management​

The system offers generic representations for files and directories, enabling consistent interaction with data regardless of its physical location (local filesystem or remote object storage).

The File Class​

The File class represents a single file. It is generic (File[T]), allowing users to specify the expected format of the file's content (e.g., File[DataFrame] for a CSV file containing DataFrame data).

Capabilities:

  • Path Management: Stores the path to the file, which can be local (e.g., /tmp/data.csv) or remote (e.g., s3://my-bucket/data.csv).
  • Asynchronous and Synchronous I/O: Provides open() (async) and open_sync() (sync) context managers to access the file content as a file-like object. Users are responsible for reading from or writing to these file handles.
  • Download/Upload:
    • download(): Asynchronously downloads the file to a local path.
    • from_local(): Asynchronously uploads a local file to a remote store, returning a new File instance.
    • new_remote(): Creates a File instance pointing to a new, randomly generated remote path, useful for writing directly to remote storage.
    • from_existing_remote(): Creates a File instance referencing an already existing remote file.
  • Existence Check: exists_sync() checks if the file exists.
  • Hashing: Includes hash and hash_method attributes, which can be used to store or compute a hash of the file's content. This is crucial for data integrity checks and caching mechanisms.

Usage Examples:

from pandas import DataFrame
from flyte.io import File

# Async usage: Read a remote CSV file
async def read_csv_async():
    csv_file = File[DataFrame](path="s3://my-bucket/data.csv")
    async with csv_file.open() as f:
        content = await f.read()
        # Process content, e.g., df = pd.read_csv(f)

# Sync usage: Read a local CSV file
def read_csv_sync():
    csv_file = File[DataFrame](path="/tmp/local_data.csv")
    with csv_file.open_sync() as f:
        content = f.read()
        # Process content

# Upload a local file
async def upload_file():
    remote_file = await File[DataFrame].from_local("/tmp/local_data.csv", "s3://my-bucket/uploaded_data.csv")

# Write directly to a new remote file
async def write_remote_file():
    import pandas as pd
    df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})
    file = File.new_remote()
    async with file.open("wb") as f:
        df.to_csv(f)
    return file

The FileTransformer handles the serialization of File objects to and from the platform's Literal representation, specifically as a BlobType.SINGLE literal, storing the file's URI and format.

The Dir Class​

The Dir class represents a directory, similar to File but for collections of files. It is also generic (Dir[T]), indicating the format of files expected within the directory.

Capabilities:

  • Path Management: Stores the path to the directory.
  • Asynchronous and Synchronous Traversal:
    • walk(): Asynchronously iterates through files in the directory, optionally recursively.
    • walk_sync(): Synchronously iterates through files in the directory.
  • Listing Files: list_files() (async) and list_files_sync() (sync) retrieve a non-recursive list of files.
  • Download/Upload:
    • download(): Asynchronously downloads the entire directory to a local path.
    • from_local(): Asynchronously uploads a local directory to a remote store.
    • from_existing_remote(): Creates a Dir instance referencing an existing remote directory.
  • Existence Check: exists() (async) and exists_sync() (sync) check if the directory exists.
  • File Access: get_file() (async) and get_file_sync() (sync) retrieve a specific File object within the directory.
  • Hashing: Includes a hash attribute for caching and discovery.

Usage Examples:

from pandas import DataFrame
from flyte.io import Dir

# Async usage: Walk through files in a remote directory
async def process_dir_async():
    data_dir = Dir[DataFrame](path="s3://my-bucket/data/")
    async for file in data_dir.walk():
        local_path = await file.download()
        # Process the downloaded file

# Sync usage: List files in a local directory
def list_dir_sync():
    local_dir = Dir[DataFrame](path="/tmp/my_data_dir/")
    files = local_dir.list_files_sync()
    for file in files:
        print(f"Found file: {file.path}")

# Upload a local directory
async def upload_directory():
    remote_dir = await Dir[DataFrame].from_local("/tmp/local_data_dir/", "s3://my-bucket/uploaded_data/")

The DirTransformer handles the serialization of Dir objects to and from the platform's Literal representation, specifically as a BlobType.MULTIPART literal, storing the directory's URI and format.

Important Consideration: For both File and Dir, the type transformers do not automatically upload or download data. Users are explicitly responsible for invoking methods like download(), from_local(), or using open() to manage I/O.

Structured Data with DataFrame​

The DataFrame class provides a flexible and extensible way to handle structured datasets, such as Pandas DataFrames or PyArrow Tables, within the platform. It acts as a wrapper around the actual DataFrame object, managing its URI, file format, and metadata.

Core DataFrame Class​

The DataFrame class itself is a user-facing wrapper. It holds:

  • uri: The storage location of the dataset (local or remote).
  • file_format: The format of the stored data (e.g., "parquet", "csv").
  • val: An optional attribute to hold the in-memory Python DataFrame object.

Capabilities:

  • Lazy Loading: The actual DataFrame object (val) is loaded only when explicitly requested via open().all().
  • Format and Protocol Agnostic: Supports various underlying DataFrame libraries (e.g., Pandas, PyArrow) and storage protocols (e.g., S3, GCS, local filesystem) through a pluggable handler system.
  • Column Schema: Can carry column type information, enabling schema enforcement and column subsetting during deserialization.
  • Serialization/Deserialization: Handles conversion to and from Flyte's StructuredDataset literal type.

Usage Examples:

import pandas as pd
from flyte.io import DataFrame
from flytekit.types.structured.structured import StructuredDataset

# Define a task that takes and returns a DataFrame
# The type annotation can include column information for schema enforcement
# and format specification.
@task
async def process_data(input_df: Annotated[DataFrame, {"col1": int, "col2": str}, "parquet"]) -> DataFrame:
    # Load the DataFrame into a Pandas DataFrame
    df = await input_df.open(pd.DataFrame).all()
    print(f"Processing DataFrame with columns: {df.columns.tolist()}")
    # Perform operations
    df["col3"] = df["col1"] * 2
    # Return the modified DataFrame. The system will automatically encode it.
    return DataFrame(val=df)

# Example of creating a DataFrame from an existing remote URI
@task
async def get_remote_data() -> DataFrame:
    return DataFrame(uri="s3://my-bucket/processed_data.parquet", file_format="parquet")

# Example of iterating over a large DataFrame (if supported by decoder)
@task
async def iterate_data(input_df: DataFrame):
    async for chunk in input_df.open(pd.DataFrame).iter():
        print(f"Processing chunk of size: {len(chunk)}")

Extensibility with DataFrameTransformerEngine​

The DataFrameTransformerEngine is a meta-transformer that manages the conversion of DataFrame objects. It uses a pluggable system of DataFrameEncoder and DataFrameDecoder implementations to support different DataFrame libraries and file formats.

  • DataFrameEncoder: Defines how a Python DataFrame object (e.g., pd.DataFrame) is converted into a Flyte StructuredDataset literal. This involves writing the data to a specified URI in a particular format.
    • Example: PandasToParquetEncodingHandler converts a Pandas DataFrame to Parquet format.
  • DataFrameDecoder: Defines how a Flyte StructuredDataset literal (a URI and format) is converted back into a Python DataFrame object. This involves reading data from the URI and parsing it into the target DataFrame library.
    • Example: ParquetToPandasDecodingHandler reads Parquet data into a Pandas DataFrame.

Developers can register custom encoders and decoders for new DataFrame types or formats using DataFrameTransformerEngine.register().

Handler Discovery: The engine intelligently selects the appropriate encoder/decoder based on the Python type, storage protocol (e.g., "s3", "file"), and file format (e.g., "parquet", "csv"). It supports fallback mechanisms for generic formats or protocols.

Type System Extensibility (TypeEngine and TypeTransformer)​

The TypeEngine is the central component for managing the conversion between native Python types and Flyte's platform-agnostic Literal representation (based on Protocol Buffers). This enables tasks written in Python to seamlessly exchange data with tasks written in other languages or with the Flyte control plane.

The TypeEngine​

The TypeEngine acts as a registry and dispatcher for TypeTransformer instances.

Key Capabilities:

  • Type Registration: register() and register_additional_type() methods allow developers to extend the type system by associating Python types with custom TypeTransformer implementations.
  • Python to Literal Conversion:
    • to_literal_type(): Converts a Python type hint (e.g., int, List[str], DataFrame) into a Flyte LiteralType (a schema-like definition).
    • to_literal(): Converts a Python value (e.g., 42, ["a", "b"], a DataFrame instance) into a Flyte Literal (the actual data payload). This process can involve serialization to binary formats or uploading to remote storage.
  • Literal to Python Conversion:
    • guess_python_type(): Infers the Python type from a Flyte LiteralType.
    • to_python_value(): Converts a Flyte Literal back into a native Python value. This can involve downloading data, deserializing binary formats, and reconstructing Python objects.
  • Complex Type Handling: Includes built-in transformers for common complex Python types:
    • ListTransformer: Handles typing.List[T].
    • DictTransformer: Handles typing.Dict[str, T] (as LiteralMap) and untyped dictionaries (as SimpleType.STRUCT with MessagePack).
    • UnionTransformer: Handles typing.Union[T1, T2, ...], including Optional[T].
  • Specialized Type Handling:
    • FileTransformer and DirTransformer: As described above.
    • DataFrameTransformerEngine: As described above.
    • FlytePickleTransformer: A fallback transformer that serializes any unrecognized Python type using cloudpickle and stores it as a BlobType.SINGLE literal. This allows arbitrary Python objects to be passed, though it's generally discouraged for cross-language compatibility.
    • PydanticTransformer: Automatically handles Pydantic BaseModel instances, converting them to/from binary literals using MessagePack and generating JSON schemas for metadata.
    • DataclassTransformer: Handles Python dataclass instances similarly to Pydantic models, using MessagePack and JSON schema generation.
    • EnumTransformer: Converts Python enum.Enum types to/from string literals.
    • ProtobufTransformer: Handles Google Protocol Buffer Message types, converting them to/from SimpleType.STRUCT literals.
    • SimpleTransformer: A utility for quickly defining transformers for simple types with custom serialization/deserialization functions.
    • RestrictedTypeTransformer: Prevents certain Python types from being used as task inputs/outputs, raising RestrictedTypeError if attempted.
  • HTML Rendering: to_html() provides a mechanism to render Python values as HTML, useful for UI display.
  • Offloaded Literals: Supports offloading large literals to remote storage, with unwrap_offloaded_literal() handling the retrieval.

The TypeTransformer Interface​

TypeTransformer is an abstract base class that defines the contract for converting a specific Python type.

Key Methods for Implementers:

  • get_literal_type(t: Type[T]) -> LiteralType: Returns the Flyte LiteralType corresponding to the Python type t.
  • to_literal(python_val: T, python_type: Type[T], expected: LiteralType) -> Literal: Converts a Python value python_val of type python_type into a Flyte Literal.
  • to_python_value(lv: Literal, expected_python_type: Type[T]) -> Optional[T]: Converts a Flyte Literal lv into a Python value of expected_python_type.
  • guess_python_type(literal_type: LiteralType) -> Type[T]: (Optional) Attempts to infer the Python type from a Flyte LiteralType.
  • assert_type(t: Type[T], v: T): Performs type validation.
  • from_binary_idl(binary_idl_object: Binary, expected_python_type: Type[T]) -> Optional[T]: Handles deserialization from binary literal payloads (e.g., MessagePack).
  • to_html(python_val: T, expected_python_type: Type[T]) -> str: Converts the Python value to an HTML string for display.

Integration Points: When defining custom types or integrating external libraries, developers implement a TypeTransformer and register it with the TypeEngine. This allows the platform to automatically handle serialization and deserialization for these custom types.

Data Integrity and Caching (Hashing)​

The system incorporates hashing mechanisms to ensure data integrity and enable efficient caching of data artifacts.

  • HashMethod: A protocol defining the interface for hash accumulators (update and result).
  • HashlibAccumulator: An implementation of HashMethod that wraps Python's hashlib for common hashing algorithms (e.g., SHA256, MD5).
  • PrecomputedValue: Allows providing a precomputed hash string, useful when the hash is known beforehand.
  • HashingReader / AsyncHashingReader: Synchronous and asynchronous file-like objects that wrap an underlying file handle. As data is read, they update an associated HashMethod accumulator, allowing the hash of the read content to be computed on the fly.
  • HashingWriter / AsyncHashingWriter: Similar to readers, these wrap file handles for writing. As data is written, they update an associated HashMethod accumulator, computing the hash of the written content.

Usage: The File and Dir classes utilize these hashing utilities. When a File or Dir is created or uploaded, a hash can be associated with it. This hash can then be used by the platform for caching task outputs or for data discovery, ensuring that tasks only re-execute if their inputs (or their hashes) have changed.

Storage Configuration​

The storage configuration layer provides a unified way to interact with various cloud storage providers and local filesystems, leveraging fsspec for abstraction.

  • Storage: The base class for general storage configuration, including retry logic (retries, backoff) and debug settings (enable_debug).
  • Provider-Specific Configurations:
    • S3: Configures access to Amazon S3, including endpoint, access keys, and secret keys. It supports sandbox environments (e.g., MinIO) and anonymous access.
    • GCS: Configures access to Google Cloud Storage, with options like gsutil_parallelism.
    • ABFS: Configures access to Azure Blob Storage, including account name, keys, and client/tenant IDs.

Configuration: These storage configurations can be automatically loaded from environment variables using the auto() class method, or explicitly configured.

get_fsspec_kwargs(): Each storage configuration class provides this method to translate its settings into keyword arguments compatible with fsspec filesystem constructors. This allows the underlying data persistence layer to seamlessly connect to different storage backends.

RemoteFSPathResolver: This utility manages mappings between internal flyte:// URIs and their actual remote storage paths, ensuring consistent path resolution within the platform.

Best Practices and Considerations​

  • Asynchronous I/O: For performance-critical tasks, especially when dealing with remote storage, prioritize the asynchronous open(), download(), walk(), and from_local() methods provided by File and Dir.
  • Explicit I/O: Remember that File and Dir transformers do not perform automatic data transfers. Always explicitly call download(), from_local(), or use the open() context manager to interact with the data.
  • Type Annotations: Leverage Python's type hints extensively. The TypeEngine relies heavily on accurate type annotations to correctly infer and transform data. For DataFrame, use Annotated to specify column schemas and file formats for enhanced type safety and data handling.
  • Custom Types: When introducing custom Python objects, implement a TypeTransformer and register it with the TypeEngine to ensure they are correctly serialized and deserialized. For structured data, consider implementing DataFrameEncoder and DataFrameDecoder.
  • Hashing for Caching: Utilize the hash attribute in File and Dir to enable caching and data versioning, especially for large or frequently accessed datasets.
  • Storage Configuration: Configure storage settings (e.g., S3 credentials, GCS parallelism) appropriately for your execution environment to ensure reliable and performant data access.
  • Limitations: Synchronous remote downloads/uploads for Dir are currently not fully implemented and will raise NotImplementedError. Prefer asynchronous operations for remote directory transfers.