From e8d327adb13f6e906884bd54a71ca6f7779520fd Mon Sep 17 00:00:00 2001 From: Carson Rodrigues Date: Thu, 2 Jul 2026 22:05:54 +0530 Subject: [PATCH] Fix MetadataConfigs dropping parquet shards for non-consecutive configs _from_exported_parquet_files_and_dataset_infos() groups the exported parquet files with itertools.groupby, which only groups *consecutive* equal keys. When a config (or split) appears again after another config in the exported list, groupby yields it as multiple groups; because the result is built as a dict keyed by config name, the later group overwrites the earlier one and its shard URLs are silently lost. Sort the exported files by (config, split) before grouping so all rows for a config/split are consecutive. The sort is stable, so shard order within a split is preserved. Adds a regression test. Fixes #8269 --- src/datasets/utils/metadata.py | 6 ++++++ tests/test_metadata_util.py | 24 ++++++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/src/datasets/utils/metadata.py b/src/datasets/utils/metadata.py index cceecf35d63..d341109468a 100644 --- a/src/datasets/utils/metadata.py +++ b/src/datasets/utils/metadata.py @@ -106,6 +106,12 @@ def _from_exported_parquet_files_and_dataset_infos( exported_parquet_files: list[dict[str, Any]], dataset_infos: DatasetInfosDict, ) -> "MetadataConfigs": + # itertools.groupby only groups *consecutive* equal keys, so sort by + # (config, split) first. Otherwise, if the exported list has the same + # config (or split) in non-consecutive positions, groupby yields it as + # multiple groups and the earlier shard URLs are silently dropped. The + # sort is stable, so shard order within a split is preserved. + exported_parquet_files = sorted(exported_parquet_files, key=itemgetter("config", "split")) metadata_configs = { config_name: { "data_files": [ diff --git a/tests/test_metadata_util.py b/tests/test_metadata_util.py index cf9111fa6d9..9aa3664f83a 100644 --- a/tests/test_metadata_util.py +++ b/tests/test_metadata_util.py @@ -356,3 +356,27 @@ def test_split_order_in_metadata_configs_from_exported_parquet_files_and_dataset ) split_names = [data_file["split"] for data_file in metadata_configs["default"]["data_files"]] assert split_names == ["train", "validation", "test"] + + +def test_from_exported_parquet_files_keeps_all_shards_when_configs_non_consecutive(): + # Regression for #8269: itertools.groupby only groups *consecutive* keys, so + # a config that appears non-consecutively in the exported list must not lose + # its earlier shard URLs. + base = "https://huggingface.co/datasets/ds/resolve/refs%2Fconvert%2Fparquet" + exported_parquet_files = [ + {"dataset": "ds", "config": "default", "split": "train", + "url": f"{base}/default/train/0000.parquet", "filename": "0000.parquet", "size": 1}, + {"dataset": "ds", "config": "other", "split": "train", + "url": f"{base}/other/train/0000.parquet", "filename": "0000.parquet", "size": 1}, + {"dataset": "ds", "config": "default", "split": "train", + "url": f"{base}/default/train/0001.parquet", "filename": "0001.parquet", "size": 1}, + ] + metadata_configs = MetadataConfigs._from_exported_parquet_files_and_dataset_infos( + "123", exported_parquet_files, {} + ) + assert "other" in metadata_configs + default_paths = metadata_configs["default"]["data_files"][0]["path"] + # both shards are kept (not just the last one), with the commit hash substituted + assert len(default_paths) == 2 + assert default_paths == [f"{base.replace('refs%2Fconvert%2Fparquet', '123')}/default/train/0000.parquet", + f"{base.replace('refs%2Fconvert%2Fparquet', '123')}/default/train/0001.parquet"]