Data Pipeline

Torch-RecHub provides common dataset classes, DataLoader generators, and preprocessing utilities. These components pass data that is already encoded and padded to models; category vocabularies, missing values, and split policy remain the user's responsibility.

Datasets

TorchDataset

Training/validation dataset with features and labels.

from torch_rechub.utils.data import TorchDataset
# x may be a feature-to-array/Tensor mapping or a DataFrame
dataset = TorchDataset(x, y)

x must support .items(), and each column must support row indexing. A feature dictionary or pandas.DataFrame is typical.

PredictDataset

Prediction-only dataset (features only).

from torch_rechub.utils.data import PredictDataset
dataset = PredictDataset(x)

Data Generators

DataGenerator

Build dataloaders for ranking / multi-task models.

from torch_rechub.utils.data import DataGenerator

dg = DataGenerator(x, y)
train_dl, val_dl, test_dl = dg.generate_dataloader(
    split_ratio=[0.7, 0.1],  # 70% train, 10% validation, remaining 20% test
    batch_size=256,
    num_workers=8,
)

split_ratio contains the train and validation fractions; the test set uses the remainder. If the data is already split, omit split_ratio and pass each split explicitly:

train_dl, val_dl, test_dl = dg.generate_dataloader(
    x_val=x_val,
    y_val=y_val,
    x_test=x_test,
    y_test=y_test,
    batch_size=256,
)

The automatic path uses torch.utils.data.random_split. Call torch.manual_seed(...) before creating the DataLoaders when a reproducible split is required.

MatchDataGenerator

Build dataloaders for matching/retrieval models.

from torch_rechub.utils.data import MatchDataGenerator

dg = MatchDataGenerator(x, y)
train_dl, test_dl, item_dl = dg.generate_dataloader(
    x_test_user=x_test_user,
    x_all_item=x_all_item,
    batch_size=256,
    num_workers=8,
)

Utilities

get_auto_embedding_dim

Compute embedding dim from vocab size: int(floor(6 * num_classes**0.25)).

from torch_rechub.utils.data import get_auto_embedding_dim
embed_dim = get_auto_embedding_dim(num_classes=1000)

get_loss_func

Return the default loss by task type. Classification returns torch.nn.BCELoss, so its input must be a probability in [0, 1], not an unprocessed logit. Regression returns torch.nn.MSELoss.

from torch_rechub.utils.data import get_loss_func
loss_fn = get_loss_func(task_type="classification")

Sequence Data

SequenceDataGenerator is used by next-item tasks such as HSTU and HLLM. It accepts four NumPy arrays with the same first dimension. Each batch is (seq_tokens, seq_positions, seq_time_diffs, targets).

from torch_rechub.utils.data import SequenceDataGenerator

generator = SequenceDataGenerator(
    seq_tokens,
    seq_positions,
    targets,
    seq_time_diffs,
)

# Already split data: the return value is a one-element tuple
train_dl = generator.generate_dataloader(
    batch_size=32,
    num_workers=0,
)[0]

# Automatic splitting requires three ratios that sum to 1
train_dl, val_dl, test_dl = generator.generate_dataloader(
    batch_size=32,
    split_ratio=(0.7, 0.1, 0.2),
)

Parquet Streaming Dataset

In industrial scenarios, feature engineering is typically done by PySpark on big data clusters, with data volumes reaching GB to TB scale. Using spark_df.toPandas() directly causes Driver OOM.

Torch-RecHub provides ParquetIterableDataset for streaming Parquet files generated by Spark without loading all data into memory.

Installation

Parquet data loading requires the bigdata extra:

python -m pip install "torch-rechub[bigdata]"

ParquetIterableDataset

Inherits from torch.utils.data.IterableDataset with multi-worker support.

from torch.utils.data import DataLoader
from torch_rechub.data.dataset import ParquetIterableDataset

dataset = ParquetIterableDataset(
    ["/data/train1.parquet", "/data/train2.parquet"],
    columns=["user_id", "item_id", "label"],  # Optional
    batch_size=1024,
)

loader = DataLoader(dataset, batch_size=None, num_workers=4)

for batch in loader:
    user_id = batch["user_id"]  # torch.Tensor
    item_id = batch["item_id"]  # torch.Tensor
    label = batch["label"]      # torch.Tensor

Parameters:

• file_paths: List of Parquet file paths
• columns: Column names to read; None reads all columns
• batch_size: Rows per batch (default: 1024)

Features:

• Streaming: Uses PyArrow Scanner for constant memory usage
• Multi-worker: Automatically partitions files across workers
• Type conversion: Converts PyArrow arrays to PyTorch Tensors
• Nested arrays: Supports Spark Array columns as 2D Tensors

Working with Spark

# ========== Spark Side ==========
# df.write.parquet("/data/train.parquet")

# ========== PyTorch Side ==========
import glob
from torch.utils.data import DataLoader
from torch_rechub.data.dataset import ParquetIterableDataset

file_paths = glob.glob("/data/train.parquet/*.parquet")
dataset = ParquetIterableDataset(file_paths, batch_size=2048)
loader = DataLoader(dataset, batch_size=None, num_workers=8)

Supported Types

Parquet/Arrow Type	Result
int8/16/32/64	torch.float32
float32/64	torch.float32
boolean	torch.float32
list/array	torch.Tensor (2D)

Note: Nested arrays require equal row lengths; otherwise raises ValueError.

Type limits: The current converter does not support Arrow string columns. Encode category strings as numeric IDs before writing Parquet. Every supported scalar column, including integer IDs, becomes torch.float32; models cast embedding inputs back to integer indices. Because float32 cannot represent every integer above 2^24 exactly, remap large business IDs to compact contiguous indices first.

Typical Flow

1. Define features (Dense/Sparse/Sequence).
2. Load raw data.
3. Encode categorical features (e.g., LabelEncoder).
4. Process sequences (pad/truncate).
5. Construct samples (e.g., negative sampling).
6. Use DataGenerator, MatchDataGenerator, or SequenceDataGenerator to build dataloaders for the task.
7. Train models with the trainers.