# zephyr - Marin light dataset library

We have a _lot_ (like 100) redundant data processing patterns in the Marin code
base.  Most of them involve ad-hoc use of Ray + HF datasets + manual processing.
Since most of these are working on smaller datasets, we don't have a lot of
concern about completion or stalled jobs. The lack of consistency though, means
we need to know 10 different ways to run jobs, understand how errors are
handled, and deal with a wide variety of "how do I load a jsonl" file. Given our
desire to abstract away from Ray for our cluster management, we also want to
consolidate our usage of Ray as much as possible.

To that end, we've identified a set of common patterns in our code we can
consolidate using some standard data processing patterns. We're calling this
"zephyr".

## Outline

To use Zephyr, like Beam/Flume/Ray-Data etc, you define a pipeline of operations against a Dataset:

```python
dataset = (Dataset
  .from_files("gs://my-cool-data/*.jsonl")
  .flat_map(load_jsonl) # load each file, and yield individual records in the file
  .map(convert_to_dolma_format)
  .write_jsonl("gs://ml-cool/output/processed-{shard%05d}-of-{total%05d}")
)

for filename in backend.execute(dataset):
  print(filename)
```

You focus on the individual operations, and the dataset & backend handle the
scaling and load-management for you. To process a dataset, you need a `Backend`:
we provide simple multi-processing and synchronous backends for small datasets,
as well as a Ray backend which runs tasks on a cluster. Pipelines are invoked
directly as Python functions:

```python
# experiments/my_processor.py
def my_entry_point(conf: Config, backend: zephyr.Backend):
  dataset = ...
  backend.execute(dataset)
```

## Datasets

Users primarily interact with zephyr via the Dataset class. It exposes a
standard set of constructors and transformations for construction,
transformation, and serialization.

`from_{files,list,iterator}`

Construct a dataset from a list or glob.

`flat_map(fn)`

Given a function from `item` -> `list[new_item]`, return a new `Dataset` which
contains `new_item`. A common use case is for loading files: given a dataset of
filenames, `ds.flat_map(load_jsonl)` returns a dataset which contains the
individual JSON lines.

`map_shard(fn)`

The converse of flat_map: the function receives an iterator of all items in a
shard and returns an iterator of results. This is the recommended pattern for
stateful shard processing without requiring callable classes.

Use when you need to maintain state across all items in a shard, such as
deduplication, reservoir sampling, sliding windows, or loading expensive
resources once per shard.

```python
def deduplicate_shard(items: Iterator):
    seen = set()
    for item in items:
        key = item["id"]
        if key not in seen:
            seen.add(key)
            yield item

ds = (Dataset
    .from_files("gs://data/*.jsonl")
    .flat_map(load_jsonl)
    .map_shard(deduplicate_shard)  # Dedup within each shard
    .write_jsonl("gs://output/deduped-{shard:05d}.jsonl.gz")
)
```

`map(fn)`

Map the given function over the dataset, transforming the elements.

`filter(fn)`

Return a new dataset containing only elements where `fn(item)` returns True-ish.

`take_per_shard(n)`

Take the first n items from each shard.

```python
ds = (Dataset
  .from_files("gs://data/*.jsonl")
  .flat_map(load_jsonl)
  .take(1000)
  .map(process)
)
```

`window(size)`

Given a dataset of `item`, return a dataset of `list[item]` where each list
contains up to `size` elements. If the dataset is sharded, windowing is
independently performed on each shard.

`window_by(folder_fn, initial_state=None)`

Window elements using a custom folder function that controls window boundaries.
The folder function receives `(state, item)` and returns `(should_continue, new_state)`.
When `should_continue` is False, the current window closes and a new window
starts with the item that triggered the close.

This enables custom windowing logic such as grouping by total size, time windows,
or any other stateful partitioning strategy.

```python
# Window files by total size < 10GB
ds = (Dataset
  .from_list([{"path": "f1.txt", "size": 5_000_000_000}, ...])
  .window_by(
    folder_fn=lambda total_size, item: (
      total_size + item["size"] < 10_000_000_000,
      total_size + item["size"]
    ),
    initial_state=0
  )
)
```

`batch(size)` is an alias for `window(size)` for backwards compatibility.

`write_{jsonl|parquet}(file_pattern)`

Serialize the dataset to disk at the given location. The output expects a
pattern for serializing sharded datasets, such as
"gs://path-to-data/corpus-{{shard:05d}}-of-{{total:05d}}.jsonl.gz". The `shard`
and `total` formatters are automatically replaced by the backend when
serializing. If there is only one shard being processed, the shard & total
patterns may be omitted.

`group_by(key, reducer, num_output_shards=None)`

Group items by a key function and apply a reducer to each group. The operation
is implemented as a two-phase shuffle:

```python
ds = (Dataset
  .from_files("gs://data/*.jsonl")
  .flat_map(load_jsonl)
  .group_by(
    key=lambda x: x["user_id"],
    reducer=lambda key, items: {"user": key, "count": sum(1 for _ in items)}
  )
)
```

`deduplicate(key, num_output_shards=None)`

Remove duplicate items based on a key function.

```python
ds = (Dataset
  .from_files("gs://data/*.jsonl")
  .flat_map(load_jsonl)
  .deduplicate(key=lambda x: x["id"])
)
```

`reduce(local_reducer, global_reducer=None)`

Reduce the entire dataset to a single value. The local_reducer runs on each
shard before storing results which are then reduced by the global reducer.

```python
# Count total items
total = (Dataset
  .from_list(range(1000))
  .reduce(local_reducer=lambda items: sum(1 for _ in items))
)

# Custom aggregation
stats = (Dataset
  .from_files("gs://data/*.jsonl")
  .flat_map(load_jsonl)
  .reduce(
    local_reducer=lambda items: {"count": sum(1 for _ in items), "sum": sum(x["value"] for x in items)},
    global_reducer=lambda shards: {
      "total_count": sum(s["count"] for s in shards),
      "total_sum": sum(s["sum"] for s in shards)
    }
  )
)
```

`sorted_merge_join(right, left_key, right_key, combiner=None, how="inner")`

Perform a streaming merge join between two datasets that are already sorted and
co-partitioned. Preconditions:

- Both datasets must have the same number of shards
- Corresponding shards (left[i], right[i]) must contain the same key ranges
- Items within each shard must be sorted by their join key

Only supports inner and left joins (no right or outer joins).

```python
# Both datasets grouped by the same key
docs = (Dataset
  .from_files("gs://docs/*.jsonl")
  .flat_map(load_jsonl)
  .group_by(key=lambda x: x["id"], reducer=keep_first, num_output_shards=100)
)

attrs = (Dataset
  .from_files("gs://attrs/*.jsonl")
  .flat_map(load_jsonl)
  .group_by(key=lambda x: x["id"], reducer=keep_first, num_output_shards=100)
)

joined = docs.sorted_merge_join(
  attrs,
  left_key=lambda x: x["id"],
  right_key=lambda x: x["id"],
  how="inner"  # or "left"
)
```

The default combiner merges dictionaries: `{**left, **right}`. For custom
combining or handling nulls in left joins, provide a custom combiner function.

## Backends

A `backend` takes a dataset and runs it, yielding the results from the final
stage. We provide local backends for testing and small datasets, as well as a
Ray backend which runs operations in parallel across a Ray cluster.

You can create a backend manually and specify the requirements for running
individual tasks:

```
backend = RayBackend(memory=1_000_000_000, cpu=1)
backend.execute(ds)
```

### Optimizations

By default, the Ray backend creates a separate _shard_ for each input item in a
cluster - operations like `map` are run in parallel across shards, but
sequentially within a shard.  Naive execution of a dataset would result in each
individual stage being "materialized" - this can result in memory blow-ups and
slow execution. For example, in:

```
Dataset.from_files().flat_map(load_jsonl).map(x).filter(y).map(z).write_jsonl()
```

We want to avoid computing the intermediate representation of the dataset if
possible. Fortunately for us, we only rely on trivial transformations, thus we
can easily fold these operations into a single combined map against each file
shard:

```
Dataset.from_files().map_shard(load_jsonl | x | y | z)
```

## Alternatives Considered

Obvious alternatives to our own library are Beam and Ray-data. Both of these
support basic dataset transformations and would be reasonable fits for some of
our use cases. We elected to make our own abstraction for a few reasons:

* We're unclear about our long-term usage of Ray, and didn't want to embed it everywhere in the code base at this point.
* The only fully implemented Beam backend is on GCP, and the execution model is very disjoint from e.g., our Ray cluster. This makes it suitable for big standalone jobs, but we want a dataset pattern that works for smaller jobs and is more easily tracked as well.

While we decided to make our own abstraction, there is no reason we can't add
Beam or Ray-Data as an additional _backend_ to our system. The cost of defining
our own dataset layer is quite trivial and thus we feel it's worth it to buy us
additional flexibility.


## Implementation Notes

### Repartitioning

Zephyr provides `reshard(num_shards)` to change the parallelism of a dataset
without changing the underlying order. This is useful for increasing or
decreasing the amount of parallel work. For example, a pipeline like
`Dataset.from_files([".../corpus.jsonl"]).flat_map(load_jsonl)` will only run on
a single worker by default. Using `.reshard(num_shards=20)` forces the runtime
to redistribute the data across 20 shards, allowing subsequent operations to run
in parallel.

The implementation is best-effort and redistributes chunks across shards in
round-robin fashion without splitting individual chunks.

### Resuming from Failures

For a pure map pipeline, you can leverage Zephyr's `skip_existing` functionality
to skip working on shards that have already been processed. As an optimization,
any operations _before_ the output file are also skipped.

```python
ds = (Dataset
    .from_files("gs://input/*.jsonl")
    .flat_map(load_jsonl)
    .map(expensive_transform)
    .write_jsonl("gs://output/processed-{shard:05d}-of-{total:05d}.jsonl", skip_existing=True)
)
```

For example, if you have 5 input files and outputs for shards 1 and 3 already
exist, the pipeline will:
- Process shards 0, 2, and 4 (creating new outputs)
- Skip processing shards 1 and 3 (preserving empty shards internally)
- Skip writing to existing output files for shards 1 and 3
