# JSONL

The JSONL plugin adds two typed I/O types for working with [JSON Lines](https://jsonlines.org/) 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`](https://github.com/ijl/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

```bash
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:

```
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"
    ),
)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

## 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:

```
@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
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

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

```
@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
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

## 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:

```
@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
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

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

```
@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
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

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

```
@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
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

## 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:

```
@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
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

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

```
@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
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

### 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:

```
@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
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

### Arrow batches

To hand JSONL data to columnar tooling, stream it as Arrow `RecordBatch`es 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]'`:

```
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)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/files-and-directories/jsonl.py*

## 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](https://www.union.ai/docs/v2/union/api-reference/integrations/jsonl/_index) for the full `JsonlFile` and `JsonlDir` method listings.

---
**Source**: https://github.com/unionai/unionai-docs/blob/main/content/integrations/jsonl/_index.md
**HTML**: https://www.union.ai/docs/v2/union/integrations/jsonl/