replay-rec 0.20.3rc0__py3-none-any.whl → 0.21.0rc0__py3-none-any.whl

replay/__init__.py

@@ -4,4 +4,4 @@
 # functionality removed in Python 3.12 is used in downstream packages (like lightfm)
 import setuptools as _
 
-__version__ = "0.20.3.preview"
+__version__ = "0.21.0.preview"

replay/data/dataset.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import json
+import warnings
 from collections.abc import Iterable, Sequence
 from pathlib import Path
 from typing import Callable, Optional, Union
@@ -45,6 +46,7 @@ class Dataset:
     ):
         """
         :param feature_schema: mapping of columns names and feature infos.
+            All features not specified in the schema will be assumed numerical by default.
         :param interactions: dataframe with interactions.
         :param query_features: dataframe with query features,
             defaults: ```None```.
@@ -498,6 +500,15 @@ class Dataset:
                 source=FeatureSource.QUERY_FEATURES,
                 feature_schema=updated_feature_schema,
             )
+
+        if filled_features:
+            msg = (
+                "The following features are present in the dataset but have not been specified "
+                f"by the feature schema: {[(info.column, info.feature_source.value) for info in filled_features]}. "
+                "These features will be interpreted as NUMERICAL."
+            )
+            warnings.warn(msg, stacklevel=2)
+
         return FeatureSchema(features_list=features_list + filled_features)
 
     def _fill_unlabeled_features_sources(self, feature_schema: FeatureSchema) -> list[FeatureInfo]:

replay/data/nn/__init__.py

@@ -1,6 +1,7 @@
 from replay.utils import TORCH_AVAILABLE
 
 if TORCH_AVAILABLE:
+    from .parquet import ParquetDataset, ParquetModule
     from .schema import MutableTensorMap, TensorFeatureInfo, TensorFeatureSource, TensorMap, TensorSchema
     from .sequence_tokenizer import SequenceTokenizer
     from .sequential_dataset import PandasSequentialDataset, PolarsSequentialDataset, SequentialDataset
@@ -18,6 +19,8 @@ if TORCH_AVAILABLE:
         "DEFAULT_TRAIN_PADDING_VALUE",
         "MutableTensorMap",
         "PandasSequentialDataset",
+        "ParquetDataset",
+        "ParquetModule",
         "PolarsSequentialDataset",
         "SequenceTokenizer",
         "SequentialDataset",

replay/data/nn/parquet/__init__.py

@@ -0,0 +1,22 @@
+"""
+Implementation of the ``ParquetDataset`` and its internals.
+
+``ParquetDataset`` is combination of PyTorch-compatible dataset and sampler which enables
+training and inference of models on datasets of any arbitrary size by leveraging PyArrow
+Datasets to perform batch-wise reading and processing of data from disk.
+
+``ParquetDataset`` includes support for Pytorch's distributed training framework as well as
+access to remotely stored data via PyArrow's filesystem configs.
+"""
+
+from .info.replicas import DEFAULT_REPLICAS_INFO, ReplicasInfo, ReplicasInfoProtocol
+from .parquet_dataset import ParquetDataset
+from .parquet_module import ParquetModule
+
+__all__ = [
+    "DEFAULT_REPLICAS_INFO",
+    "ParquetDataset",
+    "ParquetModule",
+    "ReplicasInfo",
+    "ReplicasInfoProtocol",
+]

replay/data/nn/parquet/collate.py

@@ -0,0 +1,29 @@
+from collections.abc import Sequence
+
+import torch
+
+from replay.data.nn.parquet.constants.batches import GeneralBatch, GeneralValue
+
+
+def dict_collate(batch: Sequence[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
+    """Simple collate function that converts a dict of values into a tensor dict."""
+    return {k: torch.cat([d[k] for d in batch], dim=0) for k in batch[0]}
+
+
+def general_collate(batch: Sequence[GeneralBatch]) -> GeneralBatch:
+    """General collate function that converts a nested dict of values into a tensor dict."""
+    result = {}
+    test_sample = batch[0]
+
+    if len(batch) == 1:
+        return test_sample
+
+    for key, test_value in test_sample.items():
+        values: Sequence[GeneralValue] = [sample[key] for sample in batch]
+        if torch.is_tensor(test_value):
+            result[key] = torch.cat(values, dim=0)
+        else:
+            assert isinstance(test_value, dict)
+            result[key] = general_collate(values)
+
+    return result

replay/data/nn/parquet/constants/batches.py

@@ -0,0 +1,8 @@
+from typing import Callable, Union
+
+import torch
+from typing_extensions import TypeAlias
+
+GeneralValue: TypeAlias = Union[torch.Tensor, "GeneralBatch"]
+GeneralBatch: TypeAlias = dict[str, GeneralValue]
+GeneralCollateFn: TypeAlias = Callable[[GeneralBatch], GeneralBatch]

replay/data/nn/parquet/constants/device.py

@@ -0,0 +1,3 @@
+import torch
+
+DEFAULT_DEVICE = torch.device("cpu")

replay/data/nn/parquet/constants/filesystem.py

@@ -0,0 +1,3 @@
+import pyarrow.fs as fs
+
+DEFAULT_FILESYSTEM = fs.LocalFileSystem()

replay/data/nn/parquet/constants/metadata.py

@@ -0,0 +1,5 @@
+SHAPE_FLAG = "shape"
+PADDING_FLAG = "padding"
+DEFAULT_PADDING = -1
+SEQUENCE_LENGTH_FLAG = "sequence_length"
+PADDING_FLAG = "padding"

replay/data/nn/parquet/fixed_batch_dataset.py

@@ -0,0 +1,157 @@
+import warnings
+from collections.abc import Iterator
+from typing import Callable, Optional, Protocol, cast
+
+import torch
+from torch.utils.data import IterableDataset
+
+from replay.data.nn.parquet.constants.batches import GeneralBatch, GeneralCollateFn
+from replay.data.nn.parquet.impl.masking import DEFAULT_COLLATE_FN
+
+
+def get_batch_size(batch: GeneralBatch, strict: bool = False) -> int:
+    """
+    Retrieves the size of the ``batch`` object.
+
+    :param batch: Batch object.
+    :param strict: If ``True``, performs additional validation. Default: ``False``.
+
+    :raises ValueError: If size mismatch is found in the batch during a strict check.
+
+    :return: Batch size.
+    """
+    batch_size: Optional[int] = None
+
+    for key, value in batch.items():
+        new_batch_size: int
+
+        if torch.is_tensor(value):
+            new_batch_size = value.size(0)
+        else:
+            assert isinstance(value, dict)
+            new_batch_size = get_batch_size(value, strict)
+
+        if batch_size is None:
+            batch_size = new_batch_size
+
+        if strict:
+            if batch_size != new_batch_size:
+                msg = f"Batch size mismatch {key}: {batch_size} != {new_batch_size}"
+                raise ValueError(msg)
+        else:
+            break
+    assert batch_size is not None
+    return cast(int, batch_size)
+
+
+def split_batches(batch: GeneralBatch, split: int) -> tuple[GeneralBatch, GeneralBatch]:
+    left: GeneralBatch = {}
+    right: GeneralBatch = {}
+
+    for key, value in batch.items():
+        if torch.is_tensor(value):
+            sub_left = value[:split, ...]
+            sub_right = value[split:, ...]
+        else:
+            sub_left, sub_right = split_batches(value, split)
+        left[key], right[key] = sub_left, sub_right
+
+    return (left, right)
+
+
+class DatasetProtocol(Protocol):
+    def __iter__(self) -> Iterator[GeneralBatch]: ...
+    @property
+    def batch_size(self) -> int: ...
+
+
+class FixedBatchSizeDataset(IterableDataset):
+    """
+    Wrapper for arbitrary datasets that fetches batches of fixed size.
+    Concatenates batches from the wrapped dataset until it reaches the specified size.
+    The last batch may be smaller than the specified size.
+    """
+
+    def __init__(
+        self,
+        dataset: DatasetProtocol,
+        batch_size: Optional[int] = None,
+        collate_fn: GeneralCollateFn = DEFAULT_COLLATE_FN,
+        strict_checks: bool = False,
+    ) -> None:
+        """
+        :param dataset: An iterable object that returns batches.
+            Generally a subclass of ``torch.utils.data.IterableDataset``.
+        :param batch_size: Desired batch size. If ``None``, will search for batch size in ``dataset.batch_size``.
+            Default: ``None``.
+        :param collate_fn: Collate function for merging batches. Default: value of ``DEFAULT_COLLATE_FN``.
+        :param strict_checks: If ``True``, additional batch size checks will be performed.
+            May affect performance. Default: ``False``.
+
+        :raises ValueError: If an invalid batch size was provided.
+        """
+        super().__init__()
+
+        self.dataset: DatasetProtocol = dataset
+
+        if batch_size is None:
+            assert hasattr(dataset, "batch_size")
+            batch_size = self.dataset.batch_size
+
+        assert isinstance(batch_size, int)
+        int_batch_size: int = cast(int, batch_size)
+
+        if int_batch_size < 1:
+            msg = f"Insufficient batch size. Got {int_batch_size=}"
+            raise ValueError(msg)
+
+        if int_batch_size < 2:
+            warnings.warn(f"Low batch size. Got {int_batch_size=}. This may cause performance issues.", stacklevel=2)
+
+        self.collate_fn: Callable = collate_fn
+        self.batch_size: int = int_batch_size
+        self.strict_checks: bool = strict_checks
+
+    def get_batch_size(self, batch: GeneralBatch) -> int:
+        return get_batch_size(batch, strict=self.strict_checks)
+
+    def __iter__(self) -> Iterator[GeneralBatch]:
+        iterator: Iterator[GeneralBatch] = iter(self.dataset)
+
+        buffer: list[GeneralBatch] = []
+        buffer_size: int = 0
+
+        while True:
+            while buffer_size < self.batch_size:
+                try:
+                    batch: GeneralBatch = next(iterator)
+                    size: int = self.get_batch_size(batch)
+
+                    buffer.append(batch)
+                    buffer_size += size
+                except StopIteration:
+                    break
+
+            if buffer_size == 0:
+                break
+
+            joined: GeneralBatch = self.collate_fn(buffer)
+            assert buffer_size == self.get_batch_size(joined)
+
+            if self.batch_size < buffer_size:
+                left, right = split_batches(joined, self.batch_size)
+                residue: int = buffer_size - self.batch_size
+                assert residue == self.get_batch_size(right)
+
+                buffer_size = residue
+                buffer = [right]
+
+                yield left
+            else:
+                buffer_size = 0
+                buffer = []
+
+                yield joined
+
+        assert buffer_size == 0
+        assert len(buffer) == 0

replay/data/nn/parquet/impl/array_1d_column.py

@@ -0,0 +1,140 @@
+from typing import Any, Union
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import torch
+
+from replay.data.nn.parquet.constants.device import DEFAULT_DEVICE
+from replay.data.nn.parquet.constants.metadata import DEFAULT_PADDING
+from replay.data.nn.parquet.metadata import (
+    Metadata,
+    get_1d_array_columns,
+    get_padding,
+    get_shape,
+)
+from replay.data.utils.typing.dtype import pyarrow_to_torch
+
+from .column_protocol import OutputType
+from .indexing import get_mask, get_offsets
+from .utils import ensure_mutable
+
+
+class Array1DColumn:
+    """
+    Representation of a 1D array column, containing a
+    list of numbers of varying length in each of its rows.
+    """
+
+    def __init__(
+        self,
+        data: torch.Tensor,
+        lengths: torch.LongTensor,
+        shape: Union[int, list[int]],
+        padding: Any = DEFAULT_PADDING,
+    ) -> None:
+        """
+        :param data: A tensor containing column data.
+        :param lengths: A tensor containing lengths of each individual row array.
+        :param shape: An integer or list of integers representing the target array shapes.
+        :param padding: Padding value to use to fill null values and match target shape.
+            Default: value of ``DEFAULT_PADDING``
+
+        :raises ValueError: If the shape provided is not one-dimensional.
+        """
+        if isinstance(shape, list) and len(shape) > 1:
+            msg = f"Array1DColumn accepts a shape of size (1,) only. Got {shape=}"
+            raise ValueError(msg)
+
+        self.padding = padding
+        self.data = data
+        self.offsets = get_offsets(lengths)
+        self.shape = shape[0] if isinstance(shape, list) else shape
+        assert self.length == torch.numel(lengths)
+
+    @property
+    def length(self) -> int:
+        return torch.numel(self.offsets) - 1
+
+    def __len__(self) -> int:
+        return self.length
+
+    @property
+    def device(self) -> torch.device:
+        assert self.data.device == self.offsets.device
+        return self.offsets.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.data.dtype
+
+    def __getitem__(self, indices: torch.LongTensor) -> OutputType:
+        indices = indices.to(device=self.device)
+        mask, output = get_mask(indices, self.offsets, self.shape)
+
+        # TODO: Test this for both 1d and 2d arrays. Add same check in 2d arrays
+        if self.data.numel() == 0:
+            mask = torch.zeros((indices.size(0), self.shape), dtype=torch.bool, device=self.device)
+            output = torch.ones((indices.size(0), self.shape), dtype=torch.bool, device=self.device) * self.padding
+            return mask, output
+
+        unmasked_values = torch.take(self.data, output)
+        masked_values = torch.where(mask, unmasked_values, self.padding)
+        assert masked_values.device == self.device
+        assert masked_values.dtype == self.dtype
+        return (mask, masked_values)
+
+
+def to_torch(array: pa.Array, device: torch.device = DEFAULT_DEVICE) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Converts a PyArrow array into a PyTorch tensor.
+
+    :param array: Original PyArrow array.
+    :param device: Target device to send the resulting tensor to. Default: value of ``DEFAULT_DEVICE``.
+
+    :return: A PyTorch tensor obtained from original array.
+    """
+    flatten = pc.list_flatten(array)
+    lengths = pc.list_value_length(array).cast(pa.int64())
+
+    # Copying to be mutable
+    flatten_torch = torch.asarray(
+        ensure_mutable(flatten.to_numpy()),
+        device=device,
+        dtype=pyarrow_to_torch(flatten.type),
+    )
+
+    # Copying to be mutable
+    lengths_torch = torch.asarray(
+        ensure_mutable(lengths.to_numpy()),
+        device=device,
+        dtype=torch.int64,
+    )
+    return (lengths_torch, flatten_torch)
+
+
+def to_array_1d_columns(
+    data: pa.RecordBatch,
+    metadata: Metadata,
+    device: torch.device = DEFAULT_DEVICE,
+) -> dict[str, Array1DColumn]:
+    """
+    Converts a PyArrow batch of data to a set of ``Array1DColums``s.
+    This function filters only those columns matching its format from the full batch.
+
+    :param data: A PyArrow batch of column data.
+    :param metadata: Metadata containing information about columns' formats.
+    :param device: Target device to send column tensors to. Default: value of ``DEFAULT_DEVICE``
+
+    :return: A dict of tensors containing dataset's numeric columns.
+    """
+    result: dict[str, Array1DColumn] = {}
+
+    for column_name in get_1d_array_columns(metadata):
+        lengths, torch_array = to_torch(data.column(column_name), device=device)
+        result[column_name] = Array1DColumn(
+            data=torch_array,
+            lengths=lengths,
+            padding=get_padding(metadata, column_name),
+            shape=get_shape(metadata, column_name),
+        )
+    return result

replay/data/nn/parquet/impl/array_2d_column.py

@@ -0,0 +1,160 @@
+from typing import Any
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import torch
+
+from replay.data.nn.parquet.constants.device import DEFAULT_DEVICE
+from replay.data.nn.parquet.constants.metadata import DEFAULT_PADDING
+from replay.data.nn.parquet.metadata import (
+    Metadata,
+    get_2d_array_columns,
+    get_padding,
+    get_shape,
+)
+from replay.data.utils.typing.dtype import pyarrow_to_torch
+
+from .column_protocol import OutputType
+from .indexing import get_mask, get_offsets
+from .utils import ensure_mutable
+
+
+class Array2DColumn:
+    """
+    Representation of a 2D array column, containing nested
+    lists of numbers of varying length in each of its rows.
+    """
+
+    def __init__(
+        self,
+        data: torch.Tensor,
+        outer_lengths: torch.LongTensor,
+        inner_lengths: torch.LongTensor,
+        shape: list[int],
+        padding: Any = DEFAULT_PADDING,
+    ) -> None:
+        """
+        :param data: A tensor containing column data.
+        :param outer_lengths: A tensor containing outer lengths (first dim) of each individual row array.
+        :param inner_lengths: A tensor containing inner lengths (second dim) of each individual row array.
+        :param shape: An integer or list of integers representing the target array shapes.
+        :param padding: Padding value to use to fill null values and match target shape.
+            Default: value of ``DEFAULT_PADDING``
+
+        :raises ValueError: If the shape provided is not two-dimensional.
+        """
+        self.padding = padding
+        self.data = data
+        self.inner_offsets = get_offsets(inner_lengths)
+        self.outer_offsets = get_offsets(outer_lengths)
+        if len(shape) != 2:
+            msg = f"Array2DColumn accepts a shape of size (2,) only. Got {shape=}"
+            raise ValueError(msg)
+        self.shape: list[int] = shape
+
+    @property
+    def length(self) -> int:
+        return torch.numel(self.outer_offsets) - 1
+
+    def __len__(self) -> int:
+        return self.length
+
+    @property
+    def device(self) -> torch.device:
+        assert self.data.device == self.inner_offsets.device
+        assert self.data.device == self.outer_offsets.device
+        return self.inner_offsets.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.data.dtype
+
+    def __getitem__(self, indices: torch.LongTensor) -> OutputType:
+        indices = indices.to(device=self.device)
+        outer_mask, outer_output = get_mask(indices, self.outer_offsets, self.shape[0])
+        left_bound = outer_output.min().item()
+        right_bound = outer_output.max().item()
+        outer_output -= left_bound
+
+        inner_indices = torch.arange(left_bound, right_bound + 1, device=indices.device)
+        inner_mask, output = get_mask(inner_indices, self.inner_offsets, self.shape[1])
+
+        final_indices = output[outer_output]
+        inner_final_mask = inner_mask[outer_output]
+
+        unmasked_values = torch.take(self.data, final_indices)
+        outer_final_mask = outer_mask.unsqueeze(-1).repeat(1, 1, unmasked_values.size(-1))
+        mask = inner_final_mask * outer_final_mask
+
+        masked_values = torch.where(mask, unmasked_values, self.padding)
+        assert masked_values.device == self.device
+        assert masked_values.dtype == self.dtype
+        return (mask, masked_values)
+
+
+def to_torch(
+    array: pa.Array,
+    device: torch.device = DEFAULT_DEVICE,
+) -> tuple[torch.LongTensor, torch.LongTensor, torch.Tensor]:
+    """
+    Converts a PyArrow array into a PyTorch tensor.
+
+    :param array: Original PyArrow array.
+    :param device: Target device to send the resulting tensor to. Default: value of ``DEFAULT_DEVICE``.
+
+    :return: A PyTorch tensor obtained from original array.
+    """
+    flatten_dim0 = pc.list_flatten(array)
+    flatten = pc.list_flatten(flatten_dim0)
+
+    outer_lengths = pc.list_value_length(array).cast(pa.int64())
+    inner_lengths = pc.list_value_length(flatten_dim0).cast(pa.int64())
+
+    # Copying to be mutable
+    flatten_torch = torch.asarray(
+        ensure_mutable(flatten.to_numpy()),
+        device=device,
+        dtype=pyarrow_to_torch(flatten.type),
+    )
+
+    # Copying to be mutable
+    outer_lengths_torch = torch.asarray(
+        ensure_mutable(outer_lengths.to_numpy()),
+        device=device,
+        dtype=torch.int64,
+    )
+    inner_lengths_torch = torch.asarray(
+        ensure_mutable(inner_lengths.to_numpy()),
+        device=device,
+        dtype=torch.int64,
+    )
+    return (outer_lengths_torch, inner_lengths_torch, flatten_torch)
+
+
+def to_array_2d_columns(
+    data: pa.RecordBatch,
+    metadata: Metadata,
+    device: torch.device = DEFAULT_DEVICE,
+) -> dict[str, Array2DColumn]:
+    """
+    Converts a PyArrow batch of data to a set of ``Array2DColums``s.
+    This function filters only those columns matching its format from the full batch.
+
+    :param data: A PyArrow batch of column data.
+    :param metadata: Metadata containing information about columns' formats.
+    :param device: Target device to send column tensors to. Default: value of ``DEFAULT_DEVICE``
+
+    :return: A dict of tensors containing dataset's numeric columns.
+    """
+    result = {}
+
+    for column_name in get_2d_array_columns(metadata):
+        outer_lengths, inner_lengths, torch_array = to_torch(data.column(column_name), device=device)
+        result[column_name] = Array2DColumn(
+            data=torch_array,
+            outer_lengths=outer_lengths,
+            inner_lengths=inner_lengths,
+            padding=get_padding(metadata, column_name),
+            shape=get_shape(metadata, column_name),
+        )
+    return result

replay/data/nn/parquet/impl/column_protocol.py

@@ -0,0 +1,17 @@
+from typing import Protocol
+
+import torch
+
+OutputType = tuple[torch.BoolTensor, torch.Tensor]
+
+
+class ColumnProtocol(Protocol):
+    def __len__(self) -> int: ...
+
+    @property
+    def length(self) -> int: ...
+
+    @property
+    def device(self) -> torch.device: ...
+
+    def __getitem__(self, indices: torch.LongTensor) -> OutputType: ...