from mixtrain import Dataset

Constructor

Dataset(name: str, *, version: int | None = None)

Creates a reference to an existing dataset on the platform. This is a lazy reference, no API call is made until you access data or schema.

dataset = Dataset("training-data")
baseline = Dataset("training-data", version=3)

Properties

Versions

Mixtrain datasets are versioned on each update.

latest = Dataset("training-data")
v3 = Dataset("training-data", version=3)

list_versions()

dataset.list_versions() -> list[DatasetVersion]

List dataset versions, newest first. Each entry includes version, snapshot_id, operation, added_records, total_records, and timestamp_ms when available.

for version in Dataset("photos").list_versions():
    print(version.version, version.operation)

added_since()

dataset.added_since(
    from_version: int | None = None,
    to_version: int | None = None,
) -> Dataset

added_since returns the rows added or changed between two versions — present at to_version and new or modified since from_version. from_version is exclusive and to_version is inclusive; from_version=None starts from the earliest version and to_version=None ends at the latest. Unchanged and deleted rows do not appear.

Dataset-triggered routines use added_since/deleted_since to process just the new rows, and have from_version and version properties to identify the version range they cover.

deleted_since()

dataset.deleted_since(
    from_version: int | None = None,
    to_version: int | None = None,
) -> Dataset

deleted_since returns the rows present at from_version and gone at to_version. from_version is exclusive and to_version is inclusive; from_version=None starts from the earliest version and to_version=None ends at the latest.

as_read_only()

dataset.as_read_only() -> Dataset

Return a read-only version of this dataset. Calling in-place mutations (append()/update()/delete()/drop()/set_column_types()/set_description()) would raise an error. Reads and non-mutating transformations will continue to work.

ro = Dataset("photos").as_read_only()
ro.sql("SELECT * FROM photos").to_pandas()      # ok
ro.append(rows)     # raises: read-only dataset

Iteration

Row iteration

for row in dataset:
    print(row)  # {"col1": value, "col2": value}

Streams rows without loading the full dataset into memory. dataset.rows(load_files=...) is the explicit form; file-backed columns arrive as Image/Video/File handles by default (load_files="auto"), as raw records with load_files=False, or as loaded content with load_files="bytes" / "path" / a per-column dict. See load_files().

batches()

dataset.batches(
    size: int = None,
    format: str = "dict",      # "dict" | "arrow"
    load_files: bool | str | dict = "auto",
) -> Iterator[dict[str, list]] | Iterator[pyarrow.RecordBatch]

Streams batches. size=None keeps the plan's natural chunk sizes. format="dict" yields columnar dicts (file columns as handles under "auto"); format="arrow" yields RecordBatches (file columns stay native record structs under "auto"). The explicit load_files values (False, "bytes", "path", per-column dict) mean the same thing in every format.

for batch in dataset.batches(64):
    print(batch)  # {"col1": [v1, v2, ...], "col2": [v1, v2, ...]}

for batch in dataset.batches(64, load_files="bytes"):
    send(batch["image"])  # list[bytes], fetched in parallel

Creating Datasets

From files

Dataset.from_file(file_path: str) -> Dataset

Load a local data file into an in-memory dataset. Use .save() to persist it to Mixtrain.

Supported formats: .csv, .parquet, .jsonl (all stream in batches)

dataset = Dataset.from_file("data.parquet")
dataset.save("training-data", description="Training dataset")

From in-memory data

# From Python dict
ds = Dataset.from_dict({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})

# From pandas DataFrame
ds = Dataset.from_pandas(df)

# From Arrow table
ds = Dataset.from_arrow(table)

# From a re-iterable batch source
ds = Dataset.from_batches(batch_factory, schema=schema)

# From HuggingFace datasets (streams by default)
ds = Dataset.from_huggingface("imdb", split="train")

# From PyTorch dataset (schema= is required)
ds = Dataset.from_torch(
    torch_dataset,
    schema=pa.schema([("data", pa.float32()), ("label", pa.int64())]),
)

These create in-memory or streaming datasets. Use .save() to persist to the platform.

Empty datasets

Dataset.empty(columns: list[str] | dict[str, type]) -> Dataset

Create a zero-row dataset to fill later with append(). You can pass a list of column names or a dict of column names to types. A list of names defaults every column to str. When passing a dictionary, the values can be Python types (str, int, float, bool, bytes), Mixtrain types (Image, Embedding, ...), type names ("image"), or PyArrow types.

from mixtrain import Dataset, Image

# All columns default to str
ds = Dataset.empty(["prompt", "photo", "score"]).save("my-dataset")

# Or specify types
ds = Dataset.empty({"prompt": str, "photo": Image, "score": float}).save("my-dataset")

If column types are not specified, the dataset will create str default columns that adopt the type of the first append(). Explicitly typed columns keep their type, and appending that doesn't match will fail with an error.

Mutating Datasets

append()

dataset.append(
    source: Dataset | Table | DataFrame | dict | list[dict],
    copy_files: bool | str = "auto",
    stats: bool = True,
    workers: int = 8,
) -> Dataset

Append rows from source to an existing platform dataset. source may be another Dataset, a PyArrow Table, a pandas DataFrame, a dict of columns, or a list of row dicts. Dict inputs are column-oriented, so each value must be a sequence. copy_files, workers, and stats behave the same as in save().

# Append rows to an existing dataset
ds = Dataset("my-dataset")

# Append multiple rows
ds.append([{"img": "fox.jpg", "label": 2}, {"img": "owl.jpg", "label": 1}])

# Append 2 rows of column-oriented data
ds.append({"img": ["owl.jpg", "bear.jpg"], "label": [1, 3]})

# The rows can come from another Dataset
new_data = Dataset("new-data")
ds.append(new_data)

update()

dataset.update(predicate: str | None = None, /, **column_values) -> Dataset

Set column values on the rows matching predicate. predicate uses the same expression language as filter() to choose the rows to update. Each keyword argument is a column assignment: column=value writes that constant onto every matched row. Predicate=None updates all rows. Columns that don't exist yet are added. At least one column=value is required.

Dataset("photos").update("score < 0", flag=True)   # update low-scoring rows with column flag set to True
Dataset("photos").update(verified=True, label=1)   # update all rows with column verified set to True and label set to 1

delete()

dataset.delete(predicate: str) -> Dataset

Delete the rows matching predicate. predicate uses the same expression language as filter(). The predicate is required — a bare delete() is refused so it can't silently wipe the dataset. Use drop() to remove the whole dataset.

Dataset("photos").delete("score < 0")   # remove low-scoring rows

save()

dataset.save(
    name: str,
    description: str = None,
    column_types: dict | str | None = "auto",
    copy_files: bool | str = "auto",
    stats: bool = True,
    workers: int = 8,
) -> Dataset

Save dataset to the platform. Files referenced by file-backed columns are handled automatically: local files are uploaded so they can be viewed and used from Mixtrain, while remote URLs are kept as references by default.

# Auto-detects column types and uploads local files (default)
ds = Dataset.from_dict({"img": ["cat.jpg", "dog.png"], "label": [0, 1]})
ds.save("my-dataset")  # img detected as Image, files uploaded

# Explicit types for exactly these columns (no inference of others)
from mixtrain import Image
ds.save("my-dataset", column_types={"photo": Image})

# Pin some columns and let save() infer the rest: chain with_column_types()
ds.with_column_types({"photo": Image}).save("my-dataset")

# Copy remote URLs into workspace storage too
ds.save("my-dataset", copy_files=True)

drop()

dataset.drop() -> dict

Remove the entire dataset from the platform. This is a destructive operation and cannot be undone.

Export Methods

collect()

dataset.collect(max_bytes: int = None, warn: bool = True) -> pyarrow.Table

Execute the plan and materialize the result as an Arrow table. Warns on large results; pass max_bytes to fail instead of materializing more than a budget.

to_arrow() / to_pandas()

dataset.to_arrow() -> pyarrow.Table
dataset.to_pandas() -> pandas.DataFrame

Materialize as an Arrow table or pandas DataFrame (equivalent to collect()).

to_tensors()

dataset.to_tensors(missing: str = "nan") -> dict[str, Tensor | list]

Convert to dict of PyTorch tensors. Uses zero-copy for numeric columns. missing controls null handling — see Missing values.

to_torch()

dataset.to_torch(
    batch_size: int = None,
    num_workers: int = 0,
    rank: int = 0,
    world_size: int = 1,
    drop_last: bool = False,
    load_files: bool | str | dict = "auto",
    prefetch_batches: int = 2,
    workers: int = 32,
    max_inflight_bytes: int = 256_000_000,
    fetch_timeout: float = 30.0,
    missing: str = "nan",
    **dataloader_kwargs,
) -> DataLoader

Get a PyTorch DataLoader. Data is sharded automatically across DataLoader workers and distributed ranks (rank/world_size) so each worker reads a disjoint partition and fetches only its own media files. Workers share one on-disk blob cache (MIXTRAIN_CACHE_DIR), so files fetched in one epoch aren't re-downloaded in the next. File-backed columns are loaded automatically (load_files= accepts the same values as load_files(); False disables); batches are pipelined with background prefetch. workers controls how many files each DataLoader worker fetches in parallel and max_inflight_bytes caps the bytes held in flight while fetching. fetch_timeout sets the per-request timeout, and MIXTRAIN_CACHE_DIR controls the local blob cache location. missing controls null handling for numeric/tensor columns — see Missing values. Extra kwargs (e.g. collate_fn, pin_memory) pass through to torch.utils.data.DataLoader. Shuffle with .shuffle() before calling to_torch() and use loader.set_epoch(epoch) to reshuffle windowed shuffles per epoch.

# Batched - yields dicts of tensors (numeric columns are zero-copy)
loader = dataset.to_torch(batch_size=32, num_workers=4)
for batch in loader:
    print(batch["features"].shape)  # torch.Size([32, ...])

# Media: images arrive as bytes — convert in collate_fn (runs in workers)
import io
from PIL import Image as PILImage

def collate(batch):
    batch["image"] = [PILImage.open(io.BytesIO(b)) for b in batch["image"]]
    return batch

loader = dataset.to_torch(batch_size=64, num_workers=4, collate_fn=collate)

See Best practices for checkpoint/resume and null handling.

to_huggingface()

dataset.to_huggingface() -> datasets.Dataset

Convert to HuggingFace Dataset.

Transformations

All transformations are lazy and return a new Dataset (immutable).

shuffle()

dataset.shuffle(seed: int = None, global_: bool = True, window: int = 10_000) -> Dataset

Randomly shuffle rows. global_=True is a full shuffle (may buffer in DuckDB); global_=False is a bounded-memory streaming window shuffle. Without a seed, a stable seed is generated so re-iteration is consistent.

sample()

dataset.sample(n: int, seed: int = None) -> Dataset

Random sample of n rows (streaming reservoir sample).

select()

dataset.select(columns: list[str]) -> Dataset

Keep only the named columns. Pushes the projection down to the storage scan when possible.

take()

dataset.take(indices: list[int]) -> Dataset

Keep only the rows at the given indices.

head() / slice()

dataset.head(n: int = 5) -> Dataset
dataset.slice(start: int, stop: int) -> Dataset

First n rows / a contiguous row range. Limits push down to the storage scan, so Dataset("big").head(5) reads only a few rows.

filter()

dataset.filter(predicate: str | Callable[[dict], bool]) -> Dataset

Filter rows with a SQL-like expression string or a Python function. Expressions are evaluated vectorized and push down into the Iceberg scan when the filter sits directly on a platform table — prefer them for large datasets; functions run per row in Python.

recent = dataset.filter("score > 0.8 and label != None")
positive = dataset.filter(lambda x: x["label"] == 1)

map() / map_batches()

dataset.map(fn: Callable[[dict], dict], *, schema: pyarrow.Schema) -> Dataset
dataset.map_batches(fn: Callable[[RecordBatch], RecordBatch], *, schema: pyarrow.Schema) -> Dataset

Apply a function per row, or per Arrow batch (faster for vectorized work). The output schema is required so the pipeline stays lazy without executing your function early.

import pyarrow as pa

ds = dataset.map(
    lambda x: {"text": x["text"], "text_len": len(x["text"])},
    schema=pa.schema([("text", pa.string()), ("text_len", pa.int64())]),
)

join()

dataset.join(other: Dataset, keys: str, join_type: str = "inner") -> Dataset

Join with another dataset

joined = users.join(orders, keys="user_id")

train_test_split()

dataset.train_test_split(*, test_size: float = 0.2, seed: int = None) -> dict[str, Dataset]

Split into deterministic, non-overlapping train and test sets.

splits = dataset.train_test_split(test_size=0.2, seed=42)
train_ds = splits["train"]
test_ds = splits["test"]

shard()

dataset.shard(index: int, count: int) -> Dataset

Restrict the dataset to one of count disjoint partitions (for manual distributed reading; to_torch() does this automatically).

File Methods

load_files()

dataset.load_files(
    to: bool | str | dict = "auto",
    columns: list[str] = None,
    workers: int = 32,
    max_inflight_bytes: int = 256_000_000,
    on_error: str = "raise",
    fetch_timeout: float = 30.0,
) -> Dataset

Load file-backed columns locally. Files are downloaded in parallel and cached (in $MIXTRAIN_CACHE_DIR), so re-iteration does not re-download them.

to controls what each file-backed column becomes:

Pass a dict for per-column values: load_files({"image": "path"}) — a dict is an exact spec — only the columns it names are loaded; other media columns are left as records (lazy handles when consumed as Python objects). on_error="null" maps failed fetches to None instead of raising.

with_column_types()

dataset.with_column_types(column_types: str | dict = "auto") -> Dataset

Set column types (image, video, audio, ...) on the schema. "auto" infers from a data sample; pass a dict for explicit mappings. save() applies this automatically, so you only need it to override inference before saving or calling load_files().

SQL Queries

sql()

dataset.sql(sql: str, table_name: str = "data", seed: int = None) -> Dataset

Execute SQL via DuckDB. The dataset is registered as data by default, or pass table_name to choose another alias. query() is an alias.

filtered = dataset.sql("SELECT * FROM data WHERE score > 0.8")
stats = dataset.sql("SELECT label, COUNT(*) as cnt FROM t GROUP BY label", table_name="t")

Inspection

explain()

dataset.explain() -> str

Render the lazy query plan.

print(Dataset("d").filter("x > 1").head(5).explain())

cache()

dataset.cache(path: str = None, load_files: bool = True) -> Dataset

Materialize the dataset locally so iteration is fast and offline-safe.

With no arguments, platform table scans are mirrored to local Parquet under $MIXTRAIN_CACHE_DIR/datasets/<name>/<snapshot_id> — reused across runs and processes, and invalidated automatically when the table gets a new snapshot. Transforms in the plan re-run against the local mirror. With path=, this dataset's full output is materialized at an explicit location and reused on later calls (the caller owns invalidation). load_files=True also fetches file-backed column contents into the local file cache.

ds = Dataset("photo-dataset").cache()   # table + files local; train offline
prep = ds.filter("image.width >= 512").cache(path="/data/prep-v2")

set_column_types()

dataset.set_column_types(column_types: dict) -> None

Update column types for rich UI rendering.

from mixtrain import Image, Audio

dataset.set_column_types({"image_url": Image, "audio_url": Audio})

set_description()

dataset.set_description(description: str) -> dict

Set the dataset's description.

Class Methods

Dataset.exists()

Check if a dataset exists.

Dataset.exists(name: str) -> bool

if not Dataset.exists("my-dataset"):
    Dataset.from_pandas(df).save("my-dataset")

Helper Functions

list_datasets()

from mixtrain import list_datasets

datasets = list_datasets()
for ds in datasets:
    print(ds.name)

get_dataset()

from mixtrain import get_dataset

dataset = get_dataset("my-dataset")

Best practices

Guidance for feeding datasets into a training loop reliably and efficiently.

Checkpoint and resume

The loader from to_torch() can checkpoint mid-epoch and resume where it left off, so a job that dies partway through an epoch doesn't restart it. Save loader.state_dict() alongside your model/optimizer state and call loader.load_state_dict(...) after restart, before iterating:

state = loader.state_dict()              # checkpoint position
# ... later, on a fresh loader ...
loader.load_state_dict(state)            # resume the interrupted epoch

Resume against the same pipeline as the checkpoint: the same dataset (pin the version with Dataset(name, version=...)), a pinned shuffle seed (shuffle(seed=...)), and the same num_workers, world_size, and batch_size. The loader checks these and raises on a mismatch rather than resuming the wrong data. With num_workers <= 1 resume is exact; with more workers it resumes from the nearest batch boundary, so a few already-seen batches may repeat.

Missing values

to_torch() and to_tensors() take a missing argument that decides how nulls in numeric and tensor columns become dense tensors:

• "nan" (default) — nulls become NaN. Float columns fill directly; integer and boolean columns are promoted to float (and filled with NaN) only when the dataset recorded the column as nullable, so a column's dtype stays the same across batches. An integer/boolean column that was never saved (so nullability is unknown) and happens to contain a null raises — re-save the dataset, cast the column to float, or filter the nulls out first.
• "error" — raise on any null in a numeric or tensor column.

Non-numeric columns keep None under either policy.

Scaling reads and uploads

• Upload throughput — raise save()'s workers when uploading many files doesn't saturate your link.
• Download-bound reads — raise from_huggingface()'s workers to read more shards concurrently and keep a later save() fed.