add aload_data to BaseReader and add implementation for SimpleDirectoryReader (#12156)

docs/docs/examples/data_connectors/simple_directory_reader.ipynb

@@ -205,38 +205,23 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "--2024-03-07 13:17:40--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
-      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.109.133, 185.199.110.133, 185.199.111.133, ...\n",
-      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.109.133|:443... connected.\n",
-      "HTTP request sent, awaiting response... 200 OK\n",
-      "Length: 75042 (73K) [text/plain]\n",
-      "Saving to: ‘data/paul_graham/paul_graham_essay1.txt’\n",
+      "--2024-03-21 16:03:25--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
+      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8003::154, 2606:50c0:8000::154, 2606:50c0:8001::154, ...\n",
+      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8003::154|:443... connected.\n",
+      "HTTP request sent, awaiting response... 404 Not Found\n",
+      "2024-03-21 16:03:25 ERROR 404: Not Found.\n",
       "\n",
-      "data/paul_graham/pa 100%[===================>]  73.28K  --.-KB/s    in 0.01s   \n",
+      "--2024-03-21 16:03:25--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
+      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8003::154, 2606:50c0:8000::154, 2606:50c0:8001::154, ...\n",
+      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8003::154|:443... connected.\n",
+      "HTTP request sent, awaiting response... 404 Not Found\n",
+      "2024-03-21 16:03:25 ERROR 404: Not Found.\n",
       "\n",
-      "2024-03-07 13:17:40 (6.64 MB/s) - ‘data/paul_graham/paul_graham_essay1.txt’ saved [75042/75042]\n",
-      "\n",
-      "--2024-03-07 13:17:40--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
-      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.109.133, 185.199.110.133, 185.199.111.133, ...\n",
-      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.109.133|:443... connected.\n",
-      "HTTP request sent, awaiting response... 200 OK\n",
-      "Length: 75042 (73K) [text/plain]\n",
-      "Saving to: ‘data/paul_graham/paul_graham_essay2.txt’\n",
-      "\n",
-      "data/paul_graham/pa 100%[===================>]  73.28K  --.-KB/s    in 0.01s   \n",
-      "\n",
-      "2024-03-07 13:17:40 (5.62 MB/s) - ‘data/paul_graham/paul_graham_essay2.txt’ saved [75042/75042]\n",
-      "\n",
-      "--2024-03-07 13:17:40--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
-      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.109.133, 185.199.110.133, 185.199.111.133, ...\n",
-      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.109.133|:443... connected.\n",
-      "HTTP request sent, awaiting response... 200 OK\n",
-      "Length: 75042 (73K) [text/plain]\n",
-      "Saving to: ‘data/paul_graham/paul_graham_essay3.txt’\n",
-      "\n",
-      "data/paul_graham/pa 100%[===================>]  73.28K  --.-KB/s    in 0.01s   \n",
-      "\n",
-      "2024-03-07 13:17:40 (6.27 MB/s) - ‘data/paul_graham/paul_graham_essay3.txt’ saved [75042/75042]\n",
+      "--2024-03-21 16:03:25--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
+      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8003::154, 2606:50c0:8000::154, 2606:50c0:8001::154, ...\n",
+      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8003::154|:443... connected.\n",
+      "HTTP request sent, awaiting response... 404 Not Found\n",
+      "2024-03-21 16:03:25 ERROR 404: Not Found.\n",
       "\n"
      ]
     }
@@ -417,6 +402,41 @@
     "print(len(all_docs))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Async execution is available through `aload_data`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "4\n"
+     ]
+    }
+   ],
+   "source": [
+    "import nest_asyncio\n",
+    "\n",
+    "nest_asyncio.apply()\n",
+    "\n",
+    "reader = SimpleDirectoryReader(\n",
+    "    input_dir=\"./data\",\n",
+    "    recursive=True,\n",
+    ")\n",
+    "\n",
+    "all_docs = await reader.aload_data()\n",
+    "\n",
+    "print(len(all_docs))"
+   ]
+  },
   {
    "attachments": {},
    "cell_type": "markdown",
@@ -436,10 +456,9 @@
     "class SimpleDirectoryReader(BaseReader):\n",
     "    \"\"\"Simple directory reader.\n",
     "\n",
-    "    Load files from file directory. \n",
+    "    Load files from file directory.\n",
     "    Automatically select the best file reader given file extensions.\n",
     "\n",
-    "\n",
     "    Args:\n",
     "        input_dir (str): Path to the directory.\n",
     "        input_files (List): List of file paths to read\n",
@@ -449,7 +468,7 @@
     "        encoding (str): Encoding of the files.\n",
     "            Default is utf-8.\n",
     "        errors (str): how encoding and decoding errors are to be handled,\n",
-    "                see https://docs.python.org/3/library/functions.html#open\n",
+    "              see https://docs.python.org/3/library/functions.html#open\n",
     "        recursive (bool): Whether to recursively search in subdirectories.\n",
     "            False by default.\n",
     "        filename_as_id (bool): Whether to use the filename as the document id.\n",
@@ -464,7 +483,10 @@
     "        file_metadata (Optional[Callable[str, Dict]]): A function that takes\n",
     "            in a filename and returns a Dict of metadata for the Document.\n",
     "            Default is None.\n",
-    "\"\"\"\n",
+    "        fs (Optional[fsspec.AbstractFileSystem]): File system to use. Defaults\n",
+    "        to using the local file system. Can be changed to use any remote file system\n",
+    "        exposed via the fsspec interface.\n",
+    "    \"\"\"\n",
     "```\n"
    ]
   }

llama-index-core/llama_index/core/async_utils.py

 """Async utils."""
+
 import asyncio
 from itertools import zip_longest
-from typing import Any, Coroutine, Iterable, List
+from typing import Any, Coroutine, Iterable, List, TypeVar
 
 
 def asyncio_module(show_progress: bool = False) -> Any:
@@ -80,12 +81,14 @@ def get_asyncio_module(show_progress: bool = False) -> Any:
 
 DEFAULT_NUM_WORKERS = 4
 
+T = TypeVar("T")
+
 
 async def run_jobs(
-    jobs: List[Coroutine],
+    jobs: List[Coroutine[Any, Any, T]],
     show_progress: bool = False,
     workers: int = DEFAULT_NUM_WORKERS,
-) -> List[Any]:
+) -> List[T]:
     """Run jobs.
 
     Args:

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

@@ -25,10 +25,21 @@ class BaseReader(ABC):
             f"{self.__class__.__name__} does not provide lazy_load_data method currently"
         )
 
+    async def alazy_load_data(
+        self, *args: Any, **load_kwargs: Any
+    ) -> Iterable[Document]:
+        """Load data from the input directory lazily."""
+        # Fake async - just calls the sync method. Override in subclasses for real async implementations.
+        return self.lazy_load_data(*args, **load_kwargs)
+
     def load_data(self, *args: Any, **load_kwargs: Any) -> List[Document]:
         """Load data from the input directory."""
         return list(self.lazy_load_data(*args, **load_kwargs))
 
+    async def aload_data(self, *args: Any, **load_kwargs: Any) -> List[Document]:
+        """Load data from the input directory."""
+        return self.load_data(*args, **load_kwargs)
+
     def load_langchain_documents(self, **load_kwargs: Any) -> List["LCDocument"]:
         """Load data in LangChain document format."""
         docs = self.load_data(**load_kwargs)

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

@@ -7,6 +7,7 @@ import multiprocessing
 import warnings
 from datetime import datetime
 from functools import reduce
+import asyncio
 from itertools import repeat
 from pathlib import Path, PurePosixPath
 import fsspec
@@ -14,6 +15,7 @@ from fsspec.implementations.local import LocalFileSystem
 from typing import Any, Callable, Dict, Generator, List, Optional, Type
 
 from llama_index.core.readers.base import BaseReader
+from llama_index.core.async_utils import run_jobs, get_asyncio_module
 from llama_index.core.schema import Document
 from tqdm import tqdm
 
@@ -434,6 +436,67 @@ class SimpleDirectoryReader(BaseReader):
 
         return documents
 
+    async def aload_file(self, input_file: Path) -> List[Document]:
+        """Load file asynchronously."""
+        # TODO: make this less redundant
+        default_file_reader_cls = SimpleDirectoryReader.supported_suffix_fn()
+        default_file_reader_suffix = list(default_file_reader_cls.keys())
+        metadata: Optional[dict] = None
+        documents: List[Document] = []
+
+        if self.file_metadata is not None:
+            metadata = self.file_metadata(str(input_file))
+
+        file_suffix = input_file.suffix.lower()
+        if (
+            file_suffix in default_file_reader_suffix
+            or file_suffix in self.file_extractor
+        ):
+            # use file readers
+            if file_suffix not in self.file_extractor:
+                # instantiate file reader if not already
+                reader_cls = default_file_reader_cls[file_suffix]
+                self.file_extractor[file_suffix] = reader_cls()
+            reader = self.file_extractor[file_suffix]
+
+            # load data -- catch all errors except for ImportError
+            try:
+                kwargs = {"extra_info": metadata}
+                if self.fs and not is_default_fs(self.fs):
+                    kwargs["fs"] = self.fs
+                docs = await reader.aload_data(input_file, **kwargs)
+            except ImportError as e:
+                # ensure that ImportError is raised so user knows
+                # about missing dependencies
+                raise ImportError(str(e))
+            except Exception as e:
+                # otherwise, just skip the file and report the error
+                print(
+                    f"Failed to load file {input_file} with error: {e}. Skipping...",
+                    flush=True,
+                )
+                return []
+
+            # iterate over docs if needed
+            if self.filename_as_id:
+                for i, doc in enumerate(docs):
+                    doc.id_ = f"{input_file!s}_part_{i}"
+
+            documents.extend(docs)
+        else:
+            # do standard read
+            fs = self.fs or get_default_fs()
+            with fs.open(input_file, errors=self.errors, encoding=self.encoding) as f:
+                data = f.read().decode(self.encoding, errors=self.errors)
+
+            doc = Document(text=data, metadata=metadata or {})
+            if self.filename_as_id:
+                doc.id_ = str(input_file)
+
+            documents.append(doc)
+
+        return documents
+
     def load_data(
         self,
         show_progress: bool = False,
@@ -497,6 +560,40 @@ class SimpleDirectoryReader(BaseReader):
 
         return self._exclude_metadata(documents)
 
+    async def aload_data(
+        self,
+        show_progress: bool = False,
+        num_workers: Optional[int] = None,
+        fs: Optional[fsspec.AbstractFileSystem] = None,
+    ) -> List[Document]:
+        """Load data from the input directory.
+
+        Args:
+            show_progress (bool): Whether to show tqdm progress bars. Defaults to False.
+            num_workers  (Optional[int]): Number of workers to parallelize data-loading over.
+            fs (Optional[fsspec.AbstractFileSystem]): File system to use. If fs was specified
+                in the constructor, it will override the fs parameter here.
+
+        Returns:
+            List[Document]: A list of documents.
+        """
+        files_to_process = self.input_files
+        fs = fs or self.fs
+
+        coroutines = [self.aload_file(input_file) for input_file in files_to_process]
+        if num_workers:
+            document_lists = await run_jobs(
+                coroutines, show_progress=show_progress, workers=num_workers
+            )
+        elif show_progress:
+            _asyncio = get_asyncio_module(show_progress=show_progress)
+            document_lists = await _asyncio.gather(*coroutines)
+        else:
+            document_lists = await asyncio.gather(*coroutines)
+        documents = [doc for doc_list in document_lists for doc in doc_list]
+
+        return self._exclude_metadata(documents)
+
     def iter_data(
         self, show_progress: bool = False
     ) -> Generator[List[Document], Any, Any]: