Title: Question: Is there a faster way to push_to_hub for large image datasets?
State: open

# Question: Is there a faster way to `push_to_hub` for large image datasets? Or could this approach be integrated?

Hi! I frequently work with large image datasets (100k-300k+ samples) and found that `dataset.push_to_hub()` can be quite slow.

cc @lhoestq - would love your thoughts on this!

I experimented with an alternative approach using `upload_large_folder` and parallel parquet conversion that gave me significant speedups:

**My benchmark (15,000 image samples, ~16GB total):**
- `dataset.push_to_hub()`: **~15 minutes**
- Alternative approach below: **~2-3 minutes**

I wanted to ask:
1. **Is there an existing way to achieve this speed** that I'm missing in the current API?
2. **If not, would something like this be worth integrating** into the library?

## The Approach

The key differences from standard `push_to_hub()`:
1. **Parallel parquet shard conversion** using ThreadPoolExecutor
2. **Using `upload_large_folder`** instead of regular upload (multi-threaded, resumable)
3. **Enabling `HF_XET_HIGH_PERFORMANCE=1`** for chunk-level deduplication

## Full Working Code

Here's the script I've been using:

```python
#!/usr/bin/env python3
"""
Fast HuggingFace Dataset Upload Script

Converts any dataset to parquet shards and uploads using upload_large_folder
for maximum speed with parallel workers and hf_xet chunk deduplication.

Usage:
    python push_to_hub_fast.py --input ./my_dataset_arrow --repo "username/my-dataset"
    python push_to_hub_fast.py --input ./data --repo "user/repo" --workers 32 --shard-size 400
"""

import argparse
import os
import shutil
import tempfile
import multiprocessing
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Optional

import yaml


def get_optimal_workers() -> int:
    """Detect optimal number of workers based on CPU cores."""
    cpu_count = multiprocessing.cpu_count()
    return max(1, cpu_count - 1)


def generate_dataset_card(dataset, num_samples: int, repo_name: str) -> str:
    """Generate a HuggingFace dataset card README with proper metadata."""
    from datasets import Value, Sequence, Image

    features = dataset.features

    def feature_to_dict(feat):
        feat_type = type(feat).__name__
        if isinstance(feat, Value):
            return {"dtype": feat.dtype}
        elif isinstance(feat, Image):
            return {"dtype": "image"}
        elif isinstance(feat, Sequence):
            if hasattr(feat.feature, "dtype"):
                return {"list": feat.feature.dtype}
            elif hasattr(feat.feature, "__iter__"):
                inner = []
                for k, v in feat.feature.items():
                    inner.append({"name": k, **feature_to_dict(v)})
                return {"list": inner}
            else:
                return {"list": feature_to_dict(feat.feature)}
        elif hasattr(feat, "__iter__") and not isinstance(feat, str):
            inner = []
            for k, v in feat.items():
                inner.append({"name": k, **feature_to_dict(v)})
            return {"list": inner}
        else:
            return {"dtype": str(feat_type).lower()}

    features_list = []
    for name, feat in features.items():
        feat_dict = {"name": name}
        feat_dict.update(feature_to_dict(feat))
        features_list.append(feat_dict)

    dataset_info = {
        "dataset_info": {
            "features": features_list,
            "splits": [{"name": "train", "num_examples": num_samples}],
        },
        "configs": [
            {
                "config_name": "default",
                "data_files": [{"split": "train", "path": "data/train-*"}],
            }
        ],
    }

    yaml_content = yaml.dump(
        dataset_info, default_flow_style=False, sort_keys=False, allow_unicode=True
    )

    # Build README string with YAML frontmatter and usage instructions
    readme = "---\n" + yaml_content + "---\n\n"
    readme += "# " + repo_name.split("/")[-1] + "\n\n"
    readme += "## Usage\n\n```python\nfrom datasets import load_dataset\n"
    readme += "dataset = load_dataset(\"" + repo_name + "\")\n```\n\n"
    readme += "- **Samples**: " + f"{num_samples:,}" + "\n"
    readme += "- **Features**: " + ", ".join(f"`{name}`" for name in features.keys())
    return readme


def push_to_hub_fast(
    input_path: str,
    repo_id: str,
    num_workers: Optional[int] = None,
    max_shard_size_mb: int = 450,
    split: str = "train",
    private: bool = False,
    token: Optional[str] = None,
):
    """
    Fast upload dataset to HuggingFace Hub.

    Args:
        input_path: Path to dataset (local arrow/disk) or HuggingFace dataset name
        repo_id: Target repository (e.g., "username/dataset-name")
        num_workers: Number of parallel workers (auto-detected if None)
        max_shard_size_mb: Target shard size in MB (default 450 for <500MB viewer limit)
        split: Dataset split to upload (default "train")
        private: Whether to make repo private
        token: HuggingFace token (uses cached login if None)
    """
    os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"

    from datasets import load_from_disk, load_dataset, DatasetDict
    from huggingface_hub import HfApi, login

    if num_workers is None:
        num_workers = get_optimal_workers()

    print(f"Push to Hub Fast")
    print(f"=" * 50)
    print(f"Input: {input_path}")
    print(f"Target: {repo_id}")
    print(f"Workers: {num_workers}")
    print(f"=" * 50)

    # Login
    if token:
        login(token=token)
    else:
        token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
        if token:
            login(token=token)
    print("✓ Logged in to HuggingFace Hub")

    # Load dataset
    print(f"\nLoading dataset from {input_path}...")
    input_path_obj = Path(input_path)

    if input_path_obj.exists():
        try:
            loaded = load_from_disk(input_path)
        except Exception:
            loaded = load_dataset("imagefolder", data_dir=input_path)
    else:
        loaded = load_dataset(input_path)

    if isinstance(loaded, DatasetDict):
        if split in loaded:
            dataset = loaded[split]
        else:
            split = list(loaded.keys())[0]
            dataset = loaded[split]
    else:
        dataset = loaded

    num_samples = len(dataset)
    print(f"✓ Dataset loaded: {num_samples:,} samples")

    # Create temp directory
    temp_dir = tempfile.mkdtemp(prefix="hf_upload_")
    upload_dir = Path(temp_dir)
    data_dir = upload_dir / "data"
    data_dir.mkdir(parents=True, exist_ok=True)

    try:
        # Generate README
        print("\nGenerating dataset card...")
        readme_content = generate_dataset_card(dataset, num_samples, repo_id)
        with open(upload_dir / "README.md", "w") as f:
            f.write(readme_content)
        print("✓ README.md generated")

        # Calculate shards
        samples_per_shard = max(100, max_shard_size_mb)
        num_shards = max(1, (num_samples + samples_per_shard - 1) // samples_per_shard)

        print(f"\nConverting to {num_shards} parquet shards...")
        print(f"Using {num_workers} parallel workers...")

        def convert_shard(shard_idx):
            shard = dataset.shard(num_shards=num_shards, index=shard_idx)
            shard_path = data_dir / f"train-{shard_idx:05d}-of-{num_shards:05d}.parquet"
            shard.to_parquet(str(shard_path))
            return shard_idx, len(shard)

        # Parallel conversion
        completed = 0
        with ThreadPoolExecutor(max_workers=num_workers) as executor:
            futures = {executor.submit(convert_shard, i): i for i in range(num_shards)}
            for future in as_completed(futures):
                shard_idx, shard_len = future.result()
                completed += 1
                print(f"  [{completed}/{num_shards}] Shard {shard_idx}: {shard_len} samples")

        parquet_files = list(data_dir.glob("*.parquet"))
        total_size = sum(f.stat().st_size for f in parquet_files)
        print(f"✓ Created {len(parquet_files)} shards ({total_size / 1e9:.2f} GB)")

        # Create repo & upload
        api = HfApi()
        api.create_repo(repo_id=repo_id, repo_type="dataset", private=private, exist_ok=True)
        print(f"\n✓ Repository ready: {repo_id}")

        print(f"\nUploading with upload_large_folder ({num_workers} workers)...")
        print("Using hf_xet HIGH_PERFORMANCE mode")

        api.upload_large_folder(
            folder_path=str(upload_dir),
            repo_id=repo_id,
            repo_type="dataset",
            num_workers=num_workers,
        )

        print(f"\n✓ Uploaded to: https://huggingface.co/datasets/{repo_id}")

        return {"status": "success", "repo_id": repo_id, "samples": num_samples}

    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)
        print("✓ Cleaned up temp files")


def main():
    parser = argparse.ArgumentParser(description="Fast upload dataset to HuggingFace Hub")
    parser.add_argument("--input", "-i", required=True, help="Input path or HF dataset name")
    parser.add_argument("--repo", "-r", required=True, help="Target HF repo")
    parser.add_argument("--workers", "-w", type=int, default=None, help="Parallel workers")
    parser.add_argument("--shard-size", "-s", type=int, default=450, help="Max shard size MB")
    parser.add_argument("--split", default="train", help="Dataset split")
    parser.add_argument("--private", action="store_true", help="Make repo private")
    parser.add_argument("--token", default=None, help="HF token")

    args = parser.parse_args()
    push_to_hub_fast(
        input_path=args.input,
        repo_id=args.repo,
        num_workers=args.workers,
        max_shard_size_mb=args.shard_size,
        split=args.split,
        private=args.private,
        token=args.token,
    )


if __name__ == "__main__":
    main()
```

## Why This Is Faster

1. **Parallel parquet conversion**: Instead of sequential shard creation, uses `ThreadPoolExecutor` to convert multiple shards simultaneously

2. **`upload_large_folder` benefits** ([docs](https://huggingface.co/docs/huggingface_hub/en/guides/upload#upload-a-large-folder)):
   - Multi-threaded uploads with `num_workers`
   - Resumable - caches progress locally
   - Resilient - auto-retries on transient errors

3. **hf_xet chunk deduplication**: With `HF_XET_HIGH_PERFORMANCE=1`, uploads reached **3.45 GB/s** vs ~25-40 MB/s with standard upload

## Test Dataset

I tested this with an image dataset containing:
- **15,000 samples** (~16GB total)
- Each sample has: image, layout data, markdown, HTML, VQA pairs
- Average ~1MB per sample

## Questions

1. Is there something in the current API that achieves similar performance that I'm missing?
2. If not, would this be a useful addition to the library? Happy to contribute a PR if so.

Thanks for the amazing library!


Comments:
This is a really interesting approach, especially combining parallel parquet
conversion with upload_large_folder and hf_xet.

One question / observation:
The shard calculation currently uses max_shard_size_mb as a proxy for
samples_per_shard. Since sample size can vary a lot across datasets,
would it make sense to estimate shard size based on byte size instead?

Also curious whether ThreadPoolExecutor was chosen intentionally
(over multiprocessing) due to to_parquet being more IO-bound.

If this direction makes sense, I’d be happy to test on a large dataset
and share numbers.


The approach makes a lot of sense :)

* parallel parquet conversions with threads for speed
* writing parquet files to disk prior to upload instead of in RAM to save some RAM
* uploading in parallel with upload_large_folder for speed


Note that there is a helper for estimating the size of the dataset: `ds._estimate_nbytes()`

This is the kind of improvements that would be helpful to have in `push_to_hub` ! Maybe one point to take into account is some datasets can be larger than disk, so we would have to go one level lower than upload_large_folder:

1. convert to Parquet on disk
2. pre-upload the file and delete it locally
3. then at the end commit the uploaded files

Contributions are welcome if someone would like to improve `push_to_hub` in `datasets` btw. 

A summary of what we could achieve with some inspiration from your approach:

* parallel parquet conversions with threads for speed
* writing parquet files to disk prior to upload instead of in RAM to save some RAM
* upload local parquet files once ready with preupload_lfs_files() and delete them to not fill the disk
* commit at the end all the uploaded files

Thanks @lhoestq! Valid point about datasets larger than disk.

Would you be open to adding `upload_large_folder` as an **opt-in** parameter? The key benefit is **resumability** - if an upload fails at 80%, it picks up where it left off. For multi-hour uploads of large image datasets, this is really valuable.

```python
dataset.push_to_hub("user/repo", use_large_folder=True)
```

Users opting in would accept the disk space trade-off for resumability + atomic commits. Default behavior stays unchanged, so larger-than-disk datasets still work fine with the current approach.

Thanks for the `_estimate_nbytes()` pointer - will use that for proper shard sizing. Happy to open a PR if this direction works!

Great questions @k281484-ctrl!

**Shard sizing**: You're right - the sample-based approach in my script was a quick hack. @lhoestq pointed out `ds._estimate_nbytes()` which is exactly what we need for proper byte-based shard calculation. Will use that in the implementation.

**ThreadPoolExecutor**: Yes, intentional choice. `to_parquet()` is I/O-bound (writing to disk), so threads work well and avoid the overhead of spawning processes + pickling data across process boundaries.

Would love to see your benchmark numbers if you get a chance to test!