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.

DuckDB

Install relflow[duckdb] to use SQL queries as named splits. DuckDBDataModule opens an existing database file read-only and streams query results as Arrow batches through the same preprocessing and batching pipeline.

Suppose your application maintains warehouse/customers.duckdb with a customer_observations table. Each row has the nested structure above, a unique customer_id, and a dataset_split column assigned by the application:

observations = """
    SELECT customer_id, country, purchases, churned
    FROM customer_observations
    WHERE dataset_split = $split
    ORDER BY customer_id
"""

data = rf.DuckDBDataModule(
    model=model,
    database="warehouse/customers.duckdb",
    train=observations,
    validate=observations,
    parameters={
        "train": {"split": "train"},
        "validate": {"split": "validate"},
    },
)

Each split accepts one SELECT statement, including a query beginning with WITH. Bind values through parameters, keyed by split: a mapping supplies named parameters such as $split; a sequence supplies positional ? parameters. Parameters are bound by DuckDB rather than interpolated into SQL.

Omitting database creates a fresh in-memory database for each scan. This is useful for queries that read files directly:

data = rf.DuckDBDataModule(
    model=model,
    predict="""
        SELECT customer_id, country, purchases
        FROM read_parquet('warehouse/scoring/*.parquet')
        ORDER BY customer_id
    """,
)

The module stores the database path, SQL, parameters, and connection settings. Each scan opens its own connection and Arrow reader in the consuming process, then closes both. Pass a database path or ":memory:"; live connections, relations, and tables registered on another connection are not inputs. Nested lists and structs remain Arrow-backed, and an empty result retains its schema. Close writable connections to a database file before starting read-only scans.

Workers and distributed ranks replay the full query, then the Arrow pipeline assigns disjoint model batches. Use a deterministic ORDER BY with a unique key and keep query results stable for the duration of each pass across workers and ranks. Data may change between passes if its Arrow schema stays stable. The module does not materialize a shared snapshot, and adding workers repeats database work.

ingress_rows defaults to 4096 and controls the Arrow batch size requested from DuckDB; model.batch_size controls the final model batches. Each connection defaults to one DuckDB thread. Override connection settings with config, such as config={"threads": 2}. The usual ArrowDataModule options, including sampling, shuffling, retention, and workers, also apply. Partition-scoped preprocessors run on streamed Arrow batches converted to eager Polars frames; dataset-scoped preprocessing is unsupported for these restartable query sources.

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.