Metadata-Version: 2.4
Name: dataio-rs
Version: 0.4.15
Classifier: Programming Language :: Rust
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Dist: numpy>=2.0,<3 ; extra == 'bench'
Requires-Dist: pillow>=10,<12 ; extra == 'bench'
Requires-Dist: torch>=2.6,<2.10 ; extra == 'bench'
Requires-Dist: torchvision>=0.21,<0.25 ; extra == 'bench'
Requires-Dist: opencv-python-headless>=4.10,<5 ; extra == 'bench'
Requires-Dist: safetensors>=0.4.5,<0.7 ; extra == 'bench'
Requires-Dist: dataio-rs[bench] ; extra == 'bench-dali'
Requires-Dist: nvidia-dali-cuda120>=1.44,<1.52 ; extra == 'bench-dali'
Requires-Dist: dataio-rs[bench] ; extra == 'bench-spdl'
Requires-Dist: spdl>=0.0.8,<0.1 ; python_full_version >= '3.10' and extra == 'bench-spdl'
Requires-Dist: dataio-rs[bench] ; extra == 'bench-tensorflow'
Requires-Dist: tensorflow-cpu>=2.18,<2.21 ; extra == 'bench-tensorflow'
Provides-Extra: bench
Provides-Extra: bench-dali
Provides-Extra: bench-spdl
Provides-Extra: bench-tensorflow
License-File: LICENSE
License-File: THIRD_PARTY_NOTICES
Summary: High-performance data processing library for ML workloads.
Home-Page: https://github.com/Mikubill/dataio-rs
License-Expression: Apache-2.0
Requires-Python: >=3.11
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
Project-URL: Package, https://pypi.org/project/dataio-rs/
Project-URL: Repository, https://github.com/Mikubill/dataio-rs

# dataio

[![CI](https://github.com/Mikubill/dataio-rs/actions/workflows/ci.yml/badge.svg)](https://github.com/Mikubill/dataio-rs/actions/workflows/ci.yml)
[![PyPI](https://img.shields.io/pypi/v/dataio-rs.svg)](https://pypi.org/project/dataio-rs/)
[![Python](https://img.shields.io/pypi/pyversions/dataio-rs.svg)](https://pypi.org/project/dataio-rs/)

A Rust-backed data-loading library for local files, object stores, tensors and
columnar data. Build reusable fetch/decode/transform/batch pipelines or read
individual objects through the same IO layer. Outputs include NumPy, PyTorch,
DLPack and Arrow.

Native IO and kernels release the GIL and share worker pools. Python source and
map callbacks run at explicit boundaries. Buffer sharing depends on ownership,
layout and device transfer; retained views keep their backing storage alive.

## Install and run

```bash
pip install dataio-rs
python -m dataio  # report this host's capabilities
```

Python 3.11 or newer is required. There are no required Python dependencies;
install NumPy, PyTorch or PyArrow when using that framework's outputs. Wheels
run without CUDA libraries. CUDA, nvJPEG, cuFile/GDS and libfabric are loaded
when their features are used.

This example runs without input files or an output framework:

```python
import dataio

records = [
    {"key": "a", "payload": b"red"},
    {"key": "b", "payload": b"sky"},
]
pipe = dataio.pipeline(
    dataio.slot("payload").decode_bytes(),
    dataio.batch(2),
)
with pipe.bind(records).run("dlpack") as loader:
    for batch in loader:
        print(batch.keys, batch["payload"].shape)  # ['a', 'b'] [2, 3]
```

Use `.run("numpy")` or `.run("torch")` to convert tensor slots to that framework.
The following examples show common tasks; file paths, datasets and callbacks
must be supplied by your application. For complete scripts, see [Examples](examples/README.md).

## Build a pipeline

A pipeline contains optional record maps, one or more slot chains, exactly one
batching stage (`batch`, `collate` or `group_by`), then optional batch transforms
and Python maps. Creating the plan does no IO; `.bind(records).run(...)` creates
the loader. Close loaders with a context manager, especially when stopping early.

```python
pipe = dataio.pipeline(
    dataio.map(lambda row: {"image": row["path"], "label": row["class_id"]}),
    dataio.slot("image")
        .decode_image(mode="rgb")
        .resize_short(256)
        .center_crop(h=224, w=224)
        .normalize("imagenet"),
    dataio.slot("label").passthrough(),
    dataio.batch(64, drop_last=True, prefetch=4),
)
with pipe.bind(records).run("torch", concurrency=64) as loader:
    for batch in loader:
        train_step(batch["image"], batch["label"])
```

Slot methods are shorthand for the factories in these namespaces:

| Namespace | Operations |
| --- | --- |
| `dataio.decoders` | Image, audio, JSON, NPY, safetensors, text, bytes and PT decoding |
| `dataio.transforms` | Resize, crop, normalize, augment and conditional transforms |
| `dataio.batch_transforms` | MixUp, CutMix, padding, shuffle and normalization |
| `dataio.cond` | Field and shape predicates for slot activation and transforms |

For example, `.decode_audio(sample_rate=16000, resampling="sinc")` enables
anti-aliased resampling; the default `"linear"` costs less computation.
[Conditional slots](examples/conditional_slots.py) can gate optional inputs
before decode; [conditional transforms](examples/conditional_branch.py) branch
on decoded shapes.

### Choose and compose sources

Bind a reusable iterable or a zero-argument record factory for repeated passes.
`from_iterable(...)` is single-pass, including when given a factory. For a finite
third-party dataset with integer row access, choose
`dataio.records(dataset, indexed=True)`. It requires `__len__` and integer
`__getitem__` and bypasses the dataset's iterator.

With `batch(N)`, indexed datasets can provide `__getitems__(indices)` to fetch
rows together. Return a list or tuple of exactly the requested rows, in order,
including duplicate indices. Fetch errors propagate without scalar retries.
See the [bulk-read protocol](docs/arch.md#indexed-and-bulk-reads) and
[runnable example](examples/callable_source.py).

Source transforms run before native decoding:

```python
def expand(record):
    for index, payload in enumerate(record["parts"]):
        yield {"key": f"{record['key']}:{index}", "payload": payload}

records = dataio.records(group_a).concat(group_b).flat_map_records(expand)
```

`map_records(fn)` replaces one record; `flat_map_records(fn)` emits zero or more
complete records; `concat(...)` appends sources without opening later ones early.
Expansion has unknown length. Use bounded shuffle for expanded or streaming
sources:

```python
records = dataio.webdataset("/data/shard-{0000..0099}.tar").shuffle(
    seed=42, buffer_size=4096,
)
```

The buffer bounds pending records, not bytes, and mixes nearby records. Omitting
`buffer_size` requests a full permutation, materializing non-indexed sources;
single-pass sources require a buffer. A fixed seed repeats the order only for an
unchanged, repeatable source. See [source contracts](docs/arch.md#record-sources).

### Adapt Python collation and sampler batches

Python maps before batching see input records. Maps after batching see the
decoded `Batch` and may return nested mappings, tuples or application objects.
They run on the consumer thread during synchronous iteration. Native batch
transforms must precede Python batch maps.

Pass Python objects through a slot to reuse an existing collator:

```python
pipe = dataio.pipeline(
    dataio.slot("sample").passthrough(),
    dataio.batch(32),
    dataio.map(lambda batch: collate_fn(batch["sample"])),
)
```

Each input record must contain `sample`. For native tensor slots, a batch map
can instead assemble nested results from `batch["slot_name"]`.

`collate()` preserves batch boundaries supplied by the source; it does not take
a Python collate function. `group_by(key)` emits one batch for each contiguous
run of a key. To adapt a sampler yielding lists of rows:

```python
records = dataio.from_sampler(
    batch_sampler,
    lambda row: {"key": row["id"], "payload": row["path"], "meta": row},
)
pipe = dataio.pipeline(
    dataio.slot("payload").decode_numpy(),
    dataio.slot("meta").passthrough(),
    dataio.group_by("batch_id", prefetch=4),
)
```

`from_sampler` adds `batch_id` to shallow copies of record mappings. Omit its
second argument when the sampler already yields valid records. The
[batch-boundary example](examples/stream_shapes.py) compares both approaches.

### Consume asynchronously

```python
with pipe.bind(records).run("numpy") as loader:
    async with loader.aiter() as batches:
        async for batch in batches:
            await consume(batch)
```

The async iterator reuses native prefetch and closes on early context exit.
Source and batch callbacks run on its worker thread. Cancellation waits for an
already-running synchronous callback to return. Do not mix sync and async
iteration on the same loader. Checkpoints track batches delivered to the caller;
see [iteration and cancellation](docs/arch.md#iteration-and-cancellation).

## Configure execution

| Setting | Scope |
| --- | --- |
| `run(..., concurrency=N)` | Per-shard stage admission; total activity also depends on prefetch |
| `batch(..., prefetch=K)` | Number of batches in flight |
| `run(..., max_prefetch_bytes=B)` | Logical bytes of pending native preallocated outputs |
| `DATAIO_CPU_WORKERS` / `DATAIO_IO_WORKERS` | Process-wide CPU / async IO worker pools |
| `DATAIO_HTTP_MAX_INFLIGHT_READS` | Process-wide GET/HEAD limit; `0` means unlimited |
| `DATAIO_OBJECT_SHARDS` | HTTP and object-store client shard fanout |

Start with `concurrency="auto"`; measure higher values for high-latency sources.
Submit order is the default. `order="completion"` returns completed batches sooner
but cannot resume mid-epoch. The output byte limit requires a fixed native layout
and direct writer; unsupported outputs and oversized batches fail before allocation.
It excludes scratch space, caches and consumer-held outputs, so it is not an RSS
limit. [Architecture](docs/arch.md#plans-and-execution) describes the scheduling
and ownership rules.

### Handle errors

`run(..., on_failure="strict")` raises on a row error. `"skip"` drops failed rows
and raises when no rows survive; an integer requires that many surviving rows.
`batch.errors` retains per-row errors. Pass `log_errors=True` to use the
`dataio.rows` logger, or provide a logging callable.

Source-iteration and Python batch-map errors propagate directly. A failed batch
map does not roll back consumed progress or its side effects. Padding and batch
normalization preserve row metadata; MixUp and CutMix invalidate it because they
combine rows. See [callback contracts](docs/arch.md#composition-and-callbacks).

### Shard and resume

```python
loader = pipe.bind(records).run(
    "torch", shard={"id": rank, "of": world_size, "pad": True}, seed=42,
)
# Save loader.state_dict() and restore with loader.load_state_dict(state).
```

`pad=True` equalizes rank batch counts. Resume replays a partial batch to align
native transform RNG. By default it re-enumerates the consumed source prefix;
source data and callbacks must replay consistently. A checkpoint identifies the
plan and tracked source configuration, not hidden callback state or data contents.

Deterministic indexed datasets can avoid fetching consumed rows:

```python
records = dataio.records(dataset, indexed=True, resume="seek")
```

Seek requires lookups determined by index. Source maps, source state hooks and
composed iterables retain replay. See the [resume contract](docs/arch.md#replay-and-indexed-resume)
and [checkpoint compatibility](docs/arch.md#checkpoint-identity).

## Read and write objects

Local paths, `s3://`, `r2://`, `gs://`, `az://`, `hf://` and HTTP(S) URIs share
these operations where the backend supports them:

```python
chunk = dataio.read("s3://bucket/file.bin", offset=0, size=1 << 20)
dataio.write("s3://bucket/out.bin", chunk)
dataio.copy("s3://bucket/out.bin", "/data/out.bin")
dataio.head("s3://bucket/out.bin")
dataio.exists("s3://bucket/out.bin")
dataio.ls("s3://bucket/prefix/")
dataio.glob("s3://bucket/**/*.safetensors")
```

`exists` returns `False` only for a missing resource. Permission failures,
timeouts and other IO errors propagate, so an inaccessible object is not
mistaken for one that needs creating.

Fill an existing writable buffer with `dataio.fetch_into(uri, buffer)`. For
composable reads and incremental writes:

```python
src = dataio.reader("s3://bucket/shard.tar.gz")
payload = src.gunzip().tar_entry("0001.jpg").read_all()
with dataio.writer("s3://bucket/checkpoint.bin") as writer:
    writer.write_chunks(huge_generator())
```

`reader.read_buffer()` returns a read-only buffer retaining native storage.
`dataio.stream(uri)` reads the response body on demand; byte, line and SSE
iterators share a cursor. Close unused streams. Streaming does not retry a body
after exposing it to the caller; only `read_all()` materializes all remaining bytes.

Configure credentials and transport before use:

```python
dataio.configure_credentials("s3", method="default")
dataio.configure_endpoint("r2", "https://<account>.r2.cloudflarestorage.com")
dataio.configure_transport(
    "https", timeout_secs=30, readall_chunk_mib=32, readall_concurrency=8,
)
```

`method="default"` uses the platform credential chain. Range settings apply to
each read; splitting requires a strong validator. Transport configuration
replaces the scheme's explicit options, so pass options together. Proxy, TLS,
cache, timeout and shared-read-limit behavior is documented in
[network configuration](docs/arch.md#network-configuration).
Repeating identical configuration preserves cached clients and active readers.
To refresh environment or file-based credentials, call `clear_credentials(scheme)`
and then reapply any credential and endpoint overrides.

## Read structured formats

```python
array = dataio.read_numpy("/data/x.npy")
image = dataio.read_image("s3://bucket/image.png", mode="rgb")
columns = dataio.read_parquet("/data/shard.parquet", columns=["key", "height"])
rows = dataio.scan_parquet("/data/shard.parquet").filter(
    lambda row: row["height"] > 512,
).collect()
for example in dataio.iter_tfrecord("/data/train.tfrecord"):
    label = example["label"]
```

Parquet and Feather/Arrow IPC scans preserve Arrow types without creating Python objects
for individual cells:

```python
scan = dataio.scan_feather("/data/table.arrow").select(["key", "value"])
stream = scan.collect_arrow_stream()  # Arrow C stream protocol
reader = scan.collect_arrow_reader()  # optional PyArrow batch reader
table = scan.collect_arrow()          # optional PyArrow table
```

For files with numeric columns, `iter_numpy_batches()` returns arrays directly
and needs only NumPy. Given a Parquet file with a numeric `features` column:

```python
scan = dataio.scan_parquet("/data/features.parquet").select(["features"])
for batch in scan.iter_numpy_batches():
    features = batch["features"]  # owned, writable, C-contiguous NumPy array
```

The same method works on `scan_feather()`. Struct columns become nested
dictionaries of arrays; fixed-size lists add array dimensions. Variable-length
numeric lists become `RaggedArray(values, offsets)`: row `i` is
`values[offsets[i]:offsets[i + 1]]`. Offsets are zero-based `int64` arrays, and
empty rows are preserved without padding. Selected null values and unsupported
types raise errors.
`dataio.iter_numpy_batches(source)` also accepts CPU Arrow table streams from
other libraries through `__arrow_c_stream__`. For example, with optional DuckDB
installed:

```python
import duckdb

query = duckdb.sql("SELECT range AS id FROM range(4)")
for batch in dataio.iter_numpy_batches(query):
    print(batch["id"])  # [0 1 2 3]
```

Pass a zero-argument factory to create a stream when iteration starts, such as
`dataio.iter_numpy_batches(scan.collect_arrow_stream)`. Reusing an existing
reader may yield no more data or fail; fresh epochs need a fresh source.
`copy=False` returns read-only arrays retaining their buffers; normalizing ragged
offsets can still require allocation. See the
[array ownership contract](docs/arch.md#columnar-numpy-batches) for supported types
and retention limits. Arrow and NumPy exports support projection and row slicing;
Parquet also supports native expressions. Python callable filters use row
collection instead. NPY decoding supports complex values, foreign byte order and
Fortran layout.

HDF5 supports integer indices, positive-step slices and one ellipsis:

```python
window = dataio.read_hdf5(
    "/data/arrays.h5", key="images", selection=(slice(10, 20), ...),
)
with dataio.open_hdf5("/data/arrays.h5") as source:
    view = source.read(key="labels", copy=False)
```

The default returns writable copies. `copy=False` gives read-only output that
survives session close, though decoding or conversion can allocate. Local chunked
selections decode intersecting chunks; remote HDF5 still downloads the file.

### Read or fill safetensors weights

Single-file and sharded checkpoints use the same API:

```python
dataio.write_safetensors(
    "checkpoints/model/", state_dict, max_shard_size="5GB", metadata={"format": "pt"},
)
subset = dataio.read_safetensors(
    "checkpoints/model/", keys=["layer.weight", "layer.bias"], framework="torch",
)
with dataio.open_safetensors("weights.safetensors") as source:
    tensor = source.read(key="layer.weight", framework="torch")
    source.read_into(destinations, chunk_bytes=16 << 20)
```

`destinations` maps tensor names to existing NumPy arrays or Torch tensors.
`read_into` validates shape, dtype, layout and overlap before writing, then fills
CPU or CUDA targets in bounded chunks. It avoids another complete weight copy
and does not use GDS. IO failure can leave targets partially written; inference
Torch targets require an inference-mode context.

NumPy views are read-only; local Torch mappings use copy-on-write. The
[ownership reference](docs/arch.md#buffer-ownership-and-dlpack) explains copies,
CUDA staging and buffer lifetime. Complex64 outputs are supported. On local
GDS-ready storage, pass `framework="torch", device="cuda:0"` to
`read_safetensors` for direct GPU reads; unavailable GDS falls back to CPU with
a warning. `gds_available()` and `gds_info()` report the host's GDS support.

### Convert distributed checkpoints

```python
dataio.scan_dcp("s3://bucket/checkpoint", keys="model.", bf16=True).rename(
    lambda name: name.removeprefix("model.")
).write("weights.safetensors", strategy="stream")
```

Selection uses original names; rename collisions fail before payload reads.
Streaming batches whole tensors, so one tensor can exceed the batch target.
Dense row-block shards and transposed storage are supported; storage transforms
and other sharding layouts are rejected. DCP metadata uses Python pickle: load
only trusted checkpoints. NumPy collection rejects BF16 instead of changing dtype.

## Transfer with RDMA

`dataio.rdma` provides one-sided reads and writes through libfabric. It loads the
library on first use; `available()` can also succeed with TCP software providers,
so inspect `info()` or `diagnose()` to establish physical RDMA capability.

```python
import dataio.rdma as rdma

with rdma.serve(bind="0.0.0.0:0", buffer_size=1 << 20) as server:
    with rdma.connect(server.local_addr) as client:
        rdma.write(client, offset=0, data=b"hello rdma")
        assert rdma.read(client, offset=0, size=10) == b"hello rdma"
```

`serve_file` exposes file ranges; `read_into` fills caller-owned CPU or CUDA
buffers. GPUDirect RDMA requires a compatible HMEM provider and libfabric API
1.9 or newer. See [RDMA transport](docs/arch.md#rdma-transport) for buffer,
provider and diagnostic details.

## API index

| Task | Entry points |
| --- | --- |
| Build a pipeline | `pipeline`, `slot`, `slots`, `chain`, `batch`, `collate`, `group_by` |
| Extend processing | `map`, `custom_op`, `register_output_adapter` |
| Bind records | `records`, `files`, `urls`, `manifest`, `from_callable`, `from_iterable`, `from_iterator`, `from_sampler`, `webdataset`, `archive` |
| Work with bytes | `read`, `fetch_into`, `write`, `copy`, `sync`, `head`, `stream`, `ls`, `glob`, `exists`, `delete` |
| Compose IO | `reader`, `writer`, `scan` |
| Work with tensors | `read_safetensors`, `open_safetensors`, `write_safetensors`, `scan_dcp`, `read_dcp`, `scan_pt`, `read_pt` |
| Work with arrays/tables | `read_*` / `scan_*` for `numpy`, `npz`, `parquet`, `feather`, `tfrecord`, `hdf5`; `read_image`, `iter_tfrecord`, `open_hdf5` |
| Configure IO | `configure_credentials`, `configure_endpoint`, `configure_transport` |
| Use GDS | `gds_open`, `gds_register_buffer`, `gds_read_into_tensor` |

The [documentation index](docs/README.md) links to execution contracts,
benchmarks and planned capabilities. Use the [benchmark guide](docs/benchmark.md)
to compare the output and workload your application actually uses.

## Develop

The toolchain is pinned in `rust-toolchain.toml`. Build the editable extension,
then run the relevant tests:

```bash
uv run --no-sync --with maturin maturin develop --profile staging
cargo test
uv run --no-sync python -m pytest python/tests -q
```

Source builds omit RDMA by default. Install libfabric development headers and add
`--features rdma` to enable it; published wheels include the feature. Native IO,
object and format crates can build without Python. Build a wheel with:

```bash
uv run --no-sync --with maturin maturin build --profile staging --out dist
```

See [Benchmarking](docs/benchmark.md) for workload validation and comparisons.

## License

[Apache-2.0](LICENSE). Bundled dependency licenses are in
[THIRD_PARTY_NOTICES](THIRD_PARTY_NOTICES).

