perf: faster fs data loading

Author: percevalw

edsnlp/data/json.py

@@ -36,7 +36,9 @@ def __init__(
         self.files = (
             [
                 file
-                for file in walk_match(self.fs, os.path.dirname(self.path), "*.json*")
+                for file in walk_match(
+                    self.fs, os.path.dirname(self.path), ".*[.]json.*"
+                )
                 if keep_ipynb_checkpoints or ".ipynb_checkpoints" not in str(file)
             ]
             if self.fs.isdir(self.path)
@@ -128,7 +130,7 @@ def __init__(
         )
         if path_exists:
             if self.fs.isdir(self.path):
-                files = [f for f in walk_match(self.fs, self.path, "*.json*")]
+                files = [f for f in walk_match(self.fs, self.path, ".*[.]json.*")]
                 if files:
                     if not overwrite:
                         raise FileExistsError(

edsnlp/data/parquet.py

@@ -83,15 +83,15 @@ def __init__(
         self.fs, self.path = normalize_fs_path(filesystem, path)
 
         # Check that filesystem has the same protocol as indicated by path
-        self.fs.makedir(self.path, create_parents=True)
+        self.fs.makedirs(self.path, exist_ok=True)
+
         dataset: pyarrow.dataset.FileSystemDataset = (  # type: ignore
             pyarrow.dataset.dataset(
                 self.path,
                 format="parquet",
                 filesystem=self.fs,
             )
         )
-        self.fs.makedirs(self.path, exist_ok=True)
         if len(list(dataset.files)):
             if not overwrite:
                 raise FileExistsError(

edsnlp/data/standoff.py

@@ -19,13 +19,11 @@
 from loguru import logger
 
 from edsnlp import registry
-from edsnlp.core import PipelineProtocol
 from edsnlp.core.lazy_collection import LazyCollection
 from edsnlp.data.base import BaseReader, BaseWriter
 from edsnlp.data.converters import (
     FILENAME,
     AttributesMappingArg,
-    SequenceStr,
     get_dict2doc_converter,
     get_doc2dict_converter,
 )
@@ -50,7 +48,8 @@ def __init__(self, ann_file, line):
 
 
 def parse_standoff_file(
-    path: str,
+    txt_path: str,
+    ann_paths: List[str],
     merge_spaced_fragments: bool = True,
     fs: FileSystem = LOCAL_FS,
 ) -> Dict:
@@ -74,28 +73,19 @@ def parse_standoff_file(
     -------
     Iterator[Dict]
     """
-    ann_filenames = []
-    for filename in walk_match(
-        fs,
-        os.path.dirname(path),
-        os.path.basename(path).replace(".txt", ".a*"),
-        recursive=False,
-    ):
-        ann_filenames.append(filename)
-
     entities = {}
     relations = []
     events = {}
 
-    with fs.open(path, "r") as f:
+    with fs.open(txt_path, "r") as f:
         text = f.read()
 
-    if not len(ann_filenames):
+    if not len(ann_paths):
         return {
             "text": text,
         }
 
-    for ann_file in ann_filenames:
+    for ann_file in ann_paths:
         with fs.open(ann_file, "r") as f:
             for line_idx, line in enumerate(f):
                 try:
@@ -303,34 +293,33 @@ def __init__(
     ):
         super().__init__()
         self.fs, self.path = normalize_fs_path(filesystem, path)
-        self.files: List[str] = [
+        files = {
             file
-            for file in walk_match(self.fs, self.path, "*.txt")
+            for file in walk_match(self.fs, self.path, ".*[.](txt|a*)")
             if (keep_ipynb_checkpoints or ".ipynb_checkpoints" not in str(file))
-            and (
-                keep_txt_only_docs
-                or walk_match(
-                    self.fs,
-                    os.path.dirname(file),
-                    os.path.basename(file).replace(".txt", ".a*"),
-                    recursive=False,
-                )
-            )
+        }
+        ann_files = {}
+        for f in files:
+            name, ext = os.path.splitext(f)
+            if ext.startswith(".a"):
+                ann_files.setdefault(name, []).append(f)
+        self.files = [
+            (file, ann_files.get(file.replace(".txt", ""), []))
+            for file in files
+            if file.endswith(".txt")
+            and (keep_txt_only_docs or file.replace(".txt", "") in ann_files)
         ]
         assert len(self.files), f"No .txt files found in the BRAT directory {self.path}"
-        for file in self.files:
-            if not self.fs.exists(file):
-                raise FileNotFoundError(f"File {file} does not exist")
         logger.info(f"The BRAT directory contains {len(self.files)} .txt files.")
 
     def read_main(self) -> Iterable[Tuple[str, int]]:
         return ((f, 1) for f in self.files)
 
     def read_worker(self, fragment: List[str]):
         tasks = []
-        for file in fragment:
-            anns = parse_standoff_file(str(file), fs=self.fs)
-            anns[FILENAME] = os.path.relpath(file, self.path).rsplit(".", 1)[0]
+        for txt_path, ann_paths in fragment:
+            anns = parse_standoff_file(txt_path, ann_paths, fs=self.fs)
+            anns[FILENAME] = os.path.relpath(txt_path, self.path).rsplit(".", 1)[0]
             anns["doc_id"] = anns[FILENAME]
             tasks.append(anns)
         return tasks
@@ -350,9 +339,8 @@ def __init__(
 
         if self.fs.exists(self.path):
             unsafe_exts = Counter(
-                os.path.splitext(f)[1] for f in walk_match(self.fs, self.path, "*.txt")
-            ) + Counter(
-                os.path.splitext(f)[1] for f in walk_match(self.fs, self.path, "*.a*")
+                os.path.splitext(f)[1]
+                for f in walk_match(self.fs, self.path, ".*[.](txt|a.*)")
             )
             if unsafe_exts and not overwrite:
                 raise FileExistsError(

edsnlp/utils/file_system.py

@@ -1,8 +1,9 @@
-import fnmatch
 import os
+import re
 from pathlib import Path
 from typing import Optional, Tuple, Union
 
+import fsspec.implementations.local
 import pyarrow.fs
 from fsspec import AbstractFileSystem
 from fsspec import __version__ as fsspec_version
@@ -22,19 +23,12 @@ def walk_match(
     fs: FileSystem,
     root: str,
     file_pattern: str,
-    recursive: bool = True,
 ) -> list:
-    if fsspec_version >= "2023.10.0":
-        # Version fixes fsspec glob https://github.com/fsspec/filesystem_spec/pull/1329
-        glob_str = os.path.join(root, "**" if recursive else "", file_pattern)
-        return fs.glob(glob_str)
     return [
         os.path.join(dirpath, f)
-        for dirpath, dirnames, files in fs.walk(
-            root,
-            maxdepth=None if recursive else 1,
-        )
-        for f in fnmatch.filter(files, file_pattern)
+        for dirpath, dirnames, files in fs.walk(root)
+        for f in files
+        if re.match(file_pattern, f)
     ]
 
 
@@ -43,15 +37,21 @@ def normalize_fs_path(
     path: Union[str, Path],
 ) -> Tuple[AbstractFileSystem, str]:
     has_protocol = isinstance(path, str) and "://" in path
+    filesystem = (
+        ArrowFSWrapper(filesystem)
+        if isinstance(filesystem, pyarrow.fs.FileSystem)
+        else filesystem
+    )
 
     # We need to detect the fs from the path
     if filesystem is None or has_protocol:
         uri: str = path if has_protocol else f"file://{os.path.abspath(path)}"
-        inferred_fs, fs_path = pyarrow.fs.FileSystem.from_uri(uri)
+        inferred_fs, fs_path = fsspec.core.url_to_fs(uri)
+        inferred_fs: fsspec.AbstractFileSystem
         filesystem = filesystem or inferred_fs
-        assert inferred_fs.type_name == filesystem.type_name, (
-            f"Protocol {inferred_fs.type_name} in path does not match "
-            f"filesystem {filesystem.type_name}"
+        assert inferred_fs.protocol == filesystem.protocol, (
+            f"Protocol {inferred_fs.protocol} in path does not match "
+            f"filesystem {filesystem.protocol}"
         )
         path = fs_path  # path without protocol