LanceDB schemas

Opt a protobuf message into LanceDB codegen with (pgml.lancedb).enable = true. py-gen-ml then emits lancedb.pydantic.LanceModel classes for that message and every nested message type reachable from it, plus small helpers to create a table with the generated schema.

Feature-row messages should also set (pgml.kind) = FEATURE_ROW so other generators can discover the same contract. See Message kinds. Kind does not replace the LanceDB opt-in.

Install the extra

The generator itself ships with py-gen-ml. Generated modules import lancedb, so install the optional extra in projects that use them:

pip install 'py-gen-ml[lancedb]'

Or with uv:

uv add 'py-gen-ml[lancedb]'

The torch DataLoader example below also needs PyTorch (pip install torch).

Enable the generator when you run codegen (it is off by default):

py-gen-ml path/to/schema.proto --generators=base,patch,sweep,cli_args,lancedb

Concepts

Message opt-in

Annotate the root row type with (pgml.lancedb) (and preferably (pgml.kind) = FEATURE_ROW):

Option	Meaning
`enable`	Required. When `true`, emit LanceModels for this message and its nested types.
`table_name`	Optional default table name. Falls back to the message name in snake_case.

Nested messages do not need their own (pgml.lancedb) option. Walking field types from each enabled root builds the set of models to emit.

Vector columns

LanceDB fixed-size vectors use lancedb.pydantic.Vector(dim). Mark a field with (pgml.lancedb_field).vector_dim = N (typically on repeated float):

repeated float embedding = 2 [(pgml.lancedb_field).vector_dim = 8];

That field is generated as embedding: Vector(8) instead of List[float].

Merge keys

Mark join columns for LanceDB merge_insert with (pgml.lancedb_field).merge_key = true:

string id = 1 [(pgml.lancedb_field) = { merge_key: true }];

Codegen emits *_merge_on() -> list[str] listing those fields. Pass that to py_gen_ml.bridges.merge_rows(table, rows, on=...) for upserts. When no field is marked, use py_gen_ml.bridges.append_rows(table, rows) instead.

What gets generated

For a proto lancedb_demo.proto, enabling the generator writes lancedb_demo_lancedb.py containing:

One LanceModel subclass per message in the nested closure
*_table_name() returning the configured (or default) table name
*_merge_on() returning fields marked merge_key
create_*_table(db, ...) calling db.create_table(..., schema=<RootModel>, exist_ok=True)

Enums are stored as str. Scalar / list / nested-message fields follow the usual Python mapping that LanceDB converts to Arrow types.

Torch Dataset wrappers are not generated. LanceDB tables already implement PyTorch's Dataset contract, so you pass the table to torch.utils.data.DataLoader directly (see LanceDB's PyTorch integration).

Annotate a schema

syntax = "proto3";

package demo;

import "py_gen_ml/extensions.proto";

// Metadata nested under each training row.
message SampleMeta {
    string label = 1;
    int32 split_id = 2;
}

// Root table schema for a LanceDB dataset of embedding rows.
message EmbeddingSample {
    option (pgml.kind) = FEATURE_ROW;
    option (pgml.lancedb) = {
        enable: true;
        table_name: "embedding_samples";
    };

    string id = 1 [(pgml.lancedb_field) = { merge_key: true }];
    // Fixed-size vector column for similarity search / training.
    repeated float embedding = 2 [(pgml.lancedb_field).vector_dim = 8];
    SampleMeta meta = 3;
}

Generated models

# Autogenerated code. DO NOT EDIT.
import typing
from lancedb.db import DBConnection
from lancedb.pydantic import LanceModel, Vector
from lancedb.table import LanceTable


class SampleMeta(LanceModel):
    """Metadata nested under each training row."""

    label: str
    split_id: int


class EmbeddingSample(LanceModel):
    """Root table schema for a LanceDB dataset of embedding rows."""

    id: str
    embedding: Vector(8)
    """Fixed-size vector column for similarity search / training."""

    meta: SampleMeta


def embedding_sample_table_name() -> str:
    """Default LanceDB table name for :class:`EmbeddingSample`."""
    return 'embedding_samples'


def embedding_sample_merge_on() -> typing.List[str]:
    """Join columns for ``merge_insert`` / ``merge_rows`` (fields with ``(pgml.lancedb_field).merge_key``)."""
    return ['id']


def create_embedding_sample_table(
    db: DBConnection,
    *,
    name: typing.Optional[str] = None,
    exist_ok: bool = True,
    **kwargs: typing.Any
) -> LanceTable:
    """Create a LanceDB table whose schema is :class:`EmbeddingSample`.

    ``db`` is a connection from ``lancedb.connect(...)``.
    By default ``exist_ok=True`` opens the table if it already exists.
    Pass ``mode="overwrite"`` (via kwargs) to replace an existing table.
    Load rows for training via Arrow (``table.to_arrow()``) or LanceDB's
    ``Permutation`` streaming API, then hand tensors to
    ``torch.utils.data.DataLoader`` as needed.
    """
    return db.create_table(
        name or embedding_sample_table_name(),
        schema=EmbeddingSample,
        exist_ok=exist_ok,
        **kwargs
    )

End-to-end: insert rows and load with PyTorch

LanceDB's Table can be handed to torch.utils.data.DataLoader without a custom Dataset class. Batches arrive as Arrow tables; use a collate_fn to turn them into tensors.

For purely numeric scalar columns, LanceDB's lancedb.util.tbl_to_tensor is enough. Our demo schema also has strings, fixed-size vectors, and a nested struct, so the snippet uses a small collate helper that keeps those columns usable in training.

"""Insert dummy LanceDB rows and load them with LanceDB's PyTorch integration."""
from __future__ import annotations

import tempfile
from pathlib import Path

import lancedb
import pyarrow as pa
import torch
from torch.utils.data import DataLoader

from pgml_out.lancedb_demo_lancedb import (
    EmbeddingSample,
    SampleMeta,
    create_embedding_sample_table,
)


def collate_samples(batch: pa.Table) -> dict:
    """Convert an Arrow batch from LanceDB into tensors for training.

    LanceDB tables implement PyTorch's Dataset contract and yield Arrow batches
    via ``__getitems__``. For purely numeric scalar tables you can pass
    ``lancedb.util.tbl_to_tensor`` as ``collate_fn`` instead; this helper keeps
    string / vector / nested columns usable for our demo schema.
    See https://docs.lancedb.com/training/torch
    """
    return {
        'id': batch.column('id').to_pylist(),
        'embedding': torch.tensor(batch.column('embedding').to_pylist(), dtype=torch.float32),
        'label': [meta['label'] for meta in batch.column('meta').to_pylist()],
    }


def main() -> None:
    with tempfile.TemporaryDirectory() as tmp:
        db = lancedb.connect(str(Path(tmp) / 'demo.lancedb'))
        table = create_embedding_sample_table(db, mode='overwrite')

        rows = [
            EmbeddingSample(
                id=f'sample-{i}',
                embedding=[float(i)] * 8,
                meta=SampleMeta(label='even' if i % 2 == 0 else 'odd', split_id=i % 3),
            ) for i in range(6)
        ]
        table.add(rows)

        # Pass the LanceDB table straight to DataLoader — no custom Dataset class.
        loader = DataLoader(table, batch_size=2, shuffle=False, collate_fn=collate_samples)
        batches = list(loader)
        assert len(batches) == 3
        first = batches[0]
        assert first['id'] == ['sample-0', 'sample-1']
        assert first['embedding'].shape == (2, 8)
        assert first['label'] == ['even', 'odd']

        print(f'loaded {sum(len(batch["id"]) for batch in batches)} rows across {len(batches)} batches')
        for batch_idx, batch in enumerate(batches):
            print(
                f'batch {batch_idx}: ids={list(batch["id"])} '
                f'embedding_shape={tuple(batch["embedding"].shape)}'
            )


if __name__ == '__main__':
    main()

Run it from the docs snippets project (after codegen):

cd docs/snippets
uv sync --extra lancedb
uv run python -m snippets.lancedb_torch_demo

You should see three batches of two row ids each, with embedding tensors of shape (2, 8).

Going further

LanceDB's Permutation API is useful when you want column projection, splits, or shuffle without materializing a copy. For example, selecting only the embedding column reduces I/O before the DataLoader:

from lancedb.permutation import Permutation

permutation = Permutation.identity(table).select_columns(["embedding"])

See the LanceDB docs for tbl_to_tensor, torch_col format, and multi-worker loaders (num_workers, multiprocessing_context="forkserver").

Overview

This feature bridges your protobuf row schema to LanceDB table schemas without hand-maintaining parallel Pydantic models:

You declare the row shape (and nested structs / vector dims) in .proto.
py-gen-ml --generators=...,lancedb emits LanceModel classes and create_*_table helpers.
You insert validated rows (model instances or compatible dicts) into LanceDB.
For training, pass the table to torch.utils.data.DataLoader using LanceDB's PyTorch integration (with a collate suited to your column types).

Install py-gen-ml[lancedb] for the runtime import; keep lancedb in --generators whenever you regenerate so *_lancedb.py stays in sync with the proto.