JSONL

pip install flyteplugins-jsonl

import flyte
from flyteplugins.jsonl import JsonlDir, JsonlFile

env = flyte.TaskEnvironment(
    name="jsonl-examples",
    image=flyte.Image.from_debian_base(name="jsonl").with_pip_packages(
        "flyteplugins-jsonl"
    ),
)

@env.task
async def write_records() -> JsonlFile:
    """Write records to a single JSONL file."""
    out = JsonlFile.new_remote("results.jsonl")
    async with out.writer() as writer:
        for i in range(500_000):
            await writer.write({"id": i, "score": i * 0.1})
    return out

@env.task
async def read_records(data: JsonlFile) -> int:
    """Read records from a JsonlFile and return the count."""
    count = 0
    async for record in data.iter_records():
        count += 1
    return count

@env.task
async def write_bulk() -> JsonlFile:
    """Hand the writer a whole list instead of one record at a time.

    write_many() serializes every record into the writer's buffer and checks
    the flush threshold once, instead of once per record.
    """
    out = JsonlFile.new_remote("bulk.jsonl")
    records = [{"id": i, "score": i * 0.1} for i in range(200_000)]
    async with out.writer() as writer:
        await writer.write_many(records)
    return out

@env.task
async def passthrough(source: JsonlFile) -> JsonlFile:
    """Forward pre-serialized bytes without a re-encode round trip."""
    import orjson

    out = JsonlFile.new_remote("passthrough.jsonl")
    async with out.writer() as writer:
        async for record in source.iter_records():
            # write_raw() takes bytes verbatim. Unlike write() and write_many(),
            # it does not append the record separator, so include it yourself.
            await writer.write_raw(orjson.dumps(record) + b"\n")
    return out

@env.task
async def write_large_dataset() -> JsonlDir:
    """Write a large dataset to a sharded JsonlDir.

    JsonlDir automatically rotates to a new shard file once the
    current shard reaches the record or byte limit. Shards are named
    part-00000.jsonl, part-00001.jsonl, etc.
    """
    out = JsonlDir.new_remote("dataset/")
    async with out.writer(
        max_records_per_shard=100_000,
        max_bytes_per_shard=256 * 1024 * 1024,  # 256 MB
    ) as writer:
        for i in range(500_000):
            await writer.write({"index": i, "value": i * i})
    return out

@env.task
async def append_more(dataset: JsonlDir) -> JsonlDir:
    """Add shards to a directory that already has some.

    Opening a writer on a populated JsonlDir is safe: it scans for existing
    part-NNNNN shards and starts numbering at the next free index, so nothing
    already written is overwritten.
    """
    async with dataset.writer(max_records_per_shard=100_000) as writer:
        for i in range(50_000):
            await writer.write({"index": i, "appended": True})
    return dataset

@env.task
async def sum_values(dataset: JsonlDir) -> int:
    """Read all records across all shards and compute a sum.

    Iteration is transparent across shards and handles mixed
    compressed/uncompressed shards automatically. The next shard is
    prefetched in the background for higher throughput.
    """
    total = 0
    async for record in dataset.iter_records():
        total += record["value"]
    return total

@env.task
async def process_in_batches(dataset: JsonlDir) -> int:
    """Process records in batches of dicts for bulk operations."""
    total = 0
    async for batch in dataset.iter_batches(batch_size=1000):
        # Each batch is a list[dict]
        total += len(batch)
    return total

@env.task
async def write_compressed() -> JsonlFile:
    """Write a zstd-compressed JSONL file.

    Compression is activated by using a .jsonl.zst extension.
    Both reading and writing handle compression transparently.
    """
    out = JsonlFile.new_remote("results.jsonl.zst")
    async with out.writer(compression_level=3) as writer:
        for i in range(100_000):
            await writer.write({"id": i, "compressed": True})
    return out

@env.task
async def write_compressed_dir() -> JsonlDir:
    """Write zstd-compressed shards by specifying the shard extension."""
    out = JsonlDir.new_remote("compressed_dataset/")
    async with out.writer(
        shard_extension=".jsonl.zst",
        max_records_per_shard=50_000,
    ) as writer:
        for i in range(200_000):
            await writer.write({"id": i, "data": f"payload-{i}"})
    return out

async with out.writer(flush_bytes=8 << 20) as writer:
    ...

@env.task
async def read_with_prefetch_tuning(dataset: JsonlDir) -> int:
    """Tune or switch off the background shard prefetch.

    queue_size bounds the read-ahead buffer in records; lower it when records
    are large. prefetch=False reads strictly one shard at a time.
    """
    count = 0
    async for record in dataset.iter_records(prefetch=True, queue_size=1024):
        count += 1

    async for record in dataset.iter_records(prefetch=False):
        count += 1
    return count

@env.task
async def read_with_error_handling(data: JsonlFile) -> int:
    """Read records, skipping any corrupt lines instead of raising."""
    count = 0
    async for record in data.iter_records(on_error="skip"):
        count += 1
    return count

@env.task
async def read_with_custom_handler(data: JsonlFile) -> int:
    """Read records with a custom error handler that collects errors."""
    errors: list[dict] = []

    def on_error(line_number: int, raw_line: bytes, exc: Exception) -> None:
        errors.append({"line": line_number, "error": str(exc)})

    count = 0
    async for record in data.iter_records(on_error=on_error):
        count += 1
    print(f"{count} valid records, {len(errors)} errors")
    return count

arrow_env = flyte.TaskEnvironment(
    name="jsonl-arrow",
    image=flyte.Image.from_debian_base(name="jsonl-arrow").with_pip_packages(
        "flyteplugins-jsonl[arrow]"
    ),
)

@arrow_env.task
async def analyze_with_arrow(dataset: JsonlDir) -> float:
    """Stream records as Arrow RecordBatches for analytics.

    Memory usage is bounded by batch_size — the full dataset is
    never loaded into memory at once.
    """
    import pyarrow as pa

    batches = []
    async for batch in dataset.iter_arrow_batches(batch_size=65_536):
        batches.append(batch)

    table = pa.Table.from_batches(batches)
    mean_value = table.column("value").to_pylist()
    return sum(mean_value) / len(mean_value)