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:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/Lightning-AI/litdata/pull/443?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Lightning-AI) Report
Attention: Patch coverage is `50.00000%` with `3 lines` in your changes missing coverage. Please review.
> Project coverage is 72%. Comparing base [(`bd116a4`)](https://app.codecov.io/gh/Lightning-AI/litdata/commit/bd116a46f53ebb3d6133d73b6d1b008e685aad35?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Lightning-AI) to head [(`6bde708`)](https://app.codecov.io/gh/Lightning-AI/litdata/commit/6bde708db22895c3436c82e33e3360ef46b7825e?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Lightning-AI).

> :exclamation:  There is a different number of reports uploaded between BASE (bd116a4) and HEAD (6bde708). Click for more details.
> 
> <details><summary>HEAD has 12 uploads less than BASE</summary>
>
>| Flag | BASE (bd116a4) | HEAD (6bde708) |
>|------|------|------|
>|unittests|5|1|
>|3.9|3|1|
>|ubuntu-22.04|2|0|
>|3.10|2|0|
>|macos-13|2|0|
></details>

<details><summary>Additional details and impacted files</summary>


```diff
@@         Coverage Diff          @@
##           main   #443    +/-   ##
====================================
- Coverage    78%    72%    -7%     
====================================
  Files        35     35            
  Lines      5106   5097     -9     
====================================
- Hits       4006   3664   -342     
- Misses     1100   1433   +333     
```

</details>
tchaton: Hey @deependujha Nice progress ;) 
gitguardian[bot]: #### ⚠️ GitGuardian has uncovered 1 secret following the scan of your pull request.

Please consider investigating the findings and remediating the incidents. Failure to do so may lead to compromising the associated services or software components.

Since your pull request originates from a forked repository, GitGuardian is not able to associate the secrets uncovered with secret incidents on your GitGuardian dashboard.
Skipping this check run and merging your pull request will create secret incidents on your GitGuardian dashboard.

<details>
<summary>🔎 Detected hardcoded secret in your pull request</summary>
<br>

| GitGuardian id | GitGuardian status | Secret                         | Commit           | Filename        |                      |
| -------------- | ------------------ | ------------------------------ | ---------------- | --------------- | -------------------- |
| [5685611](https://dashboard.gitguardian.com/workspace/49923/incidents/5685611) | Triggered | Generic High Entropy Secret | 76efafb83ffead35b9d24f466f1ca5542a8739d7 | tests/streaming/test_resolver.py | [View secret](https://github.com/Lightning-AI/litdata/commit/76efafb83ffead35b9d24f466f1ca5542a8739d7#diff-e3f43459fd170cf2acd5cb8d4febb926b3d640a7f976118a45db37a01d034036R141) |



</details>

<details>
<summary>🛠 Guidelines to remediate hardcoded secrets</summary>
<br>

1. Understand the implications of revoking this secret by investigating where it is used in your code.
2. Replace and store your secret safely. [Learn here](https://blog.gitguardian.com/secrets-api-management) the best practices.
3. Revoke and [rotate this secret](https://docs.gitguardian.com/secrets-detection/detectors/specifics/private_key_openssh#revoke-the-secret).
4. If possible, [rewrite git history](https://blog.gitguardian.com/rewriting-git-history-cheatsheet). Rewriting git history is not a trivial act. You might completely break other contributing developers&#x27; workflow and you risk accidentally deleting legitimate data.

To avoid such incidents in the future consider

-   following these [best practices](https://blog.gitguardian.com/secrets-api-management/) for managing and storing secrets including API keys and other credentials
-   install [secret detection on pre-commit](https://docs.gitguardian.com/ggshield-docs/integrations/git-hooks/pre-commit) to catch secret before it leaves your machine and ease remediation.


</details>

---

<sup>🦉 [GitGuardian](https://dashboard.gitguardian.com/auth/login/?utm_medium=checkruns&amp;utm_source=github&amp;utm_campaign=cr1) detects secrets in your source code to help developers and security teams secure the modern development process. You are seeing this because you or someone else with access to this repository has authorized GitGuardian to scan your pull request.<br/></sup>
deependujha: Adding support for directly consuming HF datasets is an exciting direction!

For HF datasets, my current idea involves iterating through all the Parquet datasets in the HF repository and creating an index.json file that is stored in a cache (since modifying the original dataset is not feasible).

When using the streaming dataset/dataloader, we would then pass this separate index.json file from the cache.

At this point, I'm uncertain about the exact approach for handling HF datasets comprehensively. This PR is ready for review and lays the groundwork for future enhancements. We can discuss HF dataset integration in a subsequent PR.

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