JSONL

The JSONL plugin adds two typed I/O types for working with JSON Lines data as task inputs and outputs: flyteplugins.jsonl.JsonlFile for a single JSONL file and flyteplugins.jsonl.JsonlDir for a directory of sharded JSONL files. Both are backed by orjson for fast serialization and stream records one at a time, so you can process datasets that don’t fit in memory.

JsonlFile and JsonlDir extend the built-in flyte.io.File and flyte.io.Dir types, so they inherit remote-storage, upload/download, and caching behavior — they simply add JSONL-aware streaming readers and writers on top. Every read/write method has a synchronous _sync counterpart (writer_sync(), iter_records_sync()) for use in non-async tasks.

When to use this plugin

Passing line-delimited JSON datasets (LLM training/eval sets, event logs, model outputs) between tasks
Streaming records without loading an entire file into memory
Writing large outputs as automatically rotated, sharded directories
Working with compressed JSONL (.jsonl.zst) transparently

Installation

pip install flyteplugins-jsonl

Add the plugin to your task image. Installing it registers JsonlFile and JsonlDir with the Flyte type engine automatically — no explicit registration call is needed:

jsonl.py
                
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"
    ),
)

Working with `JsonlFile`

Create a writable file reference with JsonlFile.new_remote(), then stream records through the writer() context manager without holding the whole dataset in memory:

jsonl.py
                
@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

Reading is equally streaming — iter_records() yields one parsed dict per line:

jsonl.py
                
@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

Working with `JsonlDir`

JsonlDir writes a directory of shard files (part-00000.jsonl, part-00001.jsonl, …) and reads them back transparently in sorted order. Pass max_records_per_shard (or max_bytes_per_shard) to control shard rotation:

jsonl.py
                
@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

Reading iterates across all shards transparently, prefetching the next shard in the background to overlap network I/O with processing:

jsonl.py
                
@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

For bulk processing, iter_batches() yields lists of records at a time; JsonlDir also inherits all flyte.io.Dir capabilities (walk(), list_files(), download()):

jsonl.py
                
@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

Configuration and options

Compression

Give the file a .jsonl.zst (or .jsonl.zstd) extension and records are zstd-compressed transparently on write and decompressed on read. Tune the level via the writer:

jsonl.py
                
@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

For JsonlDir, set shard_extension=".jsonl.zst" on writer(). Mixed compressed and uncompressed shards within a directory are supported on read:

jsonl.py
                
@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

Error handling on read

The record iterators accept an on_error argument — "raise" (default), "skip" to drop malformed lines, or a callable (line_number, raw_line, exception) -> None for custom handling:

jsonl.py
                
@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 batches

To hand JSONL data to columnar tooling, stream it as Arrow RecordBatches with iter_arrow_batches(batch_size=...). Memory usage stays bounded by the batch size. Arrow iteration requires the optional pyarrow dependency — install it with pip install 'flyteplugins-jsonl[arrow]':

jsonl.py
                
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)

Common use cases

LLM dataset pipelines — stream prompt/completion or eval records between preprocessing, generation, and scoring tasks.
Event and log processing — read large line-delimited logs shard by shard without buffering the whole file.
Fan-out writes — produce a JsonlDir of rotated shards from a task that emits millions of records, then consume it downstream.

API reference

See the JSONL API reference for the full JsonlFile and JsonlDir method listings.