Streaming Parquet Row Groups into PyTorch Efficiently

Hello! Last lesson focused on what happens inside a recurrent model as it processes a long sequence. This lesson addresses the data path feeding that model. A recurrence is of little use if preparing its inputs first loads the entire dataset into memory—or accidentally joins two independent sequences.

You will inspect Parquet row groups, read one sequence at a time, and expose those sequences through a PyTorch IterableDataset. Allow about 40 minutes, including the short video, documentation, and code walkthrough.

A row group can be a useful sequence boundary

A Parquet file stores rows in row groups. Within each group, values are stored in column-specific chunks. That organization lets a reader request particular columns from one group without materializing every row and column in the file.

Watch the brief structural explanation, then use the diagram to locate the two row groups, their column chunks, and the footer containing metadata.

For sequence modeling, the convenient case is exactly one complete, chronologically ordered sequence per row group. Then one group can become one dataset item: a feature tensor with shape $(T,F)$, where $T$ is that sequence’s length and $F$ is its number of features. This is a data-layout requirement, not a property Parquet guarantees. A writer may put several sequences in one group or split one sequence across groups. We must check the layout rather than assume it.

Apache Arrow’s Parquet guide shows the key distinction between reading the entire file and reading an individual group.

The resulting plan is simple: open a file, inspect its group count and row counts, then read groups one at a time. Calling pq.read_table(path) instead would concatenate all groups into one Arrow table—precisely what we want to avoid for a large training file.

Read only the columns the model needs

Column selection matters as much as row-group selection. If a file contains features, targets, identifiers, timestamps, and other fields, the loader should decode only those needed for its current purpose. The same Arrow guide demonstrates selecting columns; its example uses read_table, and we will apply that idea to read_row_group.

Before implementing a loader, inspect an actual file’s schema and group sizes:

import pyarrow.parquet as pq

pf = pq.ParquetFile("train.parquet")
print(pf.schema)
print("row groups:", pf.num_row_groups)

for rg in range(pf.num_row_groups):
    print(rg, pf.metadata.row_group(rg).num_rows)

This reads metadata, not all feature values. It answers an important memory question: is the largest single group small enough to decode? Streaming by group limits peak memory only if each group fits. A file with one enormous row group cannot be made sequence-sized merely by iterating over its group numbers; it needs a finer-grained reading or rewriting strategy.

For the implementation below, assume a demonstration schema with seq_id, step, feature columns f0 and f1, and target columns y0 and y1. These are placeholders, not asserted competition column names. Substitute the actual schema after inspecting it.

Turn groups into PyTorch dataset items

A map-style PyTorch dataset answers requests such as dataset[7]. Here, our natural operation is sequential: open a Parquet file and yield each group as it is read. PyTorch calls this an IterableDataset.

Here is a complete loader for the stated one-sequence-per-group layout. Its safeguards are deliberately visible: it rejects oversized groups, groups containing multiple sequence IDs, non-increasing steps, and non-finite numeric values instead of silently training on them.

import numpy as np
import pyarrow.parquet as pq
import torch
from torch.utils.data import IterableDataset, get_worker_info


class ParquetSequences(IterableDataset):
    def __init__(self, paths, feature_cols, target_cols, *, max_rows):
        super().__init__()
        self.paths = tuple(paths)  # Pass paths in the intended file order.
        self.feature_cols = tuple(feature_cols)
        self.target_cols = tuple(target_cols)
        self.max_rows = max_rows
        if not self.paths or not self.feature_cols or not self.target_cols:
            raise ValueError("Paths, features, and targets must be nonempty.")

    def __iter__(self):
        # With multiple workers, assign each row group to exactly one worker.
        info = get_worker_info()
        worker_id = 0 if info is None else info.id
        workers = 1 if info is None else info.num_workers
        ordinal = 0

        columns = list(dict.fromkeys(
            ("seq_id", "step", *self.feature_cols, *self.target_cols)
        ))

        for path in self.paths:
            pf = pq.ParquetFile(path)  # Open inside the worker, not in __init__.
            for rg in range(pf.num_row_groups):
                mine = ordinal % workers == worker_id
                ordinal += 1
                if not mine:
                    continue

                rows = pf.metadata.row_group(rg).num_rows
                if not 0 < rows <= self.max_rows:
                    raise ValueError(f"{path}, group {rg}: invalid row count {rows}")

                table = pf.read_row_group(rg, columns=columns)

                ids = table["seq_id"].to_pylist()
                if len(set(ids)) != 1:
                    raise ValueError(f"{path}, group {rg}: mixed sequence IDs")

                steps = np.asarray(
                    table["step"].to_numpy(zero_copy_only=False)
                )
                if not np.all(steps[1:] > steps[:-1]):
                    raise ValueError(f"{path}, group {rg}: steps not increasing")

                def matrix(names):
                    return np.column_stack([
                        np.asarray(
                            table[name].to_numpy(zero_copy_only=False),
                            dtype=np.float32,
                        )
                        for name in names
                    ])

                x = matrix(self.feature_cols)
                y = matrix(self.target_cols)
                if not np.isfinite(x).all() or not np.isfinite(y).all():
                    raise ValueError(f"{path}, group {rg}: non-finite values")

                yield {
                    "x": torch.from_numpy(x),
                    "y": torch.from_numpy(y),
                    "seq_id": ids[0],
                }

Each yield contains one complete group. The first dimension of x and y is the number of time steps; their second dimensions are the numbers of selected features and targets. torch.from_numpy shares storage with its NumPy array, so it does not require another tensor-sized copy at that point. Arrow decoding, NumPy conversion, and stacking can still allocate memory. “Streaming” means memory is bounded by groups and in-flight items, not that reading is zero-copy.

The ordinal assignment prevents duplicate groups if workers are enabled. For the initial sequence-modeling pipeline, however, use num_workers=0: it is easier to debug and preserves the order in which this iterator yields groups. Multiple workers can deliver groups in a different effective order, which matters if a later model carries state between groups. They can also hold or prefetch several decoded groups at once, increasing peak memory.

Check the contract with a tiny file

This small fixture writes two distinct sequences as two row groups. It checks the loader’s output without ever collecting a large training dataset.

import os
import tempfile
import pyarrow as pa
import pyarrow.parquet as pq
from torch.utils.data import DataLoader


def make_sequence(seq_id, length):
    return pa.table({
        "seq_id": [seq_id] * length,
        "step": list(range(length)),
        "f0": [float(i) for i in range(length)],
        "f1": [float(seq_id)] * length,
        "y0": [float(i + 1) for i in range(length)],
        "y1": [float(i - 1) for i in range(length)],
    })


with tempfile.TemporaryDirectory() as directory:
    path = os.path.join(directory, "tiny.parquet")
    first = make_sequence(17, 3)
    second = make_sequence(23, 4)

    with pq.ParquetWriter(path, first.schema) as writer:
        writer.write_table(first, row_group_size=first.num_rows)
        writer.write_table(second, row_group_size=second.num_rows)

    assert pq.ParquetFile(path).num_row_groups == 2

    dataset = ParquetSequences(
        [path], ("f0", "f1"), ("y0", "y1"), max_rows=10
    )
    loader = DataLoader(dataset, batch_size=None, num_workers=0)

    observed = [
        (item["seq_id"], tuple(item["x"].shape), tuple(item["y"].shape))
        for item in loader
    ]
    assert observed == [
        (17, (3, 2), (3, 2)),
        (23, (4, 2), (4, 2)),
    ]

batch_size=None is intentional: one yielded sequence is one DataLoader item, even when lengths differ. Ordinary automatic batching would try to combine items; sequences of lengths 3 and 4 cannot simply be stacked into one rectangular tensor without a padding policy. The test stores only identifiers and shapes. In real training, iterate over loader directly rather than calling list(loader) on its tensor outputs.

The loader verifies that each individual group has one ID, but it does not prove that an ID appears in only one group. Confirm that property from the data-generation contract or a separate inspection. If a sequence spans groups, you cannot treat each group as an independent sequence and reset its recurrent state at each boundary. Conversely, if a group mixes IDs, it must be split at the true boundaries. Do not sort rows, shuffle groups, or infer boundaries casually: chronological order and independence are part of the model’s input contract.

As a rough sizing check, $20,000$ rows with 128 float32 features and two float32 targets require about $20,000(128+2)4$, or 10.4 MB, for the final numeric arrays alone. Decoded Arrow columns, temporary arrays, metadata, and any prefetched groups add to that. Choose max_rows from the available memory budget, then measure peak process memory on a representative large group.

Takeaways

Parquet row groups give you a practical unit for bounded reads, and column selection avoids decoding fields the model does not need. A PyTorch IterableDataset can yield one $(T,F)$ tensor at a time, provided you verify that each group really represents one ordered, complete sequence. Start with one DataLoader worker, inspect group sizes before reading, and keep the iteration streaming all the way into the training loop.

The next lesson will use this sequence-aware pipeline to reproduce the competition’s masked evaluation metric and build a validation procedure that avoids leakage.