flyte-sdk-data

Flyte 2 SDK Data Engineering Skill

Build ETL pipelines, data processing workflows, and data quality checks with Flyte 2.

Grounding References

Ground unfamiliar APIs in real examples. When unsure of a current Flyte 2 API, or for a pattern not shown below, and the flyte-docs search tools are available, search them first — by exact symbol (TaskEnvironment, flyte.io.File, map_task), since matching is literal substring, not semantic — then adapt a real example rather than inventing one, and cite the file or section you pulled it from. (Flyte 2 is not flytekit; priors are often wrong.)

ETL Pipeline Patterns

Basic Extract-Transform-Load

import flyte
import flyte.io

env = flyte.TaskEnvironment(
    name="etl-pipeline",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "pandas", "polars", "pyarrow", "boto3", "sqlalchemy",
    ),
)

@env.task(retries=3, cache="auto")
async def extract(source_uri: str) -> flyte.io.DataFrame:
    """Extract data from various sources."""
    import polars as pl
    if source_uri.endswith(".csv"):
        df = pl.read_csv(source_uri)
    elif source_uri.endswith(".parquet"):
        df = pl.read_parquet(source_uri)
    else:
        raise ValueError(f"Unsupported format: {source_uri}")
    return flyte.io.DataFrame(df)

@env.task(retries=2, cache="auto")
async def transform(df: flyte.io.DataFrame) -> flyte.io.DataFrame:
    """Clean and transform data."""
    inner = df.to_polars()
    cleaned = (
        inner
        .drop_nulls()
        .unique()
        .with_columns([
            pl.col("date").str.strptime(pl.Date, "%Y-%m-%d").alias("date_parsed"),
        ])
    )
    return flyte.io.DataFrame(cleaned)

@env.task(retries=1, cache="auto")
async def load(df: flyte.io.DataFrame, destination: str) -> str:
    """Load transformed data to destination."""
    inner = df.to_polars()
    if destination.endswith(".parquet"):
        inner.write_parquet(destination)
    elif destination.endswith(".csv"):
        inner.write_csv(destination)
    return destination

@env.task
async def etl_pipeline(source_uri: str, destination: str) -> dict:
    """Orchestrate the ETL pipeline."""
    raw = await extract(source_uri)
    cleaned = await transform(raw)
    loaded_path = await load(cleaned, destination)
    return {"source": source_uri, "destination": loaded_path}