import math
import os
from collections import Counter
from collections.abc import Iterator
from typing import Any, Dict
import polars as pl
import torch
from loguru import logger
from torch.utils.data import get_worker_info
from sequifier.helpers import (
PANDAS_TO_TORCH_TYPES,
columns_from_slice,
get_left_pad_lengths_from_preprocessed_data,
)
from sequifier.io.batch import SequifierBatch
from sequifier.io.folder_dataset import LazyFolderDataset
from sequifier.io.iteration_state import (
read_shared_int,
resolve_resume_worker,
skip_samples_for_batches,
)
from sequifier.io.sample_order import (
SampleOrderPlan,
curriculum_positions_from_parquet,
epoch_file_order,
logical_sample_positions,
)
from sequifier.io.window_sampling import (
build_window_batch,
target_valid_from_offsets,
validate_split_bounds_available,
)
from sequifier.typechecking import beartype
[docs]class SequifierDatasetFromFolderParquetLazy(LazyFolderDataset):
"""Streams long-format Parquet chunks into rank/worker-aligned batches."""
@beartype
def __init__(self, data_path: str, config: Any, shuffle: bool = True):
super().__init__(data_path, config, shuffle, "merge_output: False")
metadata = self.metadata
self._initialize_window_layout()
self.batch_files_info = []
raw_file_infos = list(metadata["batch_files"])
if self.file_order == "name":
raw_file_infos.sort(key=lambda item: item["path"])
for raw_file_info in raw_file_infos:
file_info = dict(raw_file_info)
file_info["stored_samples"] = int(raw_file_info["samples"])
histogram = raw_file_info.get("target_valid_from_histogram")
if histogram is None and not self.sampling_plan.legacy_single_window:
file_path = os.path.join(self.data_dir, file_info["path"])
position_rows = (
pl.scan_parquet(file_path)
.group_by(["sequenceId", "subsequenceId"])
.agg(
pl.col("leftPadLength").first(),
pl.col("startItemPosition").first(),
pl.col("splitStartItemPosition").first(),
)
.select(
pl.max_horizontal(
"leftPadLength",
pl.col("splitStartItemPosition")
- pl.col("startItemPosition"),
).alias("targetValidFrom")
)
.collect()
.get_column("targetValidFrom")
.to_list()
)
histogram = {
str(value): count for value, count in Counter(position_rows).items()
}
if self.sampling_plan.legacy_single_window:
file_info["samples"] = file_info["stored_samples"]
else:
assert histogram is not None
file_info["samples"] = (
self.sampling_plan.sample_count_from_target_valid_from_histogram(
histogram
)
)
if file_info["samples"] > 0:
self.batch_files_info.append(file_info)
self.total_samples = sum(info["samples"] for info in self.batch_files_info)
if self.total_samples == 0:
raise ValueError("No usable model windows were found in the dataset.")
self.column_torch_types = {
col: PANDAS_TO_TORCH_TYPES[config.column_data_types[col]]
for col in config.column_data_types
}
self.target_samples = self._get_target_samples()
self.total_batches = self._calculate_total_batches(self.target_samples)
logger.info(
f"Lazy Parquet Dataset mapped with {self.target_samples} samples and {self.total_batches} batches."
)
@beartype
def __iter__(
self,
) -> Iterator[SequifierBatch]:
world_size = self.world_size
rank = self.rank
worker_info = get_worker_info()
physical_worker_id = worker_info.id if worker_info is not None else 0
num_workers = worker_info.num_workers if worker_info is not None else 1
epoch = read_shared_int(self._epoch_state)
start_batch = read_shared_int(self._start_batch_state)
num_files = len(self.batch_files_info)
original_files_for_this_rank = list(range(rank, num_files, world_size))
rank_real_samples = sum(
self.batch_files_info[i]["samples"] for i in original_files_for_this_rank
)
files_for_this_rank = original_files_for_this_rank.copy()
if not files_for_this_rank:
if self.target_samples == 0:
return
files_for_this_rank = [rank % num_files]
base_samples_per_worker = self.target_samples // num_workers
remainder = self.target_samples % num_workers
worker_sample_counts = [
base_samples_per_worker + (1 if i < remainder else 0)
for i in range(num_workers)
]
worker_batch_counts = [
math.ceil(sample_count / self.batch_size)
for sample_count in worker_sample_counts
]
worker_id, skip_batches = resolve_resume_worker(
start_batch,
physical_worker_id,
num_workers,
worker_batch_counts,
)
worker_start_sample = 0
for i in range(worker_id):
worker_start_sample += worker_sample_counts[i]
worker_target_samples = worker_sample_counts[worker_id]
worker_end_sample = worker_start_sample + worker_target_samples
skipped_samples = skip_samples_for_batches(
skip_batches, self.batch_size, worker_target_samples
)
worker_start_sample += skipped_samples
worker_target_samples -= skipped_samples
if worker_target_samples <= 0:
return
file_order = epoch_file_order(
len(files_for_this_rank),
seed=self.config.seed,
epoch=epoch,
shuffle=self.shuffle,
policy=self.file_order,
)
ordered_files = [files_for_this_rank[i] for i in file_order]
extended_files = []
current_samples = 0
file_idx = 0
while current_samples < self.target_samples:
f_id = ordered_files[file_idx % len(ordered_files)]
extended_files.append(f_id)
current_samples += self.batch_files_info[f_id]["samples"]
file_idx += 1
yielded_samples = 0
global_file_start_sample = 0
sequence_columns = columns_from_slice(
slice(0, self.folder_layout.window_length),
self.folder_layout.window_length,
)
seq_buffer: Dict[str, torch.Tensor] = {}
tgt_buffer: Dict[str, torch.Tensor] = {}
meta_buffer: Dict[str, torch.Tensor] = {}
buffer_len = 0
for f_id in extended_files:
if yielded_samples >= worker_target_samples:
break
file_samples = self.batch_files_info[f_id]["samples"]
file_start = global_file_start_sample
file_end = global_file_start_sample + file_samples
global_file_start_sample += file_samples
if file_end <= worker_start_sample or file_start >= worker_end_sample:
continue
file_path = os.path.join(self.data_dir, self.batch_files_info[f_id]["path"])
df = pl.read_parquet(file_path)
left_pad_lengths = get_left_pad_lengths_from_preprocessed_data(df)
position_rows = (
df.group_by(["sequenceId", "subsequenceId"])
.agg(
pl.col("startItemPosition").first(),
pl.col("splitStartItemPosition").first(),
)
.sort(["sequenceId", "subsequenceId"])
)
sample_index = self.sampling_plan.build_index(
left_pad_lengths,
target_valid_from_offsets(
left_pad_lengths,
torch.tensor(
position_rows["startItemPosition"].to_numpy(),
dtype=torch.int64,
),
torch.tensor(
position_rows["splitStartItemPosition"].to_numpy(),
dtype=torch.int64,
),
),
)
if len(sample_index) != file_samples:
raise RuntimeError(
f"Expanded sample count mismatch for {file_path}: "
f"metadata={file_samples}, loaded={len(sample_index)}."
)
indices = SampleOrderPlan.build(
file_samples,
logical_sample_positions(
curriculum_positions_from_parquet(self.config, df, file_path),
sample_index,
),
).indices_for_epoch(
seed=self.config.seed + f_id + rank,
epoch=epoch,
shuffle=self.shuffle,
)
worker_file_start_idx = max(0, worker_start_sample - file_start)
worker_file_end_idx = min(file_samples, worker_end_sample - file_start)
worker_indices = indices[worker_file_start_idx:worker_file_end_idx]
logical_positions = torch.arange(
file_start + worker_file_start_idx,
file_start + worker_file_end_idx,
dtype=torch.int64,
)
sample_is_real = logical_positions < rank_real_samples
num_new_samples = len(worker_indices)
if num_new_samples == 0:
del df
continue
feature_partitions = {
frame.item(0, "inputCol"): frame
for frame in df.partition_by("inputCol")
}
start_item_positions = None
split_start_item_positions = None
split_end_item_positions = None
if {
"splitStartItemPosition",
"splitEndItemPosition",
} <= set(df.columns):
positions = (
df.group_by(["sequenceId", "subsequenceId"])
.agg(
pl.col("startItemPosition").first(),
pl.col("splitStartItemPosition").first(),
pl.col("splitEndItemPosition").first(),
)
.sort(["sequenceId", "subsequenceId"])
)
start_item_positions = torch.tensor(
positions["startItemPosition"].to_numpy(), dtype=torch.int64
)
split_start_item_positions = torch.tensor(
positions["splitStartItemPosition"].to_numpy(), dtype=torch.int64
)
split_end_item_positions = torch.tensor(
positions["splitEndItemPosition"].to_numpy(), dtype=torch.int64
)
validate_split_bounds_available(
start_item_positions,
split_start_item_positions,
split_end_item_positions,
file_path,
)
stored_sequences = {}
for col_name in set(self.config.input_columns + self.config.target_columns):
if col_name in feature_partitions:
stored_sequences[col_name] = torch.tensor(
feature_partitions[col_name]
.sort(["sequenceId", "subsequenceId"])
.select(sequence_columns)
.to_numpy(),
dtype=self.column_torch_types[col_name],
)
else:
raise ValueError(f"Column not found in input data: {col_name}")
new_batch = build_window_batch(
stored_sequences,
self.config.input_columns,
self.config.target_columns,
sample_index,
worker_indices,
sample_is_real,
start_item_positions=start_item_positions,
split_start_item_positions=split_start_item_positions,
split_end_item_positions=split_end_item_positions,
)
new_seq = new_batch.inputs
new_tgt = new_batch.targets
new_meta = new_batch.metadata
del df
if buffer_len == 0:
seq_buffer = new_seq
tgt_buffer = new_tgt
meta_buffer = new_meta
else:
seq_buffer = {
k: torch.cat([seq_buffer[k], new_seq[k]], dim=0) for k in seq_buffer
}
tgt_buffer = {
k: torch.cat([tgt_buffer[k], new_tgt[k]], dim=0) for k in tgt_buffer
}
if set(meta_buffer) != set(new_meta):
raise RuntimeError(
"Inconsistent leftPadLength metadata across Parquet chunks."
)
meta_buffer = {
k: torch.cat([meta_buffer[k], new_meta[k]], dim=0)
for k in meta_buffer
}
buffer_len += num_new_samples
while buffer_len >= self.batch_size:
if yielded_samples >= worker_target_samples:
break
batch_seq = {k: v[: self.batch_size] for k, v in seq_buffer.items()}
batch_tgt = {k: v[: self.batch_size] for k, v in tgt_buffer.items()}
batch_meta = {k: v[: self.batch_size] for k, v in meta_buffer.items()}
yield SequifierBatch(
inputs=batch_seq,
targets=batch_tgt,
metadata=batch_meta,
)
yielded_samples += self.batch_size
seq_buffer = {k: v[self.batch_size :] for k, v in seq_buffer.items()}
tgt_buffer = {k: v[self.batch_size :] for k, v in tgt_buffer.items()}
meta_buffer = {k: v[self.batch_size :] for k, v in meta_buffer.items()}
buffer_len -= self.batch_size
if buffer_len > 0 and yielded_samples < worker_target_samples:
remaining_needed = worker_target_samples - yielded_samples
final_yield_size = min(buffer_len, remaining_needed)
batch_seq = {k: v[:final_yield_size] for k, v in seq_buffer.items()}
batch_tgt = {k: v[:final_yield_size] for k, v in tgt_buffer.items()}
batch_meta = {k: v[:final_yield_size] for k, v in meta_buffer.items()}
yield SequifierBatch(
inputs=batch_seq,
targets=batch_tgt,
metadata=batch_meta,
)