Repository: Lightning-AI/litData

Issue:
Add streaming support for huggingface parquet dataset similar to chunk stream
## 🚀 Feature Request  

**Add Streaming Support for Hugging Face Parquet Datasets Similar to Chunk-Based Streaming**  

### Motivation  
Currently, the entire dataset must be downloaded before indexing, resulting in significant waiting time before streaming can begin. This approach is inefficient, especially for datasets with a large number of chunks.  

### Pitch  
The proposed solution is to read only the metadata of all Parquet files using byte-range requests, allowing for quick indexing. Streaming can then proceed similarly to LitData's chunk-based streaming approach.  

Reference: [Introduction to Parquet Format](https://huggingface.co/blog/cfahlgren1/intro-to-parquet-format)

Discussion:

Solution patch:
diff --git a/src/litdata/constants.py b/src/litdata/constants.py
index 49559d7..018271a 100644
--- a/src/litdata/constants.py
+++ b/src/litdata/constants.py
@@ -39,6 +39,7 @@ _AZURE_STORAGE_AVAILABLE = RequirementCache("azure.storage.blob")
 _TQDM_AVAILABLE = RequirementCache("tqdm")
 _LIGHTNING_SDK_AVAILABLE = RequirementCache("lightning_sdk")
 _HF_HUB_AVAILABLE = RequirementCache("huggingface_hub")
+_PYARROW_AVAILABLE = RequirementCache("pyarrow")
 _POLARS_AVAILABLE = RequirementCache("polars>1.0.0")
 _DEBUG = bool(int(os.getenv("DEBUG_LITDATA", "0")))
 
diff --git a/src/litdata/processing/readers.py b/src/litdata/processing/readers.py
index eff6fc5..a1ebdef 100644
--- a/src/litdata/processing/readers.py
+++ b/src/litdata/processing/readers.py
@@ -16,13 +16,10 @@ import os
 from abc import ABC, abstractmethod
 from typing import Any, List
 
-from lightning_utilities.core.imports import RequirementCache
-
+from litdata.constants import _PYARROW_AVAILABLE
 from litdata.streaming.dataloader import StreamingDataLoader
 from litdata.utilities.format import _get_tqdm_iterator_if_available
 
-_PYARROW_AVAILABLE = RequirementCache("pyarrow")
-
 
 class BaseReader(ABC):
     """The `BaseReader` interface defines how to read and preprocess data
diff --git a/src/litdata/streaming/dataset.py b/src/litdata/streaming/dataset.py
index b50e45e..f173794 100644
--- a/src/litdata/streaming/dataset.py
+++ b/src/litdata/streaming/dataset.py
@@ -100,11 +100,25 @@ class StreamingDataset(IterableDataset):
 
         if input_dir.url is not None and input_dir.url.startswith("hf://"):
             if index_path is None:
-                # no index path provide, load from cache, or try indexing on the go.
+                # No index_path was provided. Attempt to load it from cache or generate it dynamically on the fly.
                 index_path = index_hf_dataset(input_dir.url)
                 cache_dir.path = index_path
                 input_dir.path = index_path
-            item_loader = ParquetLoader()
+
+            if item_loader is not None and not isinstance(item_loader, ParquetLoader):
+                raise ValueError(
+                    "Invalid item_loader for hf://datasets. "
+                    "The item_loader must be an instance of ParquetLoader. "
+                    "Please provide a valid ParquetLoader instance."
+                )
+
+            if item_loader is not None and item_loader._low_memory and shuffle:
+                raise ValueError(
+                    "You have enabled shuffling when using low memory with ParquetLoader. "
+                    "This configuration may lead to performance issues during the training process. "
+                    "Consider disabling shuffling or using a ParquetLoader without low memory mode."
+                )
+            item_loader = item_loader or ParquetLoader()
 
         self.input_dir = input_dir
         self.cache_dir = cache_dir
@@ -548,9 +562,7 @@ class StreamingDataset(IterableDataset):
                     "The provided `item_loader` state doesn't match the current one. "
                     f"Found `{self.item_loader.state_dict()}` instead of `{state['item_loader']}`."
                 )
-            logger.warning(
-                f"Overriding state item_loader {state['item_loader']} " f"to {self.item_loader.state_dict()}."
-            )
+            logger.warning(f"Overriding state item_loader {state['item_loader']} to {self.item_loader.state_dict()}.")
             state["item_loader"] = self.item_loader.state_dict()
 
         if state["drop_last"] != self.drop_last:
diff --git a/src/litdata/streaming/downloader.py b/src/litdata/streaming/downloader.py
index 02a7904..5e595a3 100644
--- a/src/litdata/streaming/downloader.py
+++ b/src/litdata/streaming/downloader.py
@@ -15,6 +15,7 @@ import contextlib
 import os
 import shutil
 import subprocess
+import tempfile
 from abc import ABC
 from contextlib import suppress
 from typing import Any, Dict, List, Optional, Type
@@ -58,9 +59,9 @@ class Downloader(ABC):
         local_chunkpath = os.path.join(self._cache_dir, chunk_filename)
         remote_chunkpath = os.path.join(self._remote_dir, chunk_filename)
 
-        self.download_file(remote_chunkpath, local_chunkpath, chunk_filename)
+        self.download_file(remote_chunkpath, local_chunkpath)
 
-    def download_file(self, remote_chunkpath: str, local_chunkpath: str, remote_chunk_filename: str = "") -> None:
+    def download_file(self, remote_chunkpath: str, local_chunkpath: str) -> None:
         pass
 
 
@@ -74,7 +75,7 @@ class S3Downloader(Downloader):
         if not self._s5cmd_available or _DISABLE_S5CMD:
             self._client = S3Client(storage_options=self._storage_options)
 
-    def download_file(self, remote_filepath: str, local_filepath: str, remote_chunk_filename: str = "") -> None:
+    def download_file(self, remote_filepath: str, local_filepath: str) -> None:
         obj = parse.urlparse(remote_filepath)
 
         if obj.scheme != "s3":
@@ -158,7 +159,7 @@ class GCPDownloader(Downloader):
 
         super().__init__(remote_dir, cache_dir, chunks, storage_options)
 
-    def download_file(self, remote_filepath: str, local_filepath: str, remote_chunk_filename: str = "") -> None:
+    def download_file(self, remote_filepath: str, local_filepath: str) -> None:
         from google.cloud import storage
 
         obj = parse.urlparse(remote_filepath)
@@ -193,7 +194,7 @@ class AzureDownloader(Downloader):
 
         super().__init__(remote_dir, cache_dir, chunks, storage_options)
 
-    def download_file(self, remote_filepath: str, local_filepath: str, remote_chunk_filename: str = "") -> None:
+    def download_file(self, remote_filepath: str, local_filepath: str) -> None:
         from azure.storage.blob import BlobServiceClient
 
         obj = parse.urlparse(remote_filepath)
@@ -220,7 +221,7 @@ class AzureDownloader(Downloader):
 
 
 class LocalDownloader(Downloader):
-    def download_file(self, remote_filepath: str, local_filepath: str, remote_chunk_filename: str = "") -> None:
+    def download_file(self, remote_filepath: str, local_filepath: str) -> None:
         if not os.path.exists(remote_filepath):
             raise FileNotFoundError(f"The provided remote_path doesn't exist: {remote_filepath}")
 
@@ -248,32 +249,42 @@ class HFDownloader(Downloader):
             )
 
         super().__init__(remote_dir, cache_dir, chunks, storage_options)
-        from huggingface_hub import HfFileSystem
 
-        self.fs = HfFileSystem()
+    def download_file(self, remote_filepath: str, local_filepath: str) -> None:
+        """Download a file from the Hugging Face Hub.
+        The remote_filepath should be in the format `hf://<repo_type>/<repo_org>/<repo_name>/path`. For more
+        information, see
+        https://huggingface.co/docs/huggingface_hub/en/guides/hf_file_system#integrations.
+        """
+        from huggingface_hub import hf_hub_download
 
-    def download_file(self, remote_filepath: str, local_filepath: str, remote_chunk_filename: str = "") -> None:
-        # for HF dataset downloading, we don't need remote_filepath, but remote_chunk_filename
-        with suppress(Timeout), FileLock(local_filepath + ".lock", timeout=0):
-            temp_path = local_filepath + ".tmp"  # Avoid partial writes
-            try:
-                with self.fs.open(remote_chunk_filename, "rb") as cloud_file, open(temp_path, "wb") as local_file:
-                    for chunk in iter(lambda: cloud_file.read(4096), b""):  # Stream in 4KB chunks local_file.
-                        local_file.write(chunk)
+        obj = parse.urlparse(remote_filepath)
+
+        if obj.scheme != "hf":
+            raise ValueError(f"Expected obj.scheme to be `hf`, instead, got {obj.scheme} for remote={remote_filepath}")
 
-                os.rename(temp_path, local_filepath)  # Atomic move after successful write
+        if os.path.exists(local_filepath):
+            return
 
-            except Exception as e:
-                print(f"Error processing {remote_chunk_filename}: {e}")
+        with suppress(Timeout), FileLock(local_filepath + ".lock", timeout=0), tempfile.TemporaryDirectory() as tmpdir:
+            _, _, _, repo_org, repo_name, path = remote_filepath.split("/", 5)
+            repo_id = f"{repo_org}/{repo_name}"
 
-            finally:
-                # Ensure cleanup of temp file if an error occurs
-                if os.path.exists(temp_path):
-                    os.remove(temp_path)
+            downloaded_path = hf_hub_download(
+                repo_id,
+                path,
+                cache_dir=tmpdir,
+                repo_type="dataset",
+                **self._storage_options,
+            )
+            if downloaded_path != local_filepath and os.path.exists(downloaded_path):
+                temp_file_path = local_filepath + ".tmp"
+                shutil.copyfile(downloaded_path, temp_file_path)
+                os.rename(temp_file_path, local_filepath)
 
 
 class LocalDownloaderWithCache(LocalDownloader):
-    def download_file(self, remote_filepath: str, local_filepath: str, remote_chunk_filename: str = "") -> None:
+    def download_file(self, remote_filepath: str, local_filepath: str) -> None:
         remote_filepath = remote_filepath.replace("local:", "")
         super().download_file(remote_filepath, local_filepath)
 
diff --git a/src/litdata/streaming/item_loader.py b/src/litdata/streaming/item_loader.py
index 1c4b24e..ea295a9 100644
--- a/src/litdata/streaming/item_loader.py
+++ b/src/litdata/streaming/item_loader.py
@@ -10,8 +10,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import functools
+import logging
 import os
 from abc import ABC, abstractmethod
 from collections import defaultdict, namedtuple
@@ -29,6 +29,7 @@ from litdata.constants import (
     _MAX_WAIT_TIME,
     _NUMPY_DTYPES_MAPPING,
     _POLARS_AVAILABLE,
+    _PYARROW_AVAILABLE,
     _TORCH_DTYPES_MAPPING,
 )
 from litdata.streaming.serializers import Serializer
@@ -37,6 +38,8 @@ from litdata.utilities.encryption import Encryption, EncryptionLevel
 
 Interval = namedtuple("Interval", ["chunk_start", "roi_start_idx", "roi_end_idx", "chunk_end"])
 
+logger = logging.getLogger(__name__)
+
 
 class BaseItemLoader(ABC):
     """The base item loader is responsible to decide how the items within a chunk are loaded."""
@@ -527,13 +530,25 @@ class TokensLoader(BaseItemLoader):
 
 
 class ParquetLoader(BaseItemLoader):
-    def __init__(self) -> None:
+    def __init__(self, pre_load_chunk: bool = False, low_memory: bool = True) -> None:
         if not _POLARS_AVAILABLE:
             raise ModuleNotFoundError(
                 "You are using the Parquet item loader, which depends on `Polars > 1.0.0`.",
                 "Please, run: `pip install polars>1.0.0`",
             )
+        if not _PYARROW_AVAILABLE:
+            raise ModuleNotFoundError("Please, run: `pip install pyarrow`")
+
         self._chunk_filepaths: Dict[str, bool] = {}
+        self._pre_load_chunk = pre_load_chunk
+        self._low_memory = low_memory
+
+        if not self._low_memory:
+            logger.warning(
+                "You have set low_memory=False in ParquetLoader. "
+                "This may result in high memory usage when processing large Parquet chunk files. "
+                "Consider setting low_memory=True to reduce memory consumption."
+            )
 
     def setup(
         self,
@@ -548,7 +563,9 @@ class ParquetLoader(BaseItemLoader):
         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] = {}
+        self._df: Dict[int, Any] = {}
+        self._chunk_row_groups: Dict[int, Any] = {}
+        self._chunk_row_group_item_read_count: Dict[int, Any] = {}
 
     def generate_intervals(self) -> List[Interval]:
         intervals = []
@@ -566,11 +583,14 @@ class ParquetLoader(BaseItemLoader):
         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."""
+        """Preload the chunk in the background to gain some time."""
+        if not self._pre_load_chunk or self._low_memory:
+            return
+
         import polars as pl
 
-        if chunk_filepath not in self._df:
-            self._df[chunk_filepath] = pl.scan_parquet(chunk_filepath).collect()
+        if chunk_index not in self._df and os.path.exists(chunk_filepath):
+            self._df[chunk_index] = pl.scan_parquet(chunk_filepath, low_memory=True).collect()
 
     def load_item_from_chunk(
         self,
@@ -580,7 +600,7 @@ class ParquetLoader(BaseItemLoader):
         begin: int,
         filesize_bytes: int,
     ) -> Any:
-        """Returns an item loaded from a chunk."""
+        """Returns an item loaded from a parquet chunk."""
         if chunk_filepath in self._chunk_filepaths and not os.path.isfile(chunk_filepath):
             del self._chunk_filepaths[chunk_filepath]
 
@@ -593,21 +613,112 @@ class ParquetLoader(BaseItemLoader):
 
             self._chunk_filepaths[chunk_filepath] = True
 
-        return self.get_df(chunk_filepath).row(index - begin)
+        # relative index of the desired row within the chunk.
+        relative_index = index - begin
+        if self._low_memory:
+            return self._get_item_with_low_memory(chunk_index, chunk_filepath, relative_index)
 
-    def get_df(self, chunk_filepath: str) -> Any:
+        return self._get_item(chunk_index, chunk_filepath, relative_index)
+
+    def _get_item_with_low_memory(self, chunk_index: int, chunk_filepath: str, row_index: int) -> Any:
+        """Retrieve a dataframe row from a parquet chunk in low memory mode.
+
+        This method reads only the necessary row group from the parquet file using PyArrow and Polars,
+        which helps in reducing memory usage.
+
+        Args:
+            chunk_index (int): The index of the chunk to be accessed.
+            chunk_filepath (str): The file path of the parquet chunk.
+            row_index (int): The relative row index within the loaded chunk.
+
+        Returns:
+            Any: The dataframe row corresponding to the specified index.
+        """
         import polars as pl
+        import pyarrow.parquet as pq
+
+        # Load the Parquet file metadata if not already loaded
+        if chunk_index not in self._df:
+            self._df[chunk_index] = pq.ParquetFile(chunk_filepath)
+
+        # Determine the row group and the row index within the row group
+        parquet_file = self._df[chunk_index]
+        num_rows_per_row_group = parquet_file.metadata.row_group(0).num_rows
+        row_group_index = row_index // num_rows_per_row_group
+        row_index_within_group = row_index % num_rows_per_row_group
+
+        # Check if the row group is already loaded
+        if chunk_index in self._chunk_row_groups and row_group_index in self._chunk_row_groups[chunk_index]:
+            # Use the cached row group
+            row_group_df = self._chunk_row_groups[chunk_index][row_group_index]
+            # update read count
+            self._chunk_row_group_item_read_count[chunk_index][row_group_index] += 1
+        else:
+            # Load the row group and convert it to a Polars DataFrame
+            row_group = self._df[chunk_index].read_row_group(row_group_index)
+            row_group_df = pl.from_arrow(row_group)
+
+            # Cache the loaded row group
+            if chunk_index not in self._chunk_row_groups:
+                self._chunk_row_groups[chunk_index] = {}
+                self._chunk_row_group_item_read_count[chunk_index] = {}
 
-        if chunk_filepath not in self._df:
-            self._df[chunk_filepath] = pl.scan_parquet(chunk_filepath).collect()
-        return self._df[chunk_filepath]
+            self._chunk_row_groups[chunk_index][row_group_index] = row_group_df
+            self._chunk_row_group_item_read_count[chunk_index][row_group_index] = 1
+
+        # Check if the row group has been fully read and release memory if necessary
+        read_count = self._chunk_row_group_item_read_count[chunk_index][row_group_index]
+        if read_count >= num_rows_per_row_group:
+            # Release memory for the fully read row group
+            del self._chunk_row_groups[chunk_index][row_group_index]
+            del self._chunk_row_group_item_read_count[chunk_index][row_group_index]
+
+        # Return the specific row from the dataframe
+        return row_group_df.row(row_index_within_group)  # type: ignore
+
+    def _get_item(self, chunk_index: int, chunk_filepath: str, index: int) -> Any:
+        """Retrieve a dataframe row from a parquet chunk by loading the entire chunk into memory.
+
+        Note:
+            This method reads the complete parquet file using Polars. Exercise caution with large files as it
+            may significantly increase memory usage.
+
+        Args:
+            chunk_index (int): The index of the chunk to be accessed.
+            chunk_filepath (str): The file path of the parquet chunk.
+            index (int): The relative row index within the loaded chunk.
+
+        Returns:
+            Any: The dataframe row corresponding to the specified index.
+        """
+        import polars as pl
+
+        if chunk_index not in self._df:
+            self._df[chunk_index] = pl.scan_parquet(chunk_filepath, low_memory=True).collect()
+        return self._df[chunk_index].row(index)
 
     def delete(self, chunk_index: int, chunk_filepath: str) -> None:
         """Delete a chunk from the local filesystem."""
+        if chunk_index in self._df:
+            del self._df[chunk_index]
+        if chunk_index in self._chunk_row_groups:
+            del self._chunk_row_groups[chunk_index]
+
+        if chunk_index in self._chunk_row_group_item_read_count:
+            del self._chunk_row_group_item_read_count[chunk_index]
         if os.path.exists(chunk_filepath):
             os.remove(chunk_filepath)
-        if chunk_filepath in self._df:
-            del self._df[chunk_filepath]
+
+    def close(self, chunk_index: int) -> None:
+        """Release the memory-mapped file for a specific chunk index."""
+        if chunk_index in self._df:
+            del self._df[chunk_index]
+
+        if chunk_index in self._chunk_row_groups:
+            del self._chunk_row_groups[chunk_index]
+
+        if chunk_index in self._chunk_row_group_item_read_count:
+            del self._chunk_row_group_item_read_count[chunk_index]
 
     def encode_data(self, data: List[bytes], sizes: List[int], flattened: List[Any]) -> Any:
         pass
diff --git a/src/litdata/streaming/reader.py b/src/litdata/streaming/reader.py
index 41c5afc..c71a350 100644
--- a/src/litdata/streaming/reader.py
+++ b/src/litdata/streaming/reader.py
@@ -24,7 +24,7 @@ from filelock import FileLock, Timeout
 
 from litdata.constants import _DEBUG
 from litdata.streaming.config import ChunksConfig, Interval
-from litdata.streaming.item_loader import BaseItemLoader, PyTreeLoader, TokensLoader
+from litdata.streaming.item_loader import BaseItemLoader, ParquetLoader, PyTreeLoader, TokensLoader
 from litdata.streaming.sampler import ChunkedIndex
 from litdata.streaming.serializers import Serializer, _get_serializers
 from litdata.utilities.encryption import Encryption
@@ -401,7 +401,7 @@ class BinaryReader:
 
         if index.chunk_index != self._last_chunk_index:
             # Close the memory-mapped file for the last chunk index
-            if isinstance(self._item_loader, TokensLoader) and self._last_chunk_index is not None:
+            if isinstance(self._item_loader, (TokensLoader, ParquetLoader)) and self._last_chunk_index is not None:
                 self._item_loader.close(self._last_chunk_index)
 
             # track the new chunk index as the latest one
diff --git a/src/litdata/streaming/writer.py b/src/litdata/streaming/writer.py
index 0adc34d..2771106 100644
--- a/src/litdata/streaming/writer.py
+++ b/src/litdata/streaming/writer.py
@@ -566,22 +566,21 @@ def index_parquet_dataset(
     pq_dir_url: str,
     cache_dir: Optional[str] = None,
     storage_options: Optional[Dict] = {},
-    remove_after_indexing: bool = True,
     num_workers: int = 4,
 ) -> None:
     """Index a Parquet dataset from a specified URL.
 
-    This function scans all `.parquet` files in the specified directory URL, extracts metadata
-    such as file size, chunk size, and data types, and indexes them. Optionally, the files
-    can be removed after indexing.
+
+    This function scans the metadata of all `.parquet` files in the specified directory URL, extracts metadata
+    such as file size, chunk size, and data types, and indexes them. The index is saved as a JSON file in the
+    specified cache directory.
 
     Args:
         pq_dir_url (str): URL of the directory containing the Parquet files.
         cache_dir (Optional[str]): Local cache directory for storing temporary files.
             For HF dataset, index.json file will be stored here.
         storage_options (Optional[Dict]): Additional storage options for accessing the Parquet files.
-        remove_after_indexing (bool): Whether to remove files after indexing (default is True).
-        num_workers (int): Number of workers to download files and index them.
+        num_workers (int): Number of workers to download metadata of Parquet files and index them.
 
     Raises:
         ModuleNotFoundError: If the required `polars` module is not installed.
@@ -589,8 +588,6 @@ def index_parquet_dataset(
     if not _POLARS_AVAILABLE:
         raise ModuleNotFoundError("Please, run: `pip install polars`")
 
-    import polars as pl
-
     pq_chunks_info = []
     config: Dict[str, Any] = {
         "compression": None,
@@ -602,7 +599,7 @@ def index_parquet_dataset(
         "item_loader": ParquetLoader.__name__,
     }
 
-    pq_dir_class = get_parquet_indexer_cls(pq_dir_url, cache_dir, storage_options, remove_after_indexing, num_workers)
+    pq_dir_class = get_parquet_indexer_cls(pq_dir_url, cache_dir, storage_options, num_workers)
 
     if _TQDM_AVAILABLE:
         from tqdm.auto import tqdm as _tqdm
@@ -616,13 +613,11 @@ def index_parquet_dataset(
             dynamic_ncols=True,
             unit="step",
         )
-    # 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()
+
+    results = {}
+    # iterate through the directory and index each file ending with ".parquet"
+    for file_metadata, order in pq_dir_class:
+        chunk_dtypes = file_metadata["data_types"]
 
         if len(config["data_format"]) != 0 and config["data_format"] != chunk_dtypes:
             raise Exception(
@@ -631,14 +626,18 @@ def index_parquet_dataset(
             )
         config["data_format"] = chunk_dtypes
         chunk_info = {
-            "chunk_bytes": file_size,
-            "chunk_size": chunk_size,
-            "filename": file_name,
+            "chunk_bytes": file_metadata["file_size"],
+            "chunk_size": file_metadata["num_rows"],
+            "filename": file_metadata["file_name"],
             "dim": None,
         }
-        pq_chunks_info.append(chunk_info)
+        results[order] = chunk_info
         if _TQDM_AVAILABLE:
             pbar.update(1)
 
+    for i in sorted(results.keys()):
+        pq_chunks_info.append(results[i])
+
+    del results
     print(flush=True)  # to prevent truncated printing when using concurrent threads/processes
     pq_dir_class.write_index(pq_chunks_info, config)
diff --git a/src/litdata/utilities/hf_dataset.py b/src/litdata/utilities/hf_dataset.py
index f9fce4b..7740c04 100644
--- a/src/litdata/utilities/hf_dataset.py
+++ b/src/litdata/utilities/hf_dataset.py
@@ -8,20 +8,27 @@ from litdata.utilities.parquet import default_cache_dir
 
 
 def index_hf_dataset(hf_url: str) -> str:
+    """Index a Hugging Face dataset and return the cache directory path.
+
+    Args:
+        hf_url (str): The URL of the Hugging Face dataset.
+
+    Returns:
+        str: The path to the cache directory containing the index.
+    """
     if not hf_url.startswith("hf://"):
-        raise ValueError(f"Invalid Hugging Face dataset URL: {hf_url}. Expected URL to start with 'hf://'.")
+        raise ValueError(
+            f"Invalid Hugging Face dataset URL: {hf_url}. "
+            "The URL should start with 'hf://'. Please check the URL and try again."
+        )
 
     cache_dir = default_cache_dir(hf_url)
-
     cache_index_path = os.path.join(cache_dir, _INDEX_FILENAME)
-    print("=" * 50)
 
     if os.path.exists(cache_index_path):
-        print(f"Found HF index.json file in {cache_index_path}.")
+        print(f"Found existing HF {_INDEX_FILENAME} file for {hf_url} at {cache_index_path}.")
     else:
-        print("Indexing HF dataset...")
-        index_parquet_dataset(hf_url, cache_dir, remove_after_indexing=False)
-
-    print("=" * 50)
+        print(f"Indexing HF dataset from {hf_url} into {cache_index_path}.")
+        index_parquet_dataset(hf_url, cache_dir, num_workers=os.cpu_count() or 4)
 
     return cache_dir
diff --git a/src/litdata/utilities/parquet.py b/src/litdata/utilities/parquet.py
index 03bb9d9..01bfbf2 100644
--- a/src/litdata/utilities/parquet.py
+++ b/src/litdata/utilities/parquet.py
@@ -1,97 +1,66 @@
 """contains utility functions to return parquet files from local, s3, or gs."""
 
 import hashlib
+import io
 import json
 import os
-import sys
-import threading
 from abc import ABC, abstractmethod
 from concurrent.futures import ThreadPoolExecutor
-from contextlib import suppress
-from queue import Queue
-from time import sleep, time
+from time import time
 from typing import Any, Dict, Generator, List, Optional, Tuple, Union
 from urllib import parse
 
-from litdata.constants import _FSSPEC_AVAILABLE, _HF_HUB_AVAILABLE, _INDEX_FILENAME, _TQDM_AVAILABLE
+from litdata.constants import _FSSPEC_AVAILABLE, _HF_HUB_AVAILABLE, _INDEX_FILENAME, _PYARROW_AVAILABLE
 from litdata.streaming.resolver import Dir, _resolve_dir
 
 
-def delete_thread(rmq: Queue) -> None:
-    while True:
-        file_path = rmq.get()
-        if file_path is None:  # Sentinel value to exit
-            break
-        with suppress(FileNotFoundError):
-            os.remove(file_path)
-
-    # before exiting, just sleep for some time, to complete any pending deletions
-    sleep(0.5)
-
-
 class ParquetDir(ABC):
     def __init__(
         self,
         dir_path: Optional[Union[str, Dir]],
         cache_path: Optional[str] = None,
         storage_options: Optional[Dict] = {},
-        remove_after_indexing: bool = True,
         num_workers: int = 4,
     ):
         self.dir = _resolve_dir(dir_path)
         self.cache_path = cache_path
         self.storage_options = storage_options
-        self.remove_after_indexing = remove_after_indexing
         self.files: List[Any] = []
-        self.process_queue: Queue = Queue()
-        self.delete_queue: Queue = Queue()
         self.num_workers = num_workers
 
-    def __iter__(self) -> Generator[Tuple[str, str], None, None]:
-        # start worker in a separate thread, and then read values from `process_queue`, and yield
+    def __iter__(self) -> Generator[Tuple[Dict[str, Any], int], None, None]:
+        """Iterate over the Parquet files and yield their metadata.
 
-        self.is_delete_thread_running = self.dir.url is not None and self.remove_after_indexing
-
-        if self.is_delete_thread_running:
-            t = threading.Thread(
-                target=delete_thread,
-                name="delete_thread",
-                args=(self.delete_queue,),
-            )
-            t.start()
-
-        t_worker = threading.Thread(
-            target=self.worker,
-            name="worker_thread",
-            args=(),
-        )
-        t_worker.start()
-
-        while True:
-            file_name, file_path = self.process_queue.get()
-            if file_name is None and file_path is None:  # Sentinel value to exit
-                break
-            yield file_name, file_path
-            if self.is_delete_thread_running:
-                self.delete_queue.put_nowait(file_path)
-
-        if self.is_delete_thread_running:
-            self.delete_queue.put_nowait(None)  # so that it doesn't hang indefinitely
-            t.join()
+        Yields:
+            Generator[Tuple[str, int], None, None]: A generator yielding tuples of file name, file path, and order.
+        """
+        with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
+            futures = {executor.submit(self.task, _file): (order, _file) for order, _file in enumerate(self.files)}
+            for future in futures:
+                file_metadata = future.result()
+                order, _ = futures[future]
+                yield file_metadata, order
 
-        t_worker.join()  # should've completed by now. But, just to be on safe side.
+    @abstractmethod
+    def task(self, _file: Any) -> Dict[str, Any]: ...
 
     @abstractmethod
-    def task(self, _file: Any) -> None: ...
+    def write_index(self, chunks_info: List[Dict[str, Any]], config: Dict[str, Any]) -> None:
+        """Write the index file to the cache directory.
 
-    def worker(self) -> None:
-        with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
-            for _file in self.files:
-                executor.submit(self.task, _file)
-        self.process_queue.put_nowait((None, None))
+        Args:
+            chunks_info (List[Dict[str, Any]]): List of dictionaries containing chunk metadata.
+            config (Dict[str, Any]): Configuration dictionary containing metadata.
+        """
+        assert self.cache_path is not None, "Cache path is not set."
+        index_file_path = os.path.join(self.cache_path, _INDEX_FILENAME)
 
-    @abstractmethod
-    def write_index(self, chunks_info: List[Dict[str, Any]], config: Dict[str, Any]) -> None: ...
+        # 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)
+
+        print(f"Index file successfully written to: {index_file_path}")
 
 
 class LocalParquetDir(ParquetDir):
@@ -100,35 +69,47 @@ class LocalParquetDir(ParquetDir):
         dir_path: Optional[Union[str, Dir]],
         cache_path: Optional[str] = None,
         storage_options: Optional[Dict] = {},
-        remove_after_indexing: bool = True,
         num_workers: int = 4,
     ):
-        super().__init__(dir_path, cache_path, storage_options, remove_after_indexing, num_workers)
+        if not _PYARROW_AVAILABLE:
+            raise ModuleNotFoundError(
+                "The 'pyarrow' module is required for processing Parquet files. "
+                "Please install it by running: `pip install pyarrow`"
+            )
+        super().__init__(dir_path, cache_path, storage_options, num_workers)
 
         for _f in os.listdir(self.dir.path):
             if _f.endswith(".parquet"):
                 self.files.append(_f)
 
-    def task(self, _file: str) -> None:
-        assert isinstance(_file, str)
-
-        if _file.endswith(".parquet"):
-            assert self.dir.path is not None
-            assert self.dir.path != "", "Dir path can't be empty"
+    def task(self, _file: str) -> Dict[str, Any]:
+        """Extract metadata from a Parquet file on the local filesystem."""
+        import pyarrow.parquet as pq
 
-            file_path = os.path.join(self.dir.path, _file)
-            self.process_queue.put_nowait((_file, file_path))
+        assert isinstance(_file, str)
+        assert _file.endswith(".parquet")
+        assert self.dir.path is not None
+        assert self.dir.path != "", "Dir path can't be empty"
+
+        file_path = os.path.join(self.dir.path, _file)
+        parquet_file = pq.ParquetFile(file_path)
+        num_rows = parquet_file.metadata.num_rows
+        data_types = [str(col.type) for col in parquet_file.schema_arrow]
+        file_size = os.path.getsize(file_path)
+        file_name = os.path.basename(file_path)
+
+        return {
+            "file_name": file_name,
+            "num_rows": num_rows,
+            "file_size": file_size,
+            "data_types": data_types,
+        }
 
     def write_index(self, chunks_info: List[Dict[str, Any]], config: Dict[str, Any]) -> None:
-        # write to index.json file
+        """Write the index file to the cache directory."""
         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)}")
+        super().write_index(chunks_info, config)
 
 
 class CloudParquetDir(ParquetDir):
@@ -137,15 +118,18 @@ class CloudParquetDir(ParquetDir):
         dir_path: Optional[Union[str, Dir]],
         cache_path: Optional[str] = None,
         storage_options: Optional[Dict] = None,
-        remove_after_indexing: bool = True,
         num_workers: int = 4,
     ):
         if not _FSSPEC_AVAILABLE:
             raise ModuleNotFoundError(
                 "Support for Indexing cloud parquet files depends on `fsspec`.", "Please, run: `pip install fsspec`"
             )
-
-        super().__init__(dir_path, cache_path, storage_options, remove_after_indexing, num_workers)
+        if not _PYARROW_AVAILABLE:
+            raise ModuleNotFoundError(
+                "The 'pyarrow' module is required for processing Parquet files. "
+                "Please install it by running: `pip install pyarrow`"
+            )
+        super().__init__(dir_path, cache_path, storage_options, num_workers)
 
         assert self.dir.url is not None
 
@@ -169,48 +153,64 @@ class CloudParquetDir(ParquetDir):
             if _f["type"] == "file" and _f["name"].endswith(".parquet"):
                 self.files.append(_f)
 
-    def task(self, _file: Any) -> None:
-        if _file["type"] == "file" and _file["name"].endswith(".parquet"):
-            file_name = os.path.basename(_file["name"])
-            assert self.cache_path is not None
-            local_path = os.path.join(self.cache_path, file_name)
-            temp_path = local_path + ".tmp"  # Avoid partial writes
-            # if an existing temp file is present, means its corrupted
-            if os.path.exists(temp_path):
-                os.remove(temp_path)
+    def task(self, _file: Any) -> Dict[str, Any]:
+        """Extract metadata from a Parquet file on the cloud filesystem without downloading the entire file."""
+        import pyarrow.parquet as pq
+
+        # Validate inputs
+        if not isinstance(_file, dict) or "name" not in _file or "size" not in _file:
+            raise ValueError(f"Invalid file object: {_file}")
+
+        assert _file["name"].endswith(".parquet")
+
+        file_name = os.path.basename(_file["name"])
+        file_size = _file["size"]
 
-            # Download the file
-            with self.fs.open(_file["name"], "rb") as cloud_file, open(temp_path, "wb") as local_file:
-                for chunk in iter(lambda: cloud_file.read(4096), b""):  # Read in 4KB chunks
-                    local_file.write(chunk)
+        with self.fs.open(_file["name"], "rb") as f:
+            # Read footer size (last 8 bytes: 4 for footer size + 4 for magic number)
+            f.seek(file_size - 8)
+            footer_size = int.from_bytes(f.read(4), "little")
 
-            os.rename(temp_path, local_path)  # Atomic move after successful write
-            self.process_queue.put_nowait((file_name, local_path))
+            # seek to the start of the footer and read the footer data
+            footer_start = file_size - footer_size - 8
+            f.seek(footer_start)
+            footer_data = f.read(file_size - footer_start)
+
+            if len(footer_data) != file_size - footer_start:
+                raise ValueError(f"Failed to read complete footer data from {file_name}")
+
+        # Parse the footer data to extract schema information
+        with io.BytesIO(footer_data) as footer_buffer:
+            parquet_file = pq.ParquetFile(footer_buffer)
+            data_types = [str(col.type) for col in parquet_file.schema_arrow]
+            num_rows = parquet_file.metadata.num_rows
+
+        return {
+            "file_name": file_name,
+            "num_rows": num_rows,
+            "file_size": file_size,
+            "data_types": data_types,
+        }
 
     def write_index(self, chunks_info: List[Dict[str, Any]], config: Dict[str, Any]) -> None:
+        """Write the index file to the local cache directory and upload it to the cloud."""
         assert self.cache_path is not None
         assert self.dir.url is not None
 
-        # clear any `.lock` or `.tmp` chunk file if left
-        if self.is_delete_thread_running:
-            for file in os.listdir(self.cache_path):
-                if file != _INDEX_FILENAME:
-                    with suppress(FileNotFoundError):
-                        os.remove(file)
-
         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
+        # upload index file to cloud
         with open(index_file_path, "rb") as local_file, self.fs.open(cloud_index_path, "wb") as cloud_file:
             for chunk in iter(lambda: local_file.read(4096), b""):  # Read in 4KB chunks
                 cloud_file.write(chunk)
 
-        print(f"Index file written to: {cloud_index_path}")
+        print(f"Index file successfully written to: {cloud_index_path}")
         if os.path.exists(index_file_path):
             os.remove(index_file_path)
 
@@ -221,15 +221,18 @@ class HFParquetDir(ParquetDir):
         dir_path: Optional[Union[str, Dir]],
         cache_path: Optional[str] = None,
         storage_options: Optional[Dict] = None,
-        remove_after_indexing: bool = True,
         num_workers: int = 4,
     ):
         if not _HF_HUB_AVAILABLE:
             raise ModuleNotFoundError(
                 "Support for Indexing HF depends on `huggingface_hub`.", "Please, run: `pip install huggingface_hub"
             )
-
-        super().__init__(dir_path, cache_path, storage_options, remove_after_indexing, num_workers)
+        if not _PYARROW_AVAILABLE:
+            raise ModuleNotFoundError(
+                "The 'pyarrow' module is required for processing Parquet files. "
+                "Please install it by running: `pip install pyarrow`"
+            )
+        super().__init__(dir_path, cache_path, storage_options, num_workers)
 
         assert self.dir.url is not None
         assert self.dir.url.startswith("hf")
@@ -243,66 +246,79 @@ class HFParquetDir(ParquetDir):
 
         self.fs = HfFileSystem()
 
-        for _f in self.fs.ls(self.dir.url, detail=False):
-            if _f.endswith(".parquet"):
+        for _f in self.fs.ls(self.dir.url, detail=True):
+            if isinstance(_f, dict) and _f["name"].endswith(".parquet"):
                 self.files.append(_f)
 
-    def task(self, _file: str) -> None:
-        assert isinstance(_file, str)
-        assert self.cache_path is not None
-        if _file.endswith(".parquet"):
-            file_name = os.path.basename(_file)
-            local_path = os.path.join(self.cache_path, file_name)
-            os.makedirs(os.path.dirname(local_path), exist_ok=True)
-            temp_path = local_path + ".tmp"  # Avoid partial writes
-            # if an existing temp file is present, means its corrupted
-            if os.path.exists(temp_path):
-                os.remove(temp_path)
+    def task(self, _file: dict) -> Dict[str, Any]:
+        """Extract metadata from a Parquet file on Hugging Face Hub without downloading the entire file."""
+        import pyarrow.parquet as pq
+
+        # Validate inputs
+        if not isinstance(_file, dict) or "name" not in _file or "size" not in _file:
+            raise ValueError(f"Invalid file object: {_file}")
 
-            with self.fs.open(_file, "rb") as cloud_file, open(temp_path, "wb") as local_file:
-                if _TQDM_AVAILABLE:
-                    from tqdm.auto import tqdm as _tqdm
+        assert _file["name"].endswith(".parquet")
 
-                    file_size = self.fs.info(_file)["size"]
-                    pbar = _tqdm(desc=f"Downloading {file_name}", total=file_size, unit="B", unit_scale=True)
+        file_name = os.path.basename(_file["name"])
+        file_size = _file["size"]
 
-                for chunk in iter(lambda: cloud_file.read(4096), b""):  # Read in 4KB chunks
-                    local_file.write(chunk)
+        with self.fs.open(_file["name"], "rb") as f:
+            # Read footer size (last 8 bytes: 4 for footer size + 4 for magic number)
+            f.seek(file_size - 8)
+            footer_size = int.from_bytes(f.read(4), "little")
 
-                    if _TQDM_AVAILABLE:
-                        pbar.update(len(chunk))
+            # seek to the start of the footer and read the footer data
+            footer_start = file_size - footer_size - 8
+            f.seek(footer_start)
+            footer_data = f.read(file_size - footer_start)
 
-            os.rename(temp_path, local_path)  # Atomic move after successful write
-            self.process_queue.put_nowait((file_name, local_path))
+            if len(footer_data) != file_size - footer_start:
+                raise ValueError(f"Failed to read complete footer data from {file_name}")
+
+        # Parse the footer data to extract schema information
+        with io.BytesIO(footer_data) as footer_buffer:
+            parquet_file = pq.ParquetFile(footer_buffer)
+            data_types = [str(col.type) for col in parquet_file.schema_arrow]
+            num_rows = parquet_file.metadata.num_rows
+
+        return {
+            "file_name": file_name,
+            "num_rows": num_rows,
+            "file_size": file_size,
+            "data_types": data_types,
+        }
 
     def write_index(self, chunks_info: List[Dict[str, Any]], config: Dict[str, Any]) -> None:
+        """Write the index file to the cache directory."""
         assert self.cache_path is not None
         assert self.dir.url is not None
 
-        # clear any `.lock` or `.tmp` chunk file if left
-        if self.is_delete_thread_running:
-            for file in os.listdir(self.cache_path):
-                if file != _INDEX_FILENAME:
-                    with suppress(FileNotFoundError):
-                        os.remove(file)
-
-        index_file_path = os.path.join(self.cache_path, _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)
-
-        print(f"Index file written to: {index_file_path}")
+        super().write_index(chunks_info, config)
 
 
 def get_parquet_indexer_cls(
     dir_path: str,
     cache_path: Optional[str] = None,
     storage_options: Optional[Dict] = {},
-    remove_after_indexing: bool = True,
     num_workers: int = 4,
 ) -> ParquetDir:
-    args = (dir_path, cache_path, storage_options, remove_after_indexing, num_workers)
+    """Get the appropriate ParquetDir class based on the directory path scheme.
+
+    Args:
+        dir_path (str): Path to the directory containing the Parquet files.
+        cache_path (Optional[str]): Local cache directory for storing temporary files.
+        storage_options (Optional[Dict]): Additional storage options for accessing the Parquet files.
+        remove_after_indexing (bool): Whether to remove files after indexing (default is True).
+        num_workers (int): Number of workers to download metadata of Parquet files and index them.
+
+    Returns:
+        ParquetDir: An instance of the appropriate ParquetDir class.
+
+    Raises:
+        ValueError: If the provided `dir_path` does not have an associated ParquetDir class.
+    """
+    args = (dir_path, cache_path, storage_options, num_workers)
 
     obj = parse.urlparse(dir_path)
 
@@ -313,18 +329,33 @@ def get_parquet_indexer_cls(
     if obj.scheme == "hf":
         return HFParquetDir(*args)
 
-    if sys.platform == "win32":
-        return LocalParquetDir(*args)
-
+    supported_schemes = ["local", "gs", "s3", "hf"]
     raise ValueError(
-        f"The provided `dir_path` {dir_path} doesn't have a ParquetDir class associated.",
-        f"Found scheme => {obj.scheme}",
+        f"The provided `dir_path` '{dir_path}' does not have an associated ParquetDir class. "
+        f"Found scheme: '{obj.scheme}'. Supported schemes are: {', '.join(supported_schemes)}. "
+        "Please provide a valid directory path with one of the supported schemes."
     )
 
 
 def default_cache_dir(url: str) -> str:
-    # Hash the URL using SHA256 and take the first 16 characters for brevity
+    """Generate a default cache directory path based on the given URL.
+
+    The directory is created under the user's home directory at
+    ~/.cache/litdata-cache-index-pq if it does not already exist.
+
+    Args:
+        url (str): The URL to be hashed for generating the cache directory path.
+
+    Returns:
+        str: The path to the generated cache directory.
+    """
+    # Hash the URL using SHA256
     url_hash = hashlib.sha256(url.encode()).hexdigest()
+
+    # Generate the cache directory path
     cache_path = os.path.join(os.path.expanduser("~"), ".cache", "litdata-cache-index-pq", url_hash)
-    os.makedirs(cache_path, exist_ok=True)  # Ensure the directory exists
+
+    # Ensure the directory exists
+    os.makedirs(cache_path, exist_ok=True)
+
     return cache_path


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 6143dee..ec3cd72 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -181,7 +181,17 @@ def clean_pq_index_cache():
 
 
 @pytest.fixture
-def huggingface_hub_mock(monkeypatch, write_pq_data, tmp_path):
+def huggingface_hub_mock(monkeypatch):
+    huggingface_hub = ModuleType("huggingface_hub")
+    monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub)
+    hf_hub_download = ModuleType("hf_hub_download")
+    monkeypatch.setitem(sys.modules, "huggingface_hub.hf_hub_download", hf_hub_download)
+    huggingface_hub.hf_hub_download = hf_hub_download
+    return huggingface_hub
+
+
+@pytest.fixture
+def huggingface_hub_fs_mock(monkeypatch, write_pq_data, tmp_path):
     huggingface_hub = ModuleType("huggingface_hub")
     hf_file_system = ModuleType("hf_file_system")
 
@@ -196,13 +206,25 @@ def huggingface_hub_mock(monkeypatch, write_pq_data, tmp_path):
         assert os.path.exists(file_path), "file hf is trying to access, doesn't exist."
         return open(file_path, mode)
 
+    def mock_ls(*args, **kwargs):
+        files = sorted(os.listdir(os.path.join(tmp_path, "pq-dataset")))
+        return [
+            {
+                "type": "file",
+                "name": file_name,
+                "size": os.path.getsize(os.path.join(tmp_path, "pq-dataset", file_name)),
+            }
+            for file_name in files
+        ]
+
+    def mock_hf_hub_download(repo_id, filename, cache_dir, repo_type, **kwargs):
+        return os.path.join(tmp_path, "pq-dataset", os.path.basename(filename))
+
     hf_fs_mock = Mock()
-    hf_fs_mock.ls = Mock(side_effect=lambda *args, **kwargs: os.listdir(os.path.join(tmp_path, "pq-dataset")))
+    hf_fs_mock.ls = Mock(side_effect=mock_ls)
     hf_fs_mock.open = Mock(side_effect=mock_open)
-    hf_fs_mock.info = Mock(
-        side_effect=lambda filename: {"size": os.path.getsize(os.path.join(tmp_path, "pq-dataset", filename))}
-    )
     huggingface_hub.HfFileSystem = Mock(return_value=hf_fs_mock)
+    huggingface_hub.hf_hub_download = Mock(side_effect=mock_hf_hub_download)
 
     return huggingface_hub
 
@@ -220,7 +242,13 @@ def fsspec_pq_mock(monkeypatch, write_pq_data, tmp_path, fsspec_mock):
     def mock_fsspec_ls(*args, **kwargs):
         file_list = []
         for file_name in os.listdir(os.path.join(tmp_path, "pq-dataset")):
-            file_list.append({"type": "file", "name": file_name})
+            file_list.append(
+                {
+                    "type": "file",
+                    "name": file_name,
+                    "size": os.path.getsize(os.path.join(tmp_path, "pq-dataset", file_name)),
+                }
+            )
         return file_list
 
     fs_mock = Mock()
diff --git a/tests/processing/test_readers.py b/tests/processing/test_readers.py
index ff4b33e..7b81809 100644
--- a/tests/processing/test_readers.py
+++ b/tests/processing/test_readers.py
@@ -4,7 +4,8 @@ import sys
 import pytest
 
 from litdata import map
-from litdata.processing.readers import _PYARROW_AVAILABLE, BaseReader, ParquetReader
+from litdata.constants import _PYARROW_AVAILABLE
+from litdata.processing.readers import BaseReader, ParquetReader
 
 
 class DummyReader(BaseReader):
diff --git a/tests/streaming/test_downloader.py b/tests/streaming/test_downloader.py
index 52bace5..2a9c144 100644
--- a/tests/streaming/test_downloader.py
+++ b/tests/streaming/test_downloader.py
@@ -1,4 +1,5 @@
 # ruff: noqa: S604
+import contextlib
 import os
 import sys
 from unittest import mock
@@ -11,6 +12,7 @@ from litdata.streaming.downloader import (
     AzureDownloader,
     Downloader,
     GCPDownloader,
+    HFDownloader,
     LocalDownloaderWithCache,
     S3Downloader,
     get_downloader,
@@ -312,3 +314,38 @@ def test_download_with_cache(tmpdir, monkeypatch):
         os_mock.assert_called()
     finally:
         os.remove("a.txt")
+
+
+@mock.patch("litdata.streaming.downloader._HF_HUB_AVAILABLE", True)
+def test_hf_downloader(tmpdir, huggingface_hub_mock):
+    # Create a mock for hf_hub_download
+    mock_hf_hub_download = MagicMock()
+    huggingface_hub_mock.hf_hub_download = mock_hf_hub_download
+
+    # Initialize the downloader
+    storage_options = {}
+    downloader = HFDownloader("hf://datasets/sample_org/sample_repo", tmpdir, [], storage_options)
+    local_filepath = os.path.join(tmpdir, "a.txt")
+
+    # Configure the mock to return the local_filepath
+    mock_hf_hub_download.return_value = local_filepath
+
+    # Test case 1: File doesn’t exist, should download
+    with contextlib.suppress(FileNotFoundError):
+        downloader.download_file("hf://datasets/sample_org/sample_repo/a.txt", local_filepath)
+
+    # Verify that hf_hub_download was called with the correct arguments
+    huggingface_hub_mock.hf_hub_download.assert_called_once()
+
+    # Reset the mock for the next test case
+    mock_hf_hub_download.reset_mock()
+
+    # Test case 2: File exists, should skip download
+    with open(local_filepath, "w") as f:
+        f.write("dummy content")
+
+    with contextlib.suppress(FileNotFoundError):
+        downloader.download_file("hf://datasets/sample_org/sample_repo/a.txt", local_filepath)
+
+    # Verify that hf_hub_download was not called
+    mock_hf_hub_download.assert_not_called()
diff --git a/tests/streaming/test_parquet.py b/tests/streaming/test_parquet.py
index 5a37870..1d41e1c 100644
--- a/tests/streaming/test_parquet.py
+++ b/tests/streaming/test_parquet.py
@@ -2,24 +2,22 @@ import hashlib
 import json
 import os
 import sys
-import threading
-import time
 from contextlib import nullcontext
-from queue import Queue
 from types import ModuleType
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import pytest
 
 from litdata.constants import _INDEX_FILENAME
 from litdata.streaming.dataset import StreamingDataset
+from litdata.streaming.item_loader import ParquetLoader, PyTreeLoader
 from litdata.streaming.writer import index_parquet_dataset
+from litdata.utilities.hf_dataset import index_hf_dataset
 from litdata.utilities.parquet import (
     CloudParquetDir,
     HFParquetDir,
     LocalParquetDir,
     default_cache_dir,
-    delete_thread,
     get_parquet_indexer_cls,
 )
 
@@ -33,17 +31,16 @@ from litdata.utilities.parquet import (
         None,
         "s3://some_bucket/some_path",
         "gs://some_bucket/some_path",
-        "hf://some_bucket/some_path",
+        "hf://datasets/some_org/some_repo/some_path",
     ],
 )
-@pytest.mark.parametrize(("remove_after_indexing"), [True, False])
 @pytest.mark.parametrize(("num_worker"), [None, 1, 2, 4])
+@patch("litdata.utilities.parquet._HF_HUB_AVAILABLE", True)
+@patch("litdata.streaming.downloader._HF_HUB_AVAILABLE", True)
+@patch("litdata.utilities.parquet._FSSPEC_AVAILABLE", True)
 def test_parquet_index_write(
-    monkeypatch, tmp_path, pq_data, huggingface_hub_mock, fsspec_pq_mock, pq_dir_url, remove_after_indexing, num_worker
+    monkeypatch, tmp_path, pq_data, huggingface_hub_fs_mock, fsspec_pq_mock, pq_dir_url, num_worker
 ):
-    monkeypatch.setattr("litdata.utilities.parquet._HF_HUB_AVAILABLE", True)
-    monkeypatch.setattr("litdata.utilities.parquet._FSSPEC_AVAILABLE", True)
-
     if pq_dir_url is None:
         pq_dir_url = os.path.join(tmp_path, "pq-dataset")
 
@@ -57,30 +54,24 @@ def test_parquet_index_write(
 
     # call the write_parquet_index fn
     if num_worker is None:
-        index_parquet_dataset(pq_dir_url=pq_dir_url, remove_after_indexing=remove_after_indexing)
+        index_parquet_dataset(pq_dir_url=pq_dir_url)
     else:
-        index_parquet_dataset(
-            pq_dir_url=pq_dir_url, remove_after_indexing=remove_after_indexing, num_workers=num_worker
-        )
+        index_parquet_dataset(pq_dir_url=pq_dir_url, num_workers=num_worker)
 
     assert os.path.exists(index_file_path)
 
-    if remove_after_indexing and pq_dir_url.startswith(("gs://", "s3://", "hf://")):
-        time.sleep(1)  # Small delay to ensure deletion completes
-        if pq_dir_url.startswith("hf://"):
-            assert len(os.listdir(cache_dir)) == 1
-        else:
-            assert len(os.listdir(cache_dir)) == 0
+    if pq_dir_url.startswith("hf://"):
+        assert len(os.listdir(cache_dir)) == 1
+    elif pq_dir_url.startswith(("gs://", "s3://")):
+        assert len(os.listdir(cache_dir)) == 0
 
     # Read JSON file into a dictionary
     with open(index_file_path) as f:
         data = json.load(f)
-        print(f"index.json: {data}")
         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"]
 
     # no test for streaming on s3 and gs
     if pq_dir_url is None or pq_dir_url.startswith("hf://"):
@@ -96,6 +87,20 @@ def test_parquet_index_write(
             assert _ds[2] == pq_data["height"][idx]
 
 
+@pytest.mark.skipif(condition=sys.platform == "win32", reason="Fails on windows and test gets cancelled")
+@pytest.mark.usefixtures("clean_pq_index_cache")
+@patch("litdata.utilities.parquet._HF_HUB_AVAILABLE", True)
+def test_index_hf_dataset(monkeypatch, tmp_path, huggingface_hub_fs_mock):
+    with pytest.raises(ValueError, match="Invalid Hugging Face dataset URL"):
+        index_hf_dataset("invalid_url")
+
+    hf_url = "hf://datasets/some_org/some_repo/some_path"
+    cache_dir = index_hf_dataset(hf_url)
+    assert os.path.exists(cache_dir)
+    assert len(os.listdir(cache_dir)) == 1
+    assert os.path.exists(os.path.join(cache_dir, _INDEX_FILENAME))
+
+
 def test_default_cache_dir(monkeypatch):
     os = ModuleType("os")
     os.path = Mock()
@@ -140,7 +145,7 @@ def test_default_cache_dir(monkeypatch):
         ("meow://some_bucket/somepath", None, pytest.raises(ValueError, match="The provided")),
     ],
 )
-def test_get_parquet_indexer_cls(pq_url, cls, expectation, monkeypatch, fsspec_mock, huggingface_hub_mock):
+def test_get_parquet_indexer_cls(pq_url, cls, expectation, monkeypatch, fsspec_mock, huggingface_hub_fs_mock):
     os = Mock()
     os.listdir = Mock(return_value=[])
 
@@ -150,7 +155,7 @@ def test_get_parquet_indexer_cls(pq_url, cls, expectation, monkeypatch, fsspec_m
 
     hf_fs_mock = Mock()
     hf_fs_mock.ls = Mock(return_value=[])
-    huggingface_hub_mock.HfFileSystem = Mock(return_value=hf_fs_mock)
+    huggingface_hub_fs_mock.HfFileSystem = Mock(return_value=hf_fs_mock)
 
     monkeypatch.setattr("litdata.utilities.parquet.os", os)
     monkeypatch.setattr("litdata.utilities.parquet._HF_HUB_AVAILABLE", True)
@@ -160,44 +165,46 @@ def test_get_parquet_indexer_cls(pq_url, cls, expectation, monkeypatch, fsspec_m
         assert isinstance(indexer_obj, cls)
 
 
-def test_delete_thread(tmp_path):
-    # create some random files in tmp_path dir and test if delete_thread deletes them
-    tmp_files = [f"tmp_file_{i}.txt" for i in range(10)]
-    for f in tmp_files:
-        with open(os.path.join(tmp_path, f), "w") as _:
-            pass
-
-    for f in tmp_files:
-        file_path = tmp_path / f
-        assert os.path.isfile(file_path)
-
-    q = Queue()
-    dt = threading.Thread(target=delete_thread, args=(q,))
-    dt.start()
-
-    for f in tmp_files:
-        file_path = tmp_path / f
-        q.put_nowait(file_path)
-
-    q.put_nowait(None)
-    dt.join()
-
-    # check if all files are deleted
-    for f in tmp_files:
-        file_path = tmp_path / f
-        assert not os.path.exists(file_path)
+@pytest.mark.usefixtures("clean_pq_index_cache")
+@patch("litdata.utilities.parquet._HF_HUB_AVAILABLE", True)
+@patch("litdata.streaming.downloader._HF_HUB_AVAILABLE", True)
+def test_stream_hf_parquet_dataset(monkeypatch, huggingface_hub_fs_mock, pq_data):
+    hf_url = "hf://datasets/some_org/some_repo/some_path"
 
+    # Test case 1: Invalid item_loader
+    with pytest.raises(ValueError, match="Invalid item_loader for hf://datasets."):
+        StreamingDataset(hf_url, item_loader=PyTreeLoader)
 
-@pytest.mark.usefixtures("clean_pq_index_cache")
-def test_stream_hf_parquet_dataset(huggingface_hub_mock, pq_data):
-    hf_url = "hf://some_bucket/some_path"
+    # Test case 2: Streaming without passing item_loader
     ds = StreamingDataset(hf_url)
-
     assert len(ds) == 25  # 5 datasets for 5 loops
+    for i, _ds in enumerate(ds):
+        idx = i % 5
+        assert len(_ds) == 3
+        assert _ds[0] == pq_data["name"][idx]
+        assert _ds[1] == pq_data["weight"][idx]
+        assert _ds[2] == pq_data["height"][idx]
 
+    # Test case 3: Streaming with ParquetLoader as item_loader and low_memory=False
+    ds = StreamingDataset(hf_url, item_loader=ParquetLoader(low_memory=False))
+    assert len(ds) == 25
     for i, _ds in enumerate(ds):
         idx = i % 5
         assert len(_ds) == 3
         assert _ds[0] == pq_data["name"][idx]
         assert _ds[1] == pq_data["weight"][idx]
         assert _ds[2] == pq_data["height"][idx]
+
+    # Test case 4: Streaming with ParquetLoader and low_memory=True
+    ds = StreamingDataset(hf_url, item_loader=ParquetLoader(low_memory=True))
+    assert len(ds) == 25
+    for i, _ds in enumerate(ds):
+        idx = i % 5
+        assert len(_ds) == 3
+        assert _ds[0] == pq_data["name"][idx]
+        assert _ds[1] == pq_data["weight"][idx]
+        assert _ds[2] == pq_data["height"][idx]
+
+    # Test case 5: Streaming with ParquetLoader and low_memory=True and shuffle=True
+    with pytest.raises(ValueError, match="You have enabled shuffling when using low memory with ParquetLoader."):
+        StreamingDataset(hf_url, item_loader=ParquetLoader(low_memory=True), shuffle=True)