Repository: Lightning-AI/litData

Issue:
Add support for parquet files for storing the chunks
This would enable users to avoid converting their dataset if they already have their dataset as parquet folders. We would need to run an indexation but this isn't too painful.

Discussion:
@deependujha: My preliminary approach:

---
In OptimizeDataset:

Only if parquet files, Use pyarrow `read_table` function to load each parquet file one by one and then the writer will only write the number of files, and column types in `index.json` file, no chunk files created.

In reader,

All indices will remain as usual, only the reading at index `i` will be changed:
> df.slice(7, 1).to_pandas().to_dict() # parquet file 7th index value

If a parquet dataset has no index.json file, we can still call the helper function to generate index.json on the fly and then StreamingDataset takes control.

---

## Why no multithreading or multiprocessing while creating index.json file:

- Parquet files once loaded in memory are uncompressed and may exceed memory limit.

Or, we might take care of it in another PR.

---

What do you think @tchaton !@tchaton: Yes, that's what I had in mind. The main challenge will be to make the slicing and reading as fast as possible. Might be worth to use: https://github.com/pola-rs/polars@tchaton: The goal is to enable reading pyarrow HF datasets with LitData


Solution patch:
diff --git a/README.md b/README.md
index b8e5720..a369229 100644
--- a/README.md
+++ b/README.md
@@ -642,6 +642,37 @@ The `overwrite` mode will delete the existing data and start from fresh.
 
 </details>
 
+<details>
+  <summary> ✅ Index parquet datasets</summary>
+&nbsp;
+
+If your dataset is already in Parquet format, you can index it directly and use it with StreamingDataset & DataLoader.
+
+Assumption:
+Your dataset directory contains one or more Parquet files.
+
+```python
+import litdata as ld
+
+pq_data_uri = "gs://deep-litdata-parquet/my-parquet-data"
+
+ld.index_parquet_dataset(pq_data_uri)
+```
+
+When using a Streaming Dataset, ensure you use `ParquetLoader`:
+
+```python
+import litdata as ld
+from litdata.streaming.item_loader import ParquetLoader
+
+ds = ld.StreamingDataset('gs://deep-litdata-parquet/my-parquet-data', item_loader = ParquetLoader())
+
+for _ds in ds:
+    print(f"{_ds=}")
+```
+
+</details>
+
 <details>
   <summary> ✅ Use compression</summary>
 &nbsp;
diff --git a/requirements/extras.txt b/requirements/extras.txt
index f48910b..4cb819a 100644
--- a/requirements/extras.txt
+++ b/requirements/extras.txt
@@ -5,3 +5,4 @@ pyarrow
 tqdm
 lightning-sdk==0.1.46 # Must be pinned to ensure compatibility
 google-cloud-storage
+polars
diff --git a/src/litdata/__init__.py b/src/litdata/__init__.py
index f8c1ef7..af6d311 100644
--- a/src/litdata/__init__.py
+++ b/src/litdata/__init__.py
@@ -18,6 +18,7 @@ from litdata.streaming.combined import CombinedStreamingDataset
 from litdata.streaming.dataloader import StreamingDataLoader
 from litdata.streaming.dataset import StreamingDataset
 from litdata.streaming.item_loader import TokensLoader
+from litdata.streaming.writer import index_parquet_dataset
 from litdata.utilities.breakpoint import breakpoint
 from litdata.utilities.train_test_split import train_test_split
 
@@ -31,6 +32,7 @@ __all__ = [
     "walk",
     "train_test_split",
     "merge_datasets",
+    "index_parquet_dataset",
     "breakpoint",
 ]
 if RequirementCache("lightning_sdk"):
diff --git a/src/litdata/constants.py b/src/litdata/constants.py
index 1f4fb85..bdb9897 100644
--- a/src/litdata/constants.py
+++ b/src/litdata/constants.py
@@ -36,6 +36,7 @@ _GOOGLE_STORAGE_AVAILABLE = RequirementCache("google.cloud.storage")
 _AZURE_STORAGE_AVAILABLE = RequirementCache("azure.storage.blob")
 _TQDM_AVAILABLE = RequirementCache("tqdm")
 _LIGHTNING_SDK_AVAILABLE = RequirementCache("lightning_sdk")
+_POLARS_AVAILABLE = RequirementCache("polars")
 _DEBUG = bool(int(os.getenv("DEBUG", "1")))
 
 _MAX_WAIT_TIME = int(os.getenv("MAX_WAIT_TIME", "120"))
diff --git a/src/litdata/streaming/config.py b/src/litdata/streaming/config.py
index 5e1a2f3..117b2ef 100644
--- a/src/litdata/streaming/config.py
+++ b/src/litdata/streaming/config.py
@@ -70,7 +70,8 @@ class ChunksConfig:
         else:
             self._chunks = load_subsampled_chunks(subsampled_files, _original_chunks)
 
-        self._config["data_spec"] = treespec_loads(self._config["data_spec"])
+        if self._config["data_spec"] is not None:
+            self._config["data_spec"] = treespec_loads(self._config["data_spec"])
 
         assert self._chunks is not None
         self._item_loader.setup(self._config, self._chunks, serializers, region_of_interest)
@@ -229,7 +230,7 @@ class ChunksConfig:
 
         filesize_bytes = chunk["chunk_bytes"]
 
-        if self._config and self._config.get("encryption") is None:
+        if self._config and self._config.get("encryption") is None and (not local_chunkpath.endswith(".parquet")):
             filesize_bytes += (1 + chunk["chunk_size"]) * 4
 
         return local_chunkpath, begin, filesize_bytes
diff --git a/src/litdata/streaming/item_loader.py b/src/litdata/streaming/item_loader.py
index 06f2720..33ae886 100644
--- a/src/litdata/streaming/item_loader.py
+++ b/src/litdata/streaming/item_loader.py
@@ -24,7 +24,13 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import numpy as np
 import torch
 
-from litdata.constants import _FORCE_DOWNLOAD_TIME, _MAX_WAIT_TIME, _NUMPY_DTYPES_MAPPING, _TORCH_DTYPES_MAPPING
+from litdata.constants import (
+    _FORCE_DOWNLOAD_TIME,
+    _MAX_WAIT_TIME,
+    _NUMPY_DTYPES_MAPPING,
+    _POLARS_AVAILABLE,
+    _TORCH_DTYPES_MAPPING,
+)
 from litdata.streaming.serializers import Serializer
 from litdata.utilities._pytree import PyTree, tree_unflatten
 from litdata.utilities.encryption import Encryption, EncryptionLevel
@@ -412,3 +418,84 @@ class TokensLoader(BaseItemLoader):
     @classmethod
     def encode_data(cls, data: List[bytes], _: List[int], flattened: List[Any]) -> Tuple[bytes, Optional[int]]:
         return data[0], flattened[0].shape[0]
+
+
+class ParquetLoader(BaseItemLoader):
+    def __init__(self) -> None:
+        if not _POLARS_AVAILABLE:
+            raise ModuleNotFoundError("Please, run: `pip install polars`")
+        self._chunk_filepaths: Dict[str, bool] = {}
+
+    def setup(
+        self,
+        config: Dict,
+        chunks: List,
+        serializers: Dict[str, Serializer],
+        region_of_interest: Optional[List[Tuple[int, int]]] = None,
+    ) -> None:
+        self._config = config
+        self._chunks = chunks
+        self._serializers = {**serializers}
+        self._data_format = self._config["data_format"]
+        self._shift_idx = len(self._data_format) * 4
+        self.region_of_interest = region_of_interest
+        self._df: Dict[str, Any] = {}
+
+    def generate_intervals(self) -> List[Interval]:
+        intervals = []
+        begin = 0
+        end = 0
+        for idx, curr_chunk in enumerate(self._chunks):
+            end += curr_chunk["chunk_size"]
+            start_idx, end_idx = begin, end
+            if self.region_of_interest is not None:
+                start_idx = begin + self.region_of_interest[idx][0]
+                end_idx = begin + self.region_of_interest[idx][1]
+
+            intervals.append(Interval(begin, start_idx, end_idx, end))
+            begin += curr_chunk["chunk_size"]
+        return intervals
+
+    def pre_load_chunk(self, chunk_index: int, chunk_filepath: str) -> None:
+        """Logic to load the chunk in background to gain some time."""
+        pass
+
+    def load_item_from_chunk(
+        self,
+        index: int,
+        chunk_index: int,
+        chunk_filepath: str,
+        begin: int,
+        filesize_bytes: int,
+    ) -> Any:
+        """Returns an item loaded from a chunk."""
+        if chunk_filepath in self._chunk_filepaths and not os.path.isfile(chunk_filepath):
+            del self._chunk_filepaths[chunk_filepath]
+
+        if chunk_filepath not in self._chunk_filepaths:
+            exists = os.path.exists(chunk_filepath) and os.stat(chunk_filepath).st_size >= filesize_bytes
+
+            while not exists:
+                sleep(0.1)
+                exists = os.path.exists(chunk_filepath) and os.stat(chunk_filepath).st_size >= filesize_bytes
+
+            self._chunk_filepaths[chunk_filepath] = True
+
+        return self.get_df(chunk_filepath).row(index - begin)
+
+    def get_df(self, chunk_filepath: str) -> Any:
+        import polars as pl
+
+        if chunk_filepath not in self._df:
+            self._df[chunk_filepath] = pl.scan_parquet(chunk_filepath).collect()
+        return self._df[chunk_filepath]
+
+    def delete(self, chunk_index: int, chunk_filepath: str) -> None:
+        """Delete a chunk from the local filesystem."""
+        if os.path.exists(chunk_filepath):
+            os.remove(chunk_filepath)
+        if chunk_filepath in self._df:
+            del self._df[chunk_filepath]
+
+    def encode_data(self, data: List[bytes], sizes: List[int], flattened: List[Any]) -> Any:
+        pass
diff --git a/src/litdata/streaming/sampler.py b/src/litdata/streaming/sampler.py
index d135979..4c67b51 100644
--- a/src/litdata/streaming/sampler.py
+++ b/src/litdata/streaming/sampler.py
@@ -22,6 +22,26 @@ logger = logging.Logger(__name__)
 
 @dataclass
 class ChunkedIndex:
+    """Represents an index within a chunked dataset.
+
+    Attributes:
+        index (int): The global index of the data point across all chunks.
+        chunk_index (int): The index of the chunk where the data point resides.
+        chunk_indexes (Optional[List[int]]): A list specifying the range of indexes
+            allowed to read for this data point, in the form
+            [start, can_read_from, can_read_till, end]. Defaults to None.
+        is_last_index (bool): Indicates whether this is the last index in the dataset.
+            Defaults to False.
+
+    Suppose there are 3 chunk files:
+    - chunk-0.bin: Contains data points with global indexes 0-4.
+    - chunk-1.bin: Contains data points with global indexes 5-9.
+    - chunk-2.bin: Contains data points with global indexes 10-14.
+
+        A `ChunkedIndex` instance for the 6th data point (global index 5) would look like:
+            ChunkedIndex(index=5, chunk_index=1, chunk_indexes=[4,4,8,8], is_last_index=False)
+    """
+
     index: int
     chunk_index: int
     chunk_indexes: Optional[List[int]] = None
diff --git a/src/litdata/streaming/writer.py b/src/litdata/streaming/writer.py
index 10f1f11..74007e9 100644
--- a/src/litdata/streaming/writer.py
+++ b/src/litdata/streaming/writer.py
@@ -21,15 +21,16 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 
-from litdata.constants import _INDEX_FILENAME
+from litdata.constants import _INDEX_FILENAME, _POLARS_AVAILABLE
 from litdata.processing.utilities import get_worker_rank
 from litdata.streaming.compression import _COMPRESSORS, Compressor
-from litdata.streaming.item_loader import BaseItemLoader, PyTreeLoader
+from litdata.streaming.item_loader import BaseItemLoader, ParquetLoader, PyTreeLoader
 from litdata.streaming.serializers import Serializer, _get_serializers
 from litdata.utilities._pytree import PyTree, tree_flatten, treespec_dumps
 from litdata.utilities.encryption import Encryption, EncryptionLevel
 from litdata.utilities.env import _DistributedEnv, _WorkerEnv
 from litdata.utilities.format import _convert_bytes_to_int, _human_readable_bytes
+from litdata.utilities.parquet import get_parquet_indexer_cls
 
 
 @dataclass
@@ -530,3 +531,48 @@ class BinaryWriter:
             json.dump(checkPoint, f)
 
         return checkpoint_filepath
+
+
+def index_parquet_dataset(
+    pq_dir_url: str, cache_dir: Optional[str] = None, storage_options: Optional[Dict] = {}
+) -> None:
+    if not _POLARS_AVAILABLE:
+        raise ModuleNotFoundError("Please, run: `pip install polars`")
+
+    import polars as pl
+
+    pq_chunks_info = []
+    config: Dict[str, Any] = {
+        "compression": None,
+        "chunk_size": None,
+        "chunk_bytes": None,
+        "data_format": [],
+        "data_spec": None,
+        "encryption": None,
+        "item_loader": ParquetLoader.__name__,
+    }
+
+    pq_dir_class = get_parquet_indexer_cls(pq_dir_url, cache_dir, storage_options)
+    # iterate the directory and for all files ending in `.parquet` index them
+    for file_name, file_path in pq_dir_class:
+        file_size = os.path.getsize(file_path)
+        pq_polars = pl.scan_parquet(file_path)
+        chunk_dtypes = pq_polars.collect_schema().dtypes()
+        chunk_dtypes = [str(dt) for dt in chunk_dtypes]
+        chunk_size = pq_polars.select(pl.count()).collect().item()
+
+        if len(config["data_format"]) != 0 and config["data_format"] != chunk_dtypes:
+            raise Exception(
+                "The config isn't consistent between chunks. This shouldn't have happened."
+                f"Found {config}; {chunk_dtypes}."
+            )
+        config["data_format"] = chunk_dtypes
+        chunk_info = {
+            "chunk_bytes": file_size,
+            "chunk_size": chunk_size,
+            "filename": file_name,
+            "dim": None,
+        }
+        pq_chunks_info.append(chunk_info)
+
+    pq_dir_class.write_index(pq_chunks_info, config)
diff --git a/src/litdata/utilities/parquet.py b/src/litdata/utilities/parquet.py
new file mode 100644
index 0000000..bb13160
--- /dev/null
+++ b/src/litdata/utilities/parquet.py
@@ -0,0 +1,138 @@
+"""contains utility functions to return parquet files from local, s3, or gs."""
+
+import json
+import os
+from abc import ABC, abstractmethod
+from time import time
+from typing import Any, Dict, Generator, List, Optional, Tuple, Type, Union
+
+from litdata.constants import _INDEX_FILENAME
+from litdata.streaming.resolver import Dir, _resolve_dir
+
+
+class ParquetDir(ABC):
+    def __init__(
+        self,
+        dir_path: Optional[Union[str, Dir]],
+        cache_path: Optional[str] = None,
+        storage_options: Optional[Dict] = {},
+    ):
+        self.dir = _resolve_dir(dir_path)
+        self.cache_path = cache_path
+        self.storage_options = storage_options
+
+    @abstractmethod
+    def __iter__(self) -> Generator[Tuple[str, str], None, None]: ...
+
+    @abstractmethod
+    def write_index(self, chunks_info: List[Dict[str, Any]], config: Dict[str, Any]) -> None: ...
+
+
+class LocalParquetDir(ParquetDir):
+    def __init__(
+        self,
+        dir_path: Optional[Union[str, Dir]],
+        cache_path: Optional[str] = None,
+        storage_options: Optional[Dict] = {},
+    ):
+        super().__init__(dir_path, cache_path, storage_options)
+
+    def __iter__(self) -> Generator[Tuple[str, str], None, None]:
+        assert self.dir.path is not None
+        assert self.dir.path != "", "Dir path can't be empty"
+
+        for file_name in os.listdir(self.dir.path):
+            if file_name.endswith(".parquet"):
+                file_path = os.path.join(self.dir.path, file_name)
+                yield file_name, file_path
+
+    def write_index(self, chunks_info: List[Dict[str, Any]], config: Dict[str, Any]) -> None:
+        # write to index.json file
+        if self.cache_path is None:
+            self.cache_path = self.dir.path
+        assert self.cache_path is not None
+        with open(os.path.join(self.cache_path, _INDEX_FILENAME), "w") as f:
+            data = {"chunks": chunks_info, "config": config, "updated_at": str(time())}
+            json.dump(data, f, sort_keys=True)
+
+        print(f"Index file written to: {os.path.join(self.cache_path, _INDEX_FILENAME)}")
+
+
+class CloudParquetDir(ParquetDir):
+    def __init__(
+        self,
+        dir_path: Optional[Union[str, Dir]],
+        cache_path: Optional[str] = None,
+        storage_options: Optional[Dict] = None,
+    ):
+        super().__init__(dir_path, cache_path, storage_options)
+        if self.cache_path is None:
+            self.cache_path = os.path.join(os.path.expanduser("~"), ".cache", ".litdata-cache-index-pq")
+            os.makedirs(self.cache_path, exist_ok=True)  # Ensure the directory exists
+
+        import fsspec
+
+        assert self.dir.url is not None
+
+        _CLOUD_PROVIDER = ("s3", "gs")
+
+        for provider in _CLOUD_PROVIDER:
+            if self.dir.url.startswith(provider):
+                # Initialize the cloud filesystem
+                self.fs = fsspec.filesystem(provider, *self.storage_options)
+                print(f"using provider: {provider}")
+                break
+
+    def __iter__(self) -> Generator[Tuple[str, str], None, None]:
+        assert self.dir.url is not None
+        assert self.cache_path is not None
+
+        # List all files and directories in the top-level of the specified directory
+        files = self.fs.ls(self.dir.url, detail=True)
+
+        # Iterate through the items and check for Parquet files
+        for file_info in files:
+            if file_info["type"] == "file" and file_info["name"].endswith(".parquet"):
+                file_name = os.path.basename(file_info["name"])
+                assert self.cache_path is not None
+                local_path = os.path.join(self.cache_path, file_name)
+
+                # Download the file
+                with self.fs.open(file_info["name"], "rb") as cloud_file, open(local_path, "wb") as local_file:
+                    local_file.write(cloud_file.read())
+
+                yield file_name, local_path
+
+    def write_index(self, chunks_info: List[Dict[str, Any]], config: Dict[str, Any]) -> None:
+        assert self.cache_path is not None
+        assert self.dir.url is not None
+
+        index_file_path = os.path.join(self.cache_path, _INDEX_FILENAME)
+        cloud_index_path = os.path.join(self.dir.url, _INDEX_FILENAME)
+        # write to index.json file
+        with open(index_file_path, "w") as f:
+            data = {"chunks": chunks_info, "config": config, "updated_at": str(time())}
+            json.dump(data, f, sort_keys=True)
+
+        # upload file to cloud
+        with open(index_file_path, "rb") as local_file, self.fs.open(cloud_index_path, "wb") as cloud_file:
+            cloud_file.write(local_file.read())
+
+        print(f"Index file written to: {cloud_index_path}")
+
+
+_PARQUET_DIR: Dict[str, Type[ParquetDir]] = {
+    "s3://": CloudParquetDir,
+    "gs://": CloudParquetDir,
+    "local:": LocalParquetDir,
+    "": LocalParquetDir,
+}
+
+
+def get_parquet_indexer_cls(
+    dir_path: str, cache_path: Optional[str] = None, storage_options: Optional[Dict] = {}
+) -> ParquetDir:
+    for k, cls in _PARQUET_DIR.items():
+        if str(dir_path).startswith(k):
+            return cls(dir_path, cache_path, storage_options)
+    raise ValueError(f"The provided `dir_path` {dir_path} doesn't have a ParquetDir class associated.")


Test patch:
diff --git a/tests/streaming/test_writer.py b/tests/streaming/test_writer.py
index 4ad9785..1800820 100644
--- a/tests/streaming/test_writer.py
+++ b/tests/streaming/test_writer.py
@@ -22,9 +22,11 @@ import torch
 from lightning_utilities.core.imports import RequirementCache
 
 from litdata.streaming.compression import _ZSTD_AVAILABLE
+from litdata.streaming.dataset import StreamingDataset
+from litdata.streaming.item_loader import ParquetLoader
 from litdata.streaming.reader import BinaryReader
 from litdata.streaming.sampler import ChunkedIndex
-from litdata.streaming.writer import BinaryWriter
+from litdata.streaming.writer import BinaryWriter, index_parquet_dataset
 from litdata.utilities.format import _FORMAT_TO_RATIO
 
 
@@ -269,3 +271,39 @@ def test_writer_save_checkpoint(tmpdir):
     for file in os.listdir(os.path.join(cache_dir, ".checkpoints")):
         assert file.__contains__("checkpoint-0")
         assert file.endswith(".json")
+
+
+def test_parquet_index_write(tmpdir):
+    import polars as pl
+
+    os.mkdir(os.path.join(tmpdir, "data"))
+
+    for i in range(5):
+        df = pl.DataFrame(
+            {
+                "name": ["Tom", "Jerry", "Micky", "Oggy", "Doraemon"],
+                "weight": [57.9, 72.5, 53.6, 83.1, 69.4],  # (kg)
+                "height": [1.56, 1.77, 1.65, 1.75, 1.63],  # (m)
+            }
+        )
+        file_path = os.path.join(tmpdir, "data", f"tmp-{i}.parquet")
+        df.write_parquet(file_path)
+
+    index_file_path = os.path.join(tmpdir, "data", "index.json")
+    assert not os.path.exists(index_file_path)
+    # call the write_parquet_index fn
+    index_parquet_dataset(os.path.join(tmpdir, "data"))
+    assert os.path.exists(index_file_path)
+
+    # Read JSON file into a dictionary
+    with open(index_file_path) as f:
+        data = json.load(f)
+        assert len(data["chunks"]) == 5
+        for cnk in data["chunks"]:
+            assert cnk["chunk_size"] == 5
+        assert data["config"]["item_loader"] == "ParquetLoader"
+        assert data["config"]["data_format"] == ["String", "Float64", "Float64"]
+
+    ds = StreamingDataset(os.path.join(tmpdir, "data"), item_loader=ParquetLoader())
+
+    assert len(ds) == 25  # 5 datasets for 5 loops