Merge pull request #2872 from mabel-dev/clickbench-performance-regression-investigation-1

custom disk reader

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__ = 1703
+__build__ = 1704
 __author__ = "@joocer"
-__version__ = "0.26.0-beta.1703"
+__version__ = "0.26.0-beta.1704"
 
 # Store the version here so:
 # 1) we don't load dependencies by storing it in __init__.py

opteryx/compiled/io/__init__.py

@@ -0,0 +1,12 @@
+"""
+Compiled I/O operations for high-performance file access.
+"""
+
+try:
+    from .disk_reader import read_file
+    from .disk_reader import read_file_to_bytes
+
+    __all__ = ["read_file", "read_file_to_bytes"]
+except ImportError:
+    # Module not yet compiled
+    __all__ = []

opteryx/compiled/io/disk_reader.pyx

@@ -0,0 +1,115 @@
+# cython: language_level=3
+# cython: nonecheck=False
+# cython: cdivision=True
+# cython: initializedcheck=False
+# cython: infer_types=True
+# cython: wraparound=False
+# cython: boundscheck=False
+
+"""
+Ultra-fast disk reader module
+"""
+
+from cpython.buffer cimport PyBuffer_FillInfo
+from libc.stdlib cimport free
+
+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 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)
+
+    def __len__(self):
+        return self.size
+
+
+def read_file(str path, bint sequential=True, bint willneed=True, bint drop_after=False):
+    """
+    Read an entire file into memory with optimized I/O.
+    """
+    import os
+
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"File not found: {path}")
+
+    cdef size_t size = os.path.getsize(path)
+    cdef size_t out_len = 0
+
+    path_b = path.encode("utf-8")
+    cdef const char* c_path = path_b
+
+    # Allocate buffer - use bytearray for mutable buffer
+    buf = bytearray(size)
+    cdef unsigned char[::1] buf_view = buf
+    cdef unsigned char* dst = &buf_view[0]
+
+    cdef int rc = read_all_pread(c_path, dst, &out_len, sequential, willneed, drop_after)
+
+    if rc != 0:
+        raise OSError(-rc, f"Failed to read file: {path}")
+
+    return memoryview(buf)[:out_len]
+
+
+def read_file_mmap(str path):
+    """
+    Read file using memory mapping - returns an object that provides memoryview interface
+    but MUST be manually closed to avoid resource leaks.
+    """
+    import os
+
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"File not found: {path}")
+
+    path_b = path.encode("utf-8")
+    cdef const char* c_path = path_b
+    cdef unsigned char* mapped_data = NULL
+    cdef size_t size = 0
+
+    cdef int rc = read_all_mmap(c_path, &mapped_data, &size)
+
+    if rc != 0:
+        raise OSError(-rc, f"Failed to mmap file: {path}")
+
+    # Create wrapper that knows how to clean up
+    cdef MappedMemory wrapper = MappedMemory.__new__(MappedMemory)
+    wrapper.data = mapped_data
+    wrapper.size = size
+    wrapper.owned = False  # This is mmap'd memory, not malloc'd
+
+    return wrapper
+
+
+def read_file_to_bytes(str path, bint sequential=True, bint willneed=True, bint drop_after=False):
+    """
+    Read an entire file into memory as bytes.
+    """
+    mv = read_file(path, sequential, willneed, drop_after)
+    return bytes(mv)
+
+
+def unmap_memory(mem_obj):
+    """
+    Explicitly unmap memory from read_file_mmap.
+    MUST be called when done with the data to avoid resource leaks.
+    """
+    cdef int rc
+    if hasattr(mem_obj, 'data') and mem_obj.data is not None:
+        # Import the unmap function from your C code
+        rc = unmap_memory_c(mem_obj.data, mem_obj.size)
+        mem_obj.data = None
+        return rc == 0
+    return True

opteryx/connectors/disk_connector.py

@@ -127,73 +127,37 @@ def read_blob(
             OSError:
                 If an I/O error occurs while reading the file.
         """
-        # Hybrid strategy: choose mmap or read+memoryview depending on OS
-        # macOS -> mmap, Linux -> read.
-
-        # helper to use mmap path
-        def _use_mmap():
-            fd = os.open(blob_name, os.O_RDONLY)
-            try:
-                if hasattr(os, "posix_fadvise"):
-                    with contextlib.suppress(Exception):
-                        os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_WILLNEED)
-                size = os.fstat(fd).st_size
-                _map = mmap.mmap(fd, length=size, **mmap_config)
-                result = decoder(
-                    _map,
-                    just_schema=just_schema,
-                    projection=projection,
-                    selection=selection,
-                    use_threads=True,
-                )
+        from opteryx.compiled.io.disk_reader import read_file_mmap
+        from opteryx.compiled.io.disk_reader import unmap_memory
+
+        # Read using mmap for maximum speed
+        mmap_obj = read_file_mmap(blob_name)
+
+        try:
+            # Create memoryview for the decoder
+            mv = memoryview(mmap_obj)
+
+            result = decoder(
+                mv,
+                just_schema=just_schema,
+                projection=projection,
+                selection=selection,
+                use_threads=True,
+            )
 
-                self.statistics.bytes_read += size
+            self.statistics.bytes_read += len(mv)
 
-                if not just_schema:
-                    stats = self.read_blob_statistics(
-                        blob_name=blob_name, blob_bytes=_map, decoder=decoder
-                    )
-                    if self.relation_statistics is None:
-                        self.relation_statistics = stats
-
-                return result
-            finally:
-                os.close(fd)
-
-        # helper to use read()+memoryview path
-        def _use_read():
-            with open(blob_name, "rb") as f:
-                if hasattr(os, "posix_fadvise"):
-                    with contextlib.suppress(Exception):
-                        os.posix_fadvise(f.fileno(), 0, 0, os.POSIX_FADV_WILLNEED)
-
-                data = f.read()
-                size = len(data)
-                buf = memoryview(data)
-
-                result = decoder(
-                    buf,
-                    just_schema=just_schema,
-                    projection=projection,
-                    selection=selection,
-                    use_threads=True,
+            if not just_schema:
+                stats = self.read_blob_statistics(
+                    blob_name=blob_name, blob_bytes=mv, decoder=decoder
                 )
+                if self.relation_statistics is None:
+                    self.relation_statistics = stats
 
-                self.statistics.bytes_read += size
-
-                if not just_schema:
-                    stats = self.read_blob_statistics(
-                        blob_name=blob_name, blob_bytes=buf, decoder=decoder
-                    )
-                    if self.relation_statistics is None:
-                        self.relation_statistics = stats
-
-                return result
-
-        # macOS: use mmap; Linux: prefer read (observed faster on some Linux setups)
-        if platform.system() == "Darwin":
-            return _use_mmap()
-        return _use_read()
+            return result
+        finally:
+            # CRITICAL: Clean up the memory mapping
+            unmap_memory(mmap_obj)
 
     @single_item_cache
     def get_list_of_blob_names(self, *, prefix: str) -> List[str]:

pyproject.toml

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

setup.py

@@ -248,6 +248,13 @@ def rust_build(setup_kwargs: Dict[str, Any]) -> None:
             sources=["opteryx/compiled/structures/node.pyx"],
             extra_compile_args=C_COMPILE_FLAGS,
         ),
+        Extension(
+            name="opteryx.compiled.io.disk_reader",
+            sources=["opteryx/compiled/io/disk_reader.pyx", "src/cpp/disk_io.cpp"],
+            include_dirs=include_dirs + ["src/cpp"],
+            language="c++",
+            extra_compile_args=CPP_COMPILE_FLAGS,
+        ),
         Extension(
             name="opteryx.compiled.table_ops.distinct",
             sources=["opteryx/compiled/table_ops/distinct.pyx"],