Implementation:Huggingface Datatrove HuggingFaceDatasetWriter

Knowledge Sources	Huggingface_Datatrove
Domains	Data Processing, HuggingFace Hub
Last Updated	2026-02-14 17:00 GMT

Overview

HuggingFaceDatasetWriter is a pipeline writer that saves documents as Parquet files to the HuggingFace Hub, handling repository creation, LFS pre-uploads, and commit creation with retry logic.

Description

The HuggingFaceDatasetWriter class extends ParquetWriter to provide seamless uploading of large datasets directly to the HuggingFace Hub during pipeline execution. It first writes Parquet files to a local working directory (either user-specified or a temporary directory), then uploads them via the HuggingFace Hub API using LFS pre-upload for efficient large file handling.

The writer manages the full lifecycle of a Hub upload: it creates the dataset repository if it does not exist (supporting private repositories), pre-uploads LFS files as they are written (triggered when files exceed the max file size and are switched), and creates a final commit containing all uploaded file operations. The commit creation includes exponential backoff retry logic (up to 12 retries) to handle race conditions when multiple workers attempt concurrent commits to the same repository.

Key features include configurable max_file_size (defaulting to approximately 4.5 GB to stay within Hub limits), optional cleanup of local files after upload, support for a specific revision (branch) for commits, and Parquet compression options (snappy, gzip, brotli, lz4, zstd). The class warns users if the deprecated `HF_HUB_ENABLE_HF_TRANSFER` environment variable is set, recommending the newer xet-based upload mechanism instead.

Usage

Use this writer when producing very large datasets that should be published directly to the HuggingFace Hub. For smaller datasets, consider using `push_to_hub` or an `hf://datasets/...` output path instead. This writer is optimized for scenarios where data is generated incrementally and needs to be uploaded efficiently in parallel.

Code Reference

Source Location

• Repository: Huggingface_Datatrove
• File: src/datatrove/pipeline/writers/huggingface.py
• Lines: 1-146

Signature

class HuggingFaceDatasetWriter(ParquetWriter):
    def __init__(
        self,
        dataset: str,
        private: bool = True,
        local_working_dir: DataFolderLike | None = None,
        output_filename: str = None,
        compression: Literal["snappy", "gzip", "brotli", "lz4", "zstd"] | None = "snappy",
        adapter: Callable = None,
        cleanup: bool = True,
        expand_metadata: bool = True,
        max_file_size: int = round(4.5 * 2**30),
        schema: Any = None,
        revision: str | None = None,
        save_media_bytes=False,
    ):

Import

from datatrove.pipeline.writers.huggingface import HuggingFaceDatasetWriter

I/O Contract

Inputs

Name	Type	Required	Description
dataset	str	Yes	HuggingFace dataset identifier in "namespace/repo_name" format
private	bool	No	Whether to create the repository as private (default: True)
local_working_dir	DataFolderLike or None	No	Local directory for staging files before upload (default: temporary directory)
output_filename	str	No	Filename template for output Parquet files (default: "data/${rank}.parquet")
compression	str or None	No	Parquet compression codec (default: "snappy")
adapter	Callable	No	Custom function to transform Document objects (default: None)
cleanup	bool	No	Whether to delete local files after upload (default: True)
expand_metadata	bool	No	Whether to expand metadata into separate columns (default: True)
max_file_size	int	No	Maximum file size before splitting, in bytes (default: ~4.5 GB)
schema	Any	No	Parquet schema to use (default: None)
revision	str or None	No	Git branch to commit to (default: None, uses "main")
save_media_bytes	bool	No	Whether to include media bytes in output (default: False)

Outputs

Name	Type	Description
Parquet files	HuggingFace Hub	Parquet files uploaded and committed to the specified HuggingFace dataset repository

Usage Examples

Basic Usage

from datatrove.pipeline.writers.huggingface import HuggingFaceDatasetWriter

writer = HuggingFaceDatasetWriter(
    dataset="my-org/my-dataset",
    private=True,
    compression="snappy",
    expand_metadata=True,
)

Related Pages

• Principle:Huggingface_Datatrove_HuggingFace_Hub_Writing