Prevent SimpleDirectoryReader from excessive memory consumption (#18983)

Author: AstraBert

llama-index-core/llama_index/core/readers/file/base.py

@@ -14,6 +14,7 @@
 from itertools import repeat
 from pathlib import Path, PurePosixPath
 from typing import (
+    Optional,
     Any,
     Callable,
     Generator,
@@ -247,19 +248,19 @@ class SimpleDirectoryReader(BaseReader, ResourcesReaderMixin, FileSystemReaderMi
 
     def __init__(
         self,
-        input_dir: Path | str | None = None,
-        input_files: list | None = None,
-        exclude: list | None = None,
+        input_dir: Optional[Union[Path, str]] = None,
+        input_files: Optional[list] = None,
+        exclude: Optional[list] = None,
         exclude_hidden: bool = True,
         exclude_empty: bool = False,
         errors: str = "ignore",
         recursive: bool = False,
         encoding: str = "utf-8",
         filename_as_id: bool = False,
-        required_exts: list[str] | None = None,
-        file_extractor: dict[str, BaseReader] | None = None,
-        num_files_limit: int | None = None,
-        file_metadata: Callable[[str], dict] | None = None,
+        required_exts: Optional[list[str]] = None,
+        file_extractor: Optional[dict[str, BaseReader]] = None,
+        num_files_limit: Optional[int] = None,
+        file_metadata: Optional[Callable[[str], dict]] = None,
         raise_on_error: bool = False,
         fs: fsspec.AbstractFileSystem | None = None,
     ) -> None:
@@ -333,10 +334,21 @@ def _add_files(self, input_dir: Path | PurePosixPath) -> list[Path | PurePosixPa
                         rejected_files.add(_Path(str(file)))
 
         file_refs: list[str] = []
-        if self.recursive:
-            file_refs = cast(list[str], self.fs.glob(str(input_dir) + "/**/*"))
-        else:
-            file_refs = cast(list[str], self.fs.glob(str(input_dir) + "/*"))
+        limit = (
+            self.num_files_limit
+            if self.num_files_limit is not None and self.num_files_limit > 0
+            else None
+        )
+        c = 0
+        depth = 1000 if self.recursive else 1
+        for root, _, files in self.fs.walk(
+            str(input_dir), topdown=True, maxdepth=depth
+        ):
+            for file in files:
+                c += 1
+                if limit and c > limit:
+                    break
+                file_refs.append(os.path.join(root, file))
 
         for _ref in file_refs:
             # Manually check if file is hidden or directory instead of
@@ -381,9 +393,6 @@ def _add_files(self, input_dir: Path | PurePosixPath) -> list[Path | PurePosixPa
         if len(new_input_files) == 0:
             raise ValueError(f"No files found in {input_dir}.")
 
-        if self.num_files_limit is not None and self.num_files_limit > 0:
-            new_input_files = new_input_files[0 : self.num_files_limit]
-
         # print total number of files added
         logger.debug(
             f"> [SimpleDirectoryReader] Total files added: {len(new_input_files)}"

llama-index-core/tests/readers/file/test_base.py

@@ -99,7 +99,7 @@ def test_SimpleDirectoryReader_empty(data_path):
 
 def test_SimpleDirectoryReader_file_limit(data_path):
     r = SimpleDirectoryReader(input_dir=data_path, recursive=True, num_files_limit=2)
-    assert [f.name for f in r.input_files] == ["excluded_1.txt", "excluded_0.txt"]
+    assert [f.name for f in r.input_files] == ["excluded_0.txt", "file_0.md"]
 
 
 def test_SimpleDirectoryReader_list_resources(data_path):

llama-index-integrations/readers/llama-index-readers-file/tests/test_file.py

@@ -179,7 +179,7 @@ def test_num_files_limit(tmp_dir_type: Type[Union[Path, str]]) -> None:
                     assert len(reader.input_files) == 2
                     assert set(input_file_names) == {
                         "test1.txt",
-                        "test2.txt",
+                        "test3.txt",
                     }
 
                     reader = SimpleDirectoryReader(