Enable stats on disk connector #2697

Author: joocer

opteryx/__version__.py

@@ -1,4 +1,4 @@
-__build__ = 1385
+__build__ = 1386
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

opteryx/connectors/capabilities/statistics.py

@@ -3,11 +3,11 @@
 # See the License at http://www.apache.org/licenses/LICENSE-2.0
 # Distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND.
 
-
 from typing import Any
 from typing import Dict
 from typing import Optional
 
+import numpy
 from orso.schema import RelationSchema
 
 from opteryx.managers.expression import NodeType
@@ -66,6 +66,8 @@ def prune_blobs(self, blob_names: list[str], query_statistics, selection) -> lis
                 for condition in valid_conditions:
                     column_name = condition.left.source_column
                     literal_value = condition.right.value
+                    if type(literal_value) is numpy.datetime64:
+                        literal_value = str(literal_value.astype("M8[ms]"))
                     max_value = cached_stats.upper_bounds.get(column_name)
                     min_value = cached_stats.lower_bounds.get(column_name)
 

opteryx/connectors/disk_connector.py

@@ -15,12 +15,14 @@
 
 import pyarrow
 from orso.schema import RelationSchema
+from orso.tools import single_item_cache
 from orso.types import OrsoTypes
 
 from opteryx.connectors.base.base_connector import BaseConnector
 from opteryx.connectors.capabilities import LimitPushable
 from opteryx.connectors.capabilities import Partitionable
 from opteryx.connectors.capabilities import PredicatePushable
+from opteryx.connectors.capabilities import Statistics
 from opteryx.exceptions import DataError
 from opteryx.exceptions import DatasetNotFoundError
 from opteryx.exceptions import EmptyDatasetError
@@ -46,59 +48,7 @@
     mmap_config["access"] = mmap.ACCESS_READ
 
 
-def read_blob(
-    *, blob_name: str, decoder, statistics, just_schema=False, projection=None, selection=None
-):
-    """
-    Read a blob (binary large object) from disk using memory-mapped file access.
-
-    This method uses low-level file reading with memory-mapped files to
-    improve performance. It reads the entire file into memory and then
-    decodes it using the provided decoder function.
-
-    Parameters:
-        blob_name (str):
-            The name of the blob file to read.
-        decoder (callable):
-            A function to decode the memory-mapped file content.
-        just_schema (bool, optional):
-            If True, only the schema of the data is returned. Defaults to False.
-        projection (list, optional):
-            A list of fields to project. Defaults to None.
-        selection (dict, optional):
-            A dictionary of selection criteria. Defaults to None.
-        **kwargs:
-            Additional keyword arguments.
-
-    Returns:
-        The decoded blob content.
-
-    Raises:
-        FileNotFoundError:
-            If the blob file does not exist.
-        OSError:
-            If an I/O error occurs while reading the file.
-    """
-    try:
-        file_descriptor = os.open(blob_name, os.O_RDONLY | os.O_BINARY)
-        if hasattr(os, "posix_fadvise"):
-            os.posix_fadvise(file_descriptor, 0, 0, os.POSIX_FADV_WILLNEED)
-        size = os.fstat(file_descriptor).st_size
-        _map = mmap.mmap(file_descriptor, length=size, **mmap_config)
-        result = decoder(
-            _map,
-            just_schema=just_schema,
-            projection=projection,
-            selection=selection,
-            use_threads=True,
-        )
-        statistics.bytes_read += size
-        return result
-    finally:
-        os.close(file_descriptor)
-
-
-class DiskConnector(BaseConnector, Partitionable, PredicatePushable, LimitPushable):
+class DiskConnector(BaseConnector, Partitionable, PredicatePushable, LimitPushable, Statistics):
     """
     Connector for reading datasets from files on local storage.
     """
@@ -137,13 +87,74 @@ def __init__(self, **kwargs):
         Partitionable.__init__(self, **kwargs)
         PredicatePushable.__init__(self, **kwargs)
         LimitPushable.__init__(self, **kwargs)
+        Statistics.__init__(self, **kwargs)
 
         self.dataset = self.dataset.replace(".", OS_SEP)
         self.cached_first_blob = None  # Cache for the first blob in the dataset
         self.blob_list = {}
         self.rows_seen = 0
         self.blobs_seen = 0
 
+    def read_blob(
+        self, *, blob_name: str, decoder, just_schema=False, projection=None, selection=None
+    ):
+        """
+        Read a blob (binary large object) from disk using memory-mapped file access.
+
+        This method uses low-level file reading with memory-mapped files to
+        improve performance. It reads the entire file into memory and then
+        decodes it using the provided decoder function.
+
+        Parameters:
+            blob_name (str):
+                The name of the blob file to read.
+            decoder (callable):
+                A function to decode the memory-mapped file content.
+            just_schema (bool, optional):
+                If True, only the schema of the data is returned. Defaults to False.
+            projection (list, optional):
+                A list of fields to project. Defaults to None.
+            selection (dict, optional):
+                A dictionary of selection criteria. Defaults to None.
+            **kwargs:
+                Additional keyword arguments.
+
+        Returns:
+            The decoded blob content.
+
+        Raises:
+            FileNotFoundError:
+                If the blob file does not exist.
+            OSError:
+                If an I/O error occurs while reading the file.
+        """
+        try:
+            file_descriptor = os.open(blob_name, os.O_RDONLY | os.O_BINARY)
+            if hasattr(os, "posix_fadvise"):
+                os.posix_fadvise(file_descriptor, 0, 0, os.POSIX_FADV_WILLNEED)
+            size = os.fstat(file_descriptor).st_size
+            _map = mmap.mmap(file_descriptor, length=size, **mmap_config)
+            result = decoder(
+                _map,
+                just_schema=just_schema,
+                projection=projection,
+                selection=selection,
+                use_threads=True,
+            )
+            self.statistics.bytes_read += size
+
+            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(file_descriptor)
+
+    @single_item_cache
     def get_list_of_blob_names(self, *, prefix: str) -> List[str]:
         """
         List all blob files in the given directory path.
@@ -190,15 +201,19 @@ def read_dataset(
             prefix=self.dataset,
         )
 
+        if predicates is not None:
+            blob_names = self.prune_blobs(
+                blob_names=blob_names, query_statistics=self.statistics, selection=predicates
+            )
+
         remaining_rows = limit if limit is not None else float("inf")
 
         for blob_name in blob_names:
             decoder = get_decoder(blob_name)
             try:
                 if not just_schema:
-                    num_rows, _, decoded = read_blob(
+                    num_rows, _, decoded = self.read_blob(
                         blob_name=blob_name,
-                        statistics=self.statistics,
                         decoder=decoder,
                         just_schema=False,
                         projection=columns,
@@ -219,9 +234,8 @@ def read_dataset(
                     if remaining_rows <= 0:
                         break
                 else:
-                    schema = read_blob(
+                    schema = self.read_blob(
                         blob_name=blob_name,
-                        statistics=self.statistics,
                         decoder=decoder,
                         just_schema=True,
                     )

opteryx/connectors/file_connector.py

@@ -8,6 +8,7 @@
 dataset name in a query.
 """
 
+import mmap
 import os
 from typing import Dict
 from typing import Optional
@@ -20,10 +21,77 @@
 from opteryx.connectors.capabilities import LimitPushable
 from opteryx.connectors.capabilities import PredicatePushable
 from opteryx.connectors.capabilities import Statistics
-from opteryx.connectors.disk_connector import read_blob
 from opteryx.exceptions import DatasetNotFoundError
+from opteryx.utils import is_windows
 from opteryx.utils.file_decoders import get_decoder
 
+IS_WINDOWS = is_windows()
+# Define os.O_BINARY for non-Windows platforms if it's not already defined
+if not hasattr(os, "O_BINARY"):
+    os.O_BINARY = 0  # Value has no effect on non-Windows platforms
+if not hasattr(os, "O_DIRECT"):
+    os.O_DIRECT = 0  # Value has no effect on non-Windows platforms
+
+mmap_config = {}
+if not IS_WINDOWS:
+    mmap_config["flags"] = mmap.MAP_PRIVATE
+    mmap_config["prot"] = mmap.PROT_READ
+else:
+    mmap_config["access"] = mmap.ACCESS_READ
+
+
+def read_blob(
+    *, blob_name: str, decoder, just_schema=False, statistics=None, projection=None, selection=None
+):
+    """
+    Read a blob (binary large object) from disk using memory-mapped file access.
+
+    This method uses low-level file reading with memory-mapped files to
+    improve performance. It reads the entire file into memory and then
+    decodes it using the provided decoder function.
+
+    Parameters:
+        blob_name (str):
+            The name of the blob file to read.
+        decoder (callable):
+            A function to decode the memory-mapped file content.
+        just_schema (bool, optional):
+            If True, only the schema of the data is returned. Defaults to False.
+        projection (list, optional):
+            A list of fields to project. Defaults to None.
+        selection (dict, optional):
+            A dictionary of selection criteria. Defaults to None.
+        **kwargs:
+            Additional keyword arguments.
+
+    Returns:
+        The decoded blob content.
+
+    Raises:
+        FileNotFoundError:
+            If the blob file does not exist.
+        OSError:
+            If an I/O error occurs while reading the file.
+    """
+    try:
+        file_descriptor = os.open(blob_name, os.O_RDONLY | os.O_BINARY)
+        if hasattr(os, "posix_fadvise"):
+            os.posix_fadvise(file_descriptor, 0, 0, os.POSIX_FADV_WILLNEED)
+        size = os.fstat(file_descriptor).st_size
+        _map = mmap.mmap(file_descriptor, length=size, **mmap_config)
+        result = decoder(
+            _map,
+            just_schema=just_schema,
+            projection=projection,
+            selection=selection,
+            use_threads=True,
+        )
+        statistics.bytes_read += size
+
+        return result
+    finally:
+        os.close(file_descriptor)
+
 
 class FileConnector(BaseConnector, PredicatePushable, Statistics, LimitPushable):
     """