PySpark Hugging Face Data Source

Common errors

• ImportError: cannot import name 'list_datasets' from 'datasets'

  cause The `list_datasets` function, and similar Hub-interaction utilities, were moved from the `datasets` library to the `huggingface_hub` library. This is a common confusion when following older examples or tutorials.
  fix Change the import statement to `from huggingface_hub import list_datasets`.

• Not enough free disk space to download the file. The expected file size is: XXXX MB. The target location /root/.cache/huggingface/hub only has XXX MB free disk space.

  cause Hugging Face libraries default to caching downloaded models and datasets in a location (often `/root/.cache`) that might have limited storage, especially in containerized or shared environments.
  fix Before importing Hugging Face libraries, set the `HF_HUB_CACHE` environment variable to a directory with ample free space. Example: `import os; os.environ['HF_HUB_CACHE'] = '/path/to/your/large_disk_storage'`.

• org.apache.spark.api.python.PythonException: Traceback (most recent call last): ... WARN util.NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable ... WARN shortcircuit.DomainSocketFactory: The short-circuit local reads feature cannot be used because libhadoop cannot be loaded. ... Lost task 0.0 in stage ... (TID ...) ...

  cause This error typically occurs when PySpark tasks, particularly those involving `datasets.from_spark()` or similar operations that interact with distributed file systems like HDFS (via PyArrow), cannot properly access native Hadoop libraries or configured storage in a YARN or similar cluster environment.
  fix Ensure that your Spark cluster's `HADOOP_CONF_DIR` is correctly set, and that `libhadoop.so` (and other native libraries) are available on the `LD_LIBRARY_PATH` of the Spark executor nodes. Verify that PyArrow and Spark can connect to HDFS outside of `pyspark-huggingface` context. Check for Kerberos authentication issues if applicable.

Warnings

• gotcha When using `pyspark-huggingface` with PySpark 3.x (versions 3.3, 3.4, 3.5), you *must* explicitly `import pyspark_huggingface` in your code to enable the 'huggingface' data source format. This is not needed for PySpark 4+ as it's auto-imported.
  Fix: Add `import pyspark_huggingface` at the beginning of your Spark application when using PySpark 3.x.
• gotcha To read private/gated Hugging Face datasets or to write Spark DataFrames to the Hugging Face Hub, you need to authenticate. Not providing a valid token will lead to access errors.
  Fix: Authenticate using `huggingface-cli login` in your environment or pass your Hugging Face token directly via the `.option("token", "hf_xxxx")` to the `spark.read` or `spark.write` calls, or set the `HF_TOKEN` environment variable.
• gotcha For optimal performance, especially with large Parquet datasets, apply filters and select columns during the `spark.read.format('huggingface').option(...)` stage. This leverages Parquet metadata to skip unnecessary data, reducing I/O and processing.
  Fix: Use `.option("filters", '[("column_name", ">", value)]')` and `.option("columns", '["col1", "col2"]')` when loading the dataset. Example: `.option("filters", '[("language_score", ">", 0.99)]').option("columns", '["text", "language_score"]')`.
• gotcha Downloading large models or datasets can quickly exhaust disk space if the default Hugging Face cache directory (often in `/root/.cache/huggingface/hub`) is on a small root partition.
  Fix: Set the `HF_HUB_CACHE` environment variable to a path with sufficient available disk space *before* importing any `transformers` or `datasets` libraries. For example: `os.environ['HF_HUB_CACHE'] = "/path/to/large/storage"`.
• gotcha When uploading very large datasets with many shards to the Hugging Face Hub, you might encounter `HfHubHTTPError: 429 Client Error: Too Many Requests` due to hourly quotas.
  Fix: Ensure your `datasets` library (a common underlying dependency) is updated to version 2.15.0 or later, as this version includes improvements to handle such rate limits. `pip install --upgrade datasets`.

Install

• pip install pyspark-huggingface Install stable version

Imports

• pyspark_huggingface
  wrong:

  No direct classes are typically imported for DataSource usage.

  correct:

  import pyspark_huggingface

  For PySpark 3.x, an explicit import of `pyspark_huggingface` is required to enable the 'huggingface' data source format. For PySpark 4+, this is automatically handled upon installation. Interaction is primarily via `spark.read.format("huggingface")` and `spark.write.format("huggingface")`.

Quickstart

This quickstart initializes a Spark session and demonstrates how to read a public Hugging Face dataset (stanfordnlp/imdb) into a PySpark DataFrame. It also shows where to configure authentication for private datasets or write operations. For Spark 3.x users, an explicit `import pyspark_huggingface` might be necessary.

from pyspark.sql import SparkSession
import os

# Initialize Spark Session
spark = SparkSession.builder \
    .appName("HuggingFaceSpark") \
    .getOrCreate()

# For Spark 3.x, explicitly import to enable the data source:
# import pyspark_huggingface 

# Read a public dataset from Hugging Face
# Replace 'hf_token_xxxx' with your actual token if accessing private/gated datasets
hf_token = os.environ.get('HF_TOKEN', '') # Use a real token for private datasets

print("Reading stanfordnlp/imdb dataset...")
df = spark.read \
    .format("huggingface") \
    .option("token", hf_token) \
    .load("stanfordnlp/imdb")

print("Schema:")
df.printSchema()

print("First 5 rows:")
df.show(5, truncate=False)

# Example of saving a DataFrame to Hugging Face (requires write token and dataset repo name)
# try:
#     print("Saving a sample DataFrame to Hugging Face...")
#     sample_data = [("hello", "world"), ("spark", "huggingface")]
#     sample_df = spark.createDataFrame(sample_data, ["col1", "col2"])
#     sample_df.write \
#         .format("huggingface") \
#         .option("token", hf_token) \
#         .mode("overwrite") \
#         .save("your_username/your_dataset_name")
#     print("DataFrame saved successfully.")
# except Exception as e:
#     print(f"Could not save DataFrame to Hugging Face: {e}")

spark.stop()