Merge pull request #2901 from mabel-dev/#2900

New disk listing code

Author: joocer

opteryx/__version__.py

@@ -1,9 +1,9 @@
 # THIS FILE IS AUTOMATICALLY UPDATED DURING THE BUILD PROCESS
 # DO NOT EDIT THIS FILE DIRECTLY
 
-__build__ = 1739
+__build__ = 1740
 __author__ = "@joocer"
-__version__ = "0.26.2-beta.1739"
+__version__ = "0.26.2-beta.1740"
 
 # Store the version here so:
 # 1) we don't load dependencies by storing it in __init__.py

opteryx/compiled/io/disk_reader.pyx

@@ -10,25 +10,39 @@
 Ultra-fast disk reader module
 """
 
+import errno
+
 from cpython.buffer cimport PyBuffer_FillInfo
-from libc.stdlib cimport free
+from cpython.mem cimport PyMem_Free
+from cpython.mem cimport PyMem_Malloc
+from cpython.unicode cimport PyUnicode_FromString
+from libc.stddef cimport size_t
 
 cdef extern from "disk_io.h":
     int read_all_pread(const char* path, unsigned char* dst, size_t* out_len,
                        bint sequential, bint willneed, bint drop_after)
     int read_all_mmap(const char* path, unsigned char** dst, size_t* out_len)
     int unmap_memory_c(unsigned char* addr, size_t size)
 
+cdef extern from "directories.h":
+    ctypedef struct file_info_t:
+        char* name
+        int is_directory
+        int is_regular_file
+        long long size
+        long long mtime
+
+    int list_directory_c "list_directory"(const char* path, file_info_t** files, size_t* count)
+    void free_file_list(file_info_t* files, size_t count)
+    int list_matching_files_c "list_matching_files"(const char* base_path, const char** extensions,
+                                                    size_t ext_count, char*** files, size_t* count) nogil
+    void free_file_names(char** files, size_t count)
+
 cdef class MappedMemory:
     cdef unsigned char* data
     cdef size_t size
     cdef bint owned
 
-    def __dealloc__(self):
-        if self.owned and self.data != NULL:
-            # Free the allocated memory (for non-mmap case)
-            free(self.data)
-
     def __getbuffer__(self, Py_buffer* buffer, int flags):
         PyBuffer_FillInfo(buffer, self, self.data, self.size, 1, flags)
 
@@ -64,6 +78,104 @@ def read_file(str path, bint sequential=True, bint willneed=True, bint drop_afte
     return memoryview(buf)[:out_len]
 
 
+def list_directory(str path):
+    """Return directory entries using the native file system scanner."""
+
+    path_b = path.encode("utf-8")
+    cdef const char* c_path = path_b
+    cdef file_info_t* files = NULL
+    cdef size_t count = 0
+
+    cdef int rc = list_directory_c(c_path, &files, &count)
+    if rc != 0:
+        err = -rc
+        if err == errno.ENOENT:
+            raise FileNotFoundError(path)
+        raise OSError(err, f"Failed to list directory: {path}")
+
+    entries = []
+    cdef file_info_t entry
+    cdef char* name_ptr
+    cdef size_t idx
+    try:
+        for idx in range(count):
+            entry = files[idx]
+            name_ptr = entry.name
+            if name_ptr == NULL:
+                continue
+
+            py_name = PyUnicode_FromString(name_ptr)
+            entries.append(
+                (
+                    py_name,
+                    bool(entry.is_directory),
+                    bool(entry.is_regular_file),
+                    entry.size,
+                    entry.mtime,
+                )
+            )
+    finally:
+        if files != NULL:
+            free_file_list(files, count)
+
+    return entries
+
+
+def list_files(str path, extensions):
+    """Return a list of files under ``path`` matching provided extensions."""
+
+    if extensions is None:
+        raise ValueError("extensions must be provided")
+
+    ext_seq = tuple(extensions)
+    ext_count = len(ext_seq)
+
+    cdef size_t c_ext_count = ext_count
+    cdef const char** ext_array = NULL
+    cdef char** file_array = NULL
+    cdef size_t file_count = 0
+
+    ext_bytes = [ext.encode("utf-8") if isinstance(ext, str) else ext for ext in ext_seq]
+
+    if c_ext_count > 0:
+        ext_array = <const char**>PyMem_Malloc(c_ext_count * sizeof(const char*))
+        if ext_array == NULL:
+            raise MemoryError("Unable to allocate extension array")
+        for idx in range(c_ext_count):
+            ext_array[idx] = <const char*>ext_bytes[idx]
+
+    path_b = path.encode("utf-8")
+    cdef const char* c_path = path_b
+
+    cdef int rc
+    with nogil:
+        rc = list_matching_files_c(c_path, ext_array, c_ext_count, &file_array, &file_count)
+
+    if ext_array != NULL:
+        PyMem_Free(ext_array)
+
+    if rc != 0:
+        if file_array != NULL:
+            free_file_names(file_array, file_count)
+        err = -rc
+        if err == errno.ENOENT:
+            raise FileNotFoundError(path)
+        raise OSError(err, f"Failed to list files: {path}")
+
+    results = []
+    try:
+        for idx in range(file_count):
+            py_path = PyUnicode_FromString(file_array[idx])
+            if py_path is None:
+                raise MemoryError("Unable to decode file path")
+            results.append(py_path)
+    finally:
+        if file_array != NULL:
+            free_file_names(file_array, file_count)
+
+    return results
+
+
 def read_file_mmap(str path):
     """
     Read file using memory mapping - returns an object that provides memoryview interface

opteryx/connectors/disk_connector.py

@@ -15,6 +15,7 @@
 from concurrent.futures import FIRST_COMPLETED
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import wait
+from functools import lru_cache
 from typing import Dict
 from typing import List
 
@@ -38,6 +39,11 @@
 OS_SEP = os.sep
 
 
+@lru_cache(maxsize=1)
+def _get_disk_reader_module():
+    return importlib.import_module("opteryx.compiled.io.disk_reader")
+
+
 class DiskConnector(BaseConnector, Partitionable, PredicatePushable, LimitPushable, Statistics):
     """
     Connector for reading datasets from files on local storage.
@@ -121,8 +127,8 @@ def read_blob(
             OSError:
                 If an I/O error occurs while reading the file.
         """
-        disk_reader = importlib.import_module("opteryx.compiled.io.disk_reader")
-        read_file_mmap = getattr(disk_reader, "read_file_mmap")
+        disk_reader = _get_disk_reader_module()
+        read_file_mmap = disk_reader.read_file_mmap
 
         # from opteryx.compiled.io.disk_reader import unmap_memory
         # Read using mmap for maximum speed
@@ -174,12 +180,16 @@ def get_list_of_blob_names(self, *, prefix: str) -> List[str]:
         if prefix in self.blob_list:
             return self.blob_list[prefix]
 
-        blobs = sorted(
-            os.path.join(root, file)
-            for root, _, files in os.walk(prefix + OS_SEP)
-            for file in files
-            if file.endswith(TUPLE_OF_VALID_EXTENSIONS)
-        )
+        disk_reader = _get_disk_reader_module()
+        list_files = getattr(disk_reader, "list_files", None)
+        if list_files is None:
+            raise RuntimeError("Compiled disk reader missing list_files helper")
+
+        target = os.path.normpath(prefix)
+        try:
+            blobs = sorted(list_files(target, TUPLE_OF_VALID_EXTENSIONS))
+        except FileNotFoundError:
+            blobs = []
 
         self.blob_list[prefix] = blobs
         return blobs

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "opteryx"
-version = "0.26.2-beta.1739"
+version = "0.26.2-beta.1740"
 description = "Query your data, where it lives"
 requires-python = '>=3.11'
 readme = {file = "README.md", content-type = "text/markdown"}

setup.py

@@ -250,10 +250,15 @@ def rust_build(setup_kwargs: Dict[str, Any]) -> None:
         ),
         Extension(
             name="opteryx.compiled.io.disk_reader",
-            sources=["opteryx/compiled/io/disk_reader.pyx", "src/cpp/disk_io.cpp"],
+            sources=[
+                "opteryx/compiled/io/disk_reader.pyx",
+                "src/cpp/disk_io.cpp",
+                "src/cpp/directories.cpp",
+            ],
             include_dirs=include_dirs + ["src/cpp"],
             language="c++",
             extra_compile_args=CPP_COMPILE_FLAGS,
+            depends=["src/cpp/directories.h"],
         ),
         Extension(
             name="opteryx.compiled.table_ops.distinct",