@@ -156,14 +70,15 @@ It is a gated repository, users have to accept the terms of use before accessing
-```
-pip install huggingface_hub
-```
+## Set up
+
+### Installation
-You also need to install `pyarrow` to read/write Parquet / JSON / CSV / etc. files using the filesystem API provided by `huggingFace_hub`:
+To be able to read and write to Hugging Face Datasets, you need to install the `pyspark_huggingface` library:
```
-pip install pyarrow
+pip install pyspark_huggingface
```
-## Authentication
+This will also install required dependencies like `huggingface_hub` for authentication, and `pyarrow` for reading and writing datasets.
+
+### Authentication
You need to authenticate to Hugging Face to read private/gated dataset repositories or to write to your dataset repositories.
@@ -28,127 +30,39 @@ You can use the CLI for example:
huggingface-cli login
```
-It's also possible to provide your Hugging Face token with the `HF_TOKEN` environment variable or passing the `storage_options` parameter to helper functions below:
+It's also possible to provide your Hugging Face token with the `HF_TOKEN` environment variable or passing the `token` option to the reader.
+For more details about authentication, check out [this guide](https://huggingface.co/docs/huggingface_hub/quick-start#authentication).
+
+### Enable the "huggingface" Data Source
+
+PySpark 4 came with a new Data Source API which allows to use datasets from custom sources.
+If `pyspark_huggingface` is installed, PySpark auto-imports it and enables the "huggingface" Data Source.
+
+The library also backports the Data Source API for the "huggingface" Data Source for PySpark 3.5, 3.4 and 3.3.
+However in this case `pyspark_huggingface` should be imported explicitly to activate the backport and enable the "huggingface" Data Dource:
```python
-storage_options = {"token": "hf_xxx"}
+>>> import pyspark_huggingface
+huggingface datasource enabled for pyspark 3.x.x (backport from pyspark 4)
```
-For more details about authentication, check out [this guide](https://huggingface.co/docs/huggingface_hub/quick-start#authentication).
-
## Read
-PySpark doesn't have an official support for Hugging Face paths, so we provide a helper function to read datasets in a distributed manner.
+The "huggingface" Data Source allows to read datasets from Hugging Face, using `pyarrow` under the hood to stream Arrow data.
+This is compatible with all the dataset in [supported format](https://huggingface.co/docs/hub/datasets-adding#file-formats) on Hugging Face, like Parquet datasets.
-For example you can read Parquet files from Hugging Face in an optimized way using PyArrow by defining this `read_parquet` helper function:
+For example here is how to load the [stanfordnlp/imdb](https://huggingface.co/stanfordnlp/imdb) dataset:
```python
-from functools import partial
-from typing import Iterator, Optional, Union
-
-import pyarrow as pa
-import pyarrow.parquet as pq
-from huggingface_hub import HfFileSystem
-from pyspark.sql.dataframe import DataFrame
-from pyspark.sql.pandas.types import from_arrow_schema
-
-
-def _read(iterator: Iterator[pa.RecordBatch], columns: Optional[list[str]], filters: Optional[Union[list[tuple], list[list[tuple]]]], **kwargs) -> Iterator[pa.RecordBatch]:
- for batch in iterator:
- paths = batch[0].to_pylist()
- ds = pq.ParquetDataset(paths, **kwargs)
- yield from ds._dataset.to_batches(columns=columns, filter=pq.filters_to_expression(filters) if filters else None)
-
-
-def read_parquet(
- path: str,
- columns: Optional[list[str]] = None,
- filters: Optional[Union[list[tuple], list[list[tuple]]]] = None,
- **kwargs,
-) -> DataFrame:
- """
- Loads Parquet files from Hugging Face using PyArrow, returning a PySPark `DataFrame`.
-
- It reads Parquet files in a distributed manner.
-
- Access private or gated repositories using `huggingface-cli login` or passing a token
- using the `storage_options` argument: `storage_options={"token": "hf_xxx"}`
-
- Parameters
- ----------
- path : str
- Path to the file. Prefix with a protocol like `hf://` to read from Hugging Face.
- You can read from multiple files if you pass a globstring.
- columns : list, default None
- If not None, only these columns will be read from the file.
- filters : List[Tuple] or List[List[Tuple]], default None
- To filter out data.
- Filter syntax: [[(column, op, val), ...],...]
- where op is [==, =, >, >=, <, <=, !=, in, not in]
- The innermost tuples are transposed into a set of filters applied
- through an `AND` operation.
- The outer list combines these sets of filters through an `OR`
- operation.
- A single list of tuples can also be used, meaning that no `OR`
- operation between set of filters is to be conducted.
-
- **kwargs
- Any additional kwargs are passed to pyarrow.parquet.ParquetDataset.
-
- Returns
- -------
- DataFrame
- DataFrame based on parquet file.
-
- Examples
- --------
- >>> path = "hf://datasets/username/dataset/data.parquet"
- >>> pd.DataFrame({"foo": range(5), "bar": range(5, 10)}).to_parquet(path)
- >>> read_parquet(path).show()
- +---+---+
- |foo|bar|
- +---+---+
- | 0| 5|
- | 1| 6|
- | 2| 7|
- | 3| 8|
- | 4| 9|
- +---+---+
- >>> read_parquet(path, columns=["bar"]).show()
- +---+
- |bar|
- +---+
- | 5|
- | 6|
- | 7|
- | 8|
- | 9|
- +---+
- >>> sel = [("foo", ">", 2)]
- >>> read_parquet(path, filters=sel).show()
- +---+---+
- |foo|bar|
- +---+---+
- | 3| 8|
- | 4| 9|
- +---+---+
- """
- filesystem: HfFileSystem = kwargs.pop("filesystem") if "filesystem" in kwargs else HfFileSystem(**kwargs.pop("storage_options", {}))
- paths = filesystem.glob(path)
- if not paths:
- raise FileNotFoundError(f"Counldn't find any file at {path}")
- rdd = spark.sparkContext.parallelize([{"path": path} for path in paths], len(paths))
- df = spark.createDataFrame(rdd)
- arrow_schema = pq.read_schema(filesystem.open(paths[0]))
- schema = pa.schema([field for field in arrow_schema if (columns is None or field.name in columns)], metadata=arrow_schema.metadata)
- return df.mapInArrow(
- partial(_read, columns=columns, filters=filters, filesystem=filesystem, schema=arrow_schema, **kwargs),
- from_arrow_schema(schema),
- )
+>>> import pyspark_huggingface
+>>> from pyspark.sql import SparkSession
+>>> spark = SparkSession.builder.appName("demo").getOrCreate()
+>>> df = spark.read.format("huggingface").load("stanfordnlp/imdb")
```
-Here is how we can use this on the [BAAI/Infinity-Instruct](https://huggingface.co/datasets/BAAI/Infinity-Instruct) dataset.
+Here is another example with the [BAAI/Infinity-Instruct](https://huggingface.co/datasets/BAAI/Infinity-Instruct) dataset.
It is a gated repository, users have to accept the terms of use before accessing it.
+It also has multiple subsets, namely, "3M" and "7M". So we need to specify which one to load.
+Contributions are welcome to add more options here, in particular `subset` and `split`.