Data Modules

A data module connects named datasets to a model and prepares Lightning batches. Each row is one observation; nested lists and structs stay inside that row. For example, one customer record might contain several purchases:

customer_id: customer-42
country: CA
purchases:
  - amount: 86.50
    items:
      - sku: coat-17
        quantity: 1
  - amount: 24.00
    items:
      - sku: scarf-03
        quantity: 2
churned: false

The tables and frames below are supplied by your application. relflow does not load example data or choose train/validation splits for you.

Arrow

Use ArrowDataModule with existing Arrow tables:

import relflow as rf

# model matches the record structure; these are application-supplied Arrow tables.
data = rf.ArrowDataModule(
    model=model,
    train=train_table,
    validate=validation_table,
)

Each split also accepts a pa.RecordBatch, Arrow Dataset, file path or glob, rf.source(...), or a callable returning a fresh reader or iterable of Arrow batches. A reader itself is one-shot; wrap its creation in a callable when repeated epochs need it.

A file-backed setup can be just as small:

data = rf.ArrowDataModule(
    model=model,
    train="warehouse/train/*.parquet",
    validate="warehouse/validation/*.parquet",
)

rf.source(...) adds explicit format, schema, filesystem, and partitioning configuration when needed. Keep nested structs and lists in their Arrow form. Use an explicit schema when values are empty, entirely null, or ambiguous.

Polars

Use PolarsDataModule for in-memory eager frames:

data = rf.PolarsDataModule(
    model=model,
    train=train_frame,
    validate=validation_frame,
)

It converts each frame to Arrow once and uses the same pipeline. Collect a LazyFrame explicitly before passing it in.

Splits and preparation

The available splits are train, validate, test, and predict; supply at least one. Training shuffles by default. The other splits do not. The model’s batch_size sets the batches the data module produces.

Choose splits before constructing the module. For time-based evaluation, define an observation cutoff, build features only from information available then, and keep later outcomes out of inputs. relflow does not enforce these application rules. Fit application preprocessing state on training data and reuse it for validation and prediction.

Attach an ordered preprocessor with preprocessor=prepare or preprocessor=(clean, enrich). Processing happens before sampling, shuffling, and model batching. The source and processed Arrow schemas must remain stable within each split across partitions and epochs.

For prediction, retain=("customer_id",) keeps a processed identifier in the output for downstream joins. See Batch Inference.