Improve dataset .skip() speed in streaming mode

Feature request

Add extra information to the dataset_infos.json file to include the number of samples/examples in each shard, for example in a new field num_examples alongside num_bytes. The .skip() function could use this information to ignore the download of a shard when in streaming mode, which AFAICT it should speed up the skipping process.

Motivation

When resuming from a checkpoint after a crashed run, using dataset.skip() is very convenient to recover the exact state of the data and to not train again over the same examples (assuming same seed, no shuffling). However, I have noticed that for audio datasets in streaming mode this is very costly in terms of time, as shards need to be downloaded every time before skipping the right number of examples.

Your contribution

I took a look already at the code, but it seems a change like this is way deeper than I am able to manage, as it touches the library in several parts. I could give it a try but might need some guidance on the internals.

@severo

Hi! I agree skip can be inefficient to use in the current state.

To make it fast, we could use "statistics" stored in Parquet metadata and read only the chunks needed to form a dataset.

And thanks to the "datasets-server" project, which aims to store the Parquet versions of the Hub datasets (only the smaller datasets are covered currently), this solution can also be applied to datasets stored in formats other than Parquet. (cc @severo)

@mariosasko

@mariosasko do the current parquet files created by the datasets-server already have the required "statistics"? If not, please open an issue on https://github.com/huggingface/datasets-server with some details to make sure we implement it.

Yes, nothing has to be changed on the datasets-server side. What I mean by "statistics" is that we can use the "row_group" metadata embedded in a Parquet file (by default) to fetch the requested rows more efficiently.

Glad to see the feature could be of interest.

I'm sure there are many possible ways to implement this feature. I don't know enough about the datasets-server, but I guess that it is not instantaneous, in the sense that user-owned private datasets might need hours or days until they are ported to the datasets-server (if at all), which could be cumbersome. Having optionally that information in the dataset_infos.json file would make it easier for users to control the skip process a bit.

re: statistics:

>>> import pyarrow.parquet as pq
>>> import hffs
>>> fs = hffs.HfFileSystem("glue", repo_type="dataset", revision="refs/convert/parquet")
>>> metadata = pq.read_metadata("ax/glue-test.parquet", filesystem=fs)
>>> metadata
<pyarrow._parquet.FileMetaData object at 0x7f4537cec400>
  created_by: parquet-cpp-arrow version 7.0.0
  num_columns: 4
  num_rows: 1104
  num_row_groups: 2
  format_version: 1.0
  serialized_size: 2902
>>> metadata.row_group(0)
<pyarrow._parquet.RowGroupMetaData object at 0x7f45564bcbd0>
  num_columns: 4
  num_rows: 1000
  total_byte_size: 164474
>>> metadata.row_group(1)
<pyarrow._parquet.RowGroupMetaData object at 0x7f455005c400>
  num_columns: 4
  num_rows: 104
  total_byte_size: 13064

user-owned private datasets might need hours or days until they are ported to the datasets-server (if at all)

private datasets are not supported yet (huggingface/dataset-viewer#39)

@versae

@versae Dataset.push_to_hub writes shards in Parquet, so this solution would also work for such datasets (immediately after the push).

@mariosasko

@mariosasko that is right. However, there are still a good amount of datasets for which the shards are created manually. In our very specific case, we create medium-sized datasets (rarely over 100-200GB) of both text and audio, we prepare the shards by hand and then upload then. It would be great to have immediate access to this download skipping feature for them too.

From looking at Arrow's source, it seems Parquet stores metadata at the end, which means one needs to iterate over a Parquet file's data before accessing its metadata. We could mimic Dask to address this "limitation" and write metadata in a _metadata/_common_metadata file in to_parquet/push_to_hub, which we could then use to optimize reads (if present). Plus, it's handy that PyArrow can also parse these metadata files.

So if Parquet metadata needs to be in its own file anyway, why not implement this skipping feature by storing the example counts per shard in dataset_infos.json? That would allow:

Support both private and public datasets
Immediate access to the feature upon uploading of shards
Use any dataset, not only those uploaded using .push_to_hub()

A proper Parquet metadata file could still be created and "overwrite" the dataset_infos.json info in the datasets-server.

Any updates on this?

@lhoestq

i benchmarked skip() on streamed local parquet to see where the cost is today. setup: 3 files, 300,000 rows, 10 row groups per file (13.6 MB each), datasets main at 9e7496a4, pyarrow 25.0.1. the last column is a pyarrow baseline that reads the footers and then only the row group containing row n.

`skip(n)`	rows decoded	MB read	footers + 1 row group
0	10,000	13.64	0.19
100,000	110,000	27.26	0.25
290,001	300,000	40.89	0.32
past the end	300,000	40.89	0.20

the returned rows were correct in every case, and default iteration and with_format("arrow") cost the same. skip() reads and decodes every row group before n, so the cost grows with n. locally that's only ~0.02s, but when streaming from the hub it's all download.

on the earlier concern that parquet metadata sits at the end of the file: with a seekable file, pq.read_metadata only reads the footer. reading all three footers here took about 0.2 MB. it holds up over hf:// too: for HuggingFaceFW/fineweb-edu sample/10BT/000_00000.parquet (2,153 MB, 726 row groups), pq.read_metadata through xopen fetched 3.09 MB in 2 requests (~1.1s), and the footer itself is 3.03 MB. one row group there is 2.7 MB compressed, so skipping into the middle of that file could cost ~6 MB instead of the ~1 GB before it. so i don't think a separate _metadata file is needed.

a narrow version that seems to fit the current code:

let ArrowExamplesIterable ask the builder for per-source row counts. parquet can get them from footers, like _generate_num_examples in the lance, vortex and iceberg builders.
in SkipExamplesIterable, drop whole sources (files, or row groups after reshard_data_sources) whose rows are all skipped, then skip the remainder inside the first kept source like today.
only for unfiltered, untransformed parquet. fall back to the current behaviour with filters, map/filter, shuffled sources, or a state_dict that already has progress.
keep state_dict["skipped"] consistent so resuming gives the same rows.

side note, not part of this: even skip(0) reads the whole first file (13.6 MB) before yielding, because pyarrow's default pre_buffer=True fetches every row group up front. with pre_buffer=False it's 4.1 MB.

@lhoestq would you be open to a pr along these lines? i'd start with tests that pin the current skip results (row group and file boundaries, skips inside a group, past the end) before changing anything.

@lhoestq

hi @lhoestq, would the narrow scope above work for a first PR? i'd skip complete parquet sources using footer counts, keep the existing path for the remainder and for filters, transforms, shuffles or resumed progress, and leave pre_buffer unchanged. i'd keep the empty-row-group fix in #8742 separate.

yes using _generate_num_examples sounds good ! and when it's not available then it can fallback to the current implem

Navigation Menu

Repository navigation

Improve dataset `.skip()` speed in streaming mode #5380

Feature request

Motivation

Your contribution

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Improve dataset .skip() speed in streaming mode #5380

Description

Feature request

Motivation

Your contribution

Activity

mariosasko commented on Dec 20, 2022

severo commented on Dec 20, 2022

mariosasko commented on Dec 21, 2022

versae commented on Dec 27, 2022

severo commented on Dec 27, 2022

severo commented on Dec 27, 2022

mariosasko commented on Jan 13, 2023

versae commented on Jan 17, 2023

mariosasko commented on Feb 24, 2023

versae commented on Mar 3, 2023

Ueaj-Kerman commented on Nov 10, 2025

cwklurks commented on Oct 8, 2026

cwklurks commented on Oct 12, 2026

lhoestq commented on Oct 13, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Improve dataset `.skip()` speed in streaming mode #5380