fix(streaming): drop undeclared parquet columns that break batch collation

The data_files_root/bucket path reads an unversioned source (e.g. `main`), which can
carry extra annotation columns not in the dataset's feature contract — notably
`language_events`, a variable-length list (length 0..N per frame). Passed through to the
sample, these break default DataLoader collation ("each element in list of batch should
be of equal size"), which is why bucket jobs failed while the hub path (pinned to the
clean v3.0 revision) succeeded.

Drop any hf_dataset column not in meta.features after load. No-op on a clean revision;
removes language_events/language_persistent on main. Verified by reproducing the bucket
code path locally via --data_files_root hf://datasets/<repo> (parquet builder + main
columns): now decodes and collates instead of raising.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Pepijn
2026-06-09 17:24:30 +02:00
parent 2d1c17d971
commit 42d4788e4a
@@ -372,6 +372,15 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
revision=self.revision,
)
# Drop any parquet columns not declared in the dataset's feature contract. Some revisions / sources
# (e.g. an unversioned bucket holding `main`) carry extra, possibly variable-length annotation
# columns such as `language_events`; left in, they leak into the sample and break default DataLoader
# collation across frames of differing length. On a clean revision this is a no-op.
known_columns = set(self.meta.features)
extra_columns = [c for c in (self.hf_dataset.column_names or []) if c not in known_columns]
if extra_columns:
self.hf_dataset = self.hf_dataset.remove_columns(extra_columns)
self.num_shards = min(self.hf_dataset.num_shards, max_num_shards)
@property