Use parquet stats for pruning reads #2687

Author: joocer

opteryx/__version__.py

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

opteryx/config.py

@@ -169,6 +169,9 @@ def get(key: str, default: Optional[typing.Any] = None) -> Optional[typing.Any]:
 MAX_READ_BUFFER_CAPACITY: int = memory_allocation_calculation(float(get("MAX_READ_BUFFER_CAPACITY", 0.1)))
 """Read buffer pool size in either bytes or fraction of system memory."""
 
+MAX_STATISTICS_CACHE_ITEMS: int = get("MAX_STATISTICS_CACHE_ITEMS", 10_000)
+"""The number of .parquet files we cache the statistics for."""
+
 CONCURRENT_READS: int = int(get("CONCURRENT_READS", max(system_gigabytes(), 2)))
 """Number of read workers per data source."""
 

opteryx/connectors/capabilities/statistics.py

@@ -4,17 +4,110 @@
 # 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
 
 from orso.schema import RelationSchema
 
+from opteryx.managers.expression import NodeType
 from opteryx.models import RelationStatistics
+from opteryx.shared.stats_cache import StatsCache
+from opteryx.third_party.cyan4973.xxhash import hash_bytes
 
 
 class Statistics:
     def __init__(self, statistics: dict, **kwargs):
+        self.stats_cache = StatsCache()
         self.relation_statistics = RelationStatistics()
 
+    def read_blob_statistics(
+        self, blob_name: str, blob_bytes: bytes = None, decoder=None
+    ) -> Optional[Dict[str, Any]]:
+        key = hex(hash_bytes(blob_name.encode())).encode()
+        cached_stats = self.stats_cache.get(key)
+        if cached_stats is not None:
+            # If statistics are cached, return them
+            return cached_stats
+
+        cached_stats = decoder(blob_bytes, just_statistics=True)
+        self.stats_cache.set(key, cached_stats)
+        return cached_stats
+
+    def prefilter_blobs(self, blob_names: list[str], query_statistics, selection) -> list[str]:
+        new_blob_names = []
+        for blob_name in blob_names:
+            key = hex(hash_bytes(blob_name.encode())).encode()
+            cached_stats = self.stats_cache.get(key)
+            if cached_stats is None:
+                # we have no stats so we can't make a decision
+                new_blob_names.append(blob_name)
+                query_statistics.no_stats += 1
+                continue
+
+            skip_blob = False
+
+            for condition in selection:
+                if condition.left.node_type != NodeType.IDENTIFIER:
+                    continue
+                if condition.right.node_type != NodeType.LITERAL:
+                    continue
+
+                column_name = condition.left.source_column
+                literal_value = condition.right.value
+                max_value = cached_stats.upper_bounds.get(column_name, None)
+                min_value = cached_stats.lower_bounds.get(column_name, None)
+
+                if max_value is None or min_value is None:
+                    continue
+
+                if condition.value == "Eq":  # noqa: SIM102
+                    # value must be within [min, max]
+                    if literal_value < min_value or literal_value > max_value:
+                        query_statistics.blobs_pruned += 1
+                        skip_blob = True
+                        break
+
+                elif condition.value == "NotEq":  # noqa: SIM102
+                    # only prune if min == max == literal (i.e., column only contains this value)
+                    if min_value == max_value == literal_value:
+                        query_statistics.blobs_pruned += 1
+                        skip_blob = True
+                        break
+
+                elif condition.value == "Gt":  # noqa: SIM102
+                    # value must be less than max to potentially match
+                    if max_value <= literal_value:
+                        query_statistics.blobs_pruned += 1
+                        skip_blob = True
+                        break
+
+                elif condition.value == "GtEq":  # noqa: SIM102
+                    # value must be less than or equal to max to potentially match
+                    if max_value < literal_value:
+                        query_statistics.blobs_pruned += 1
+                        skip_blob = True
+                        break
+
+                elif condition.value == "Lt":  # noqa: SIM102
+                    # value must be greater than min to potentially match
+                    if min_value >= literal_value:
+                        query_statistics.blobs_pruned += 1
+                        skip_blob = True
+                        break
+
+                elif condition.value == "LtEq":  # noqa: SIM102
+                    # value must be greater than or equal to min to potentially match
+                    if min_value > literal_value:
+                        query_statistics.blobs_pruned += 1
+                        skip_blob = True
+                        break
+
+            if not skip_blob:
+                new_blob_names.append(blob_name)
+
+        return new_blob_names
+
     def map_statistics(
         self, statistics: Optional[RelationStatistics], schema: RelationSchema
     ) -> RelationSchema:

opteryx/connectors/gcp_cloudstorage_connector.py

@@ -245,8 +245,11 @@ def read_dataset(
                         selection=predicates,
                         just_schema=just_schema,
                     )
+                    stats = self.read_blob_statistics(
+                        blob_name=blob_name, blob_bytes=blob_bytes, decoder=decoder
+                    )
                     if len(blob_names) == 1:
-                        self.relation_statistics = decoder(blob_bytes, just_statistics=True)
+                        self.relation_statistics = stats
                 except Exception as err:
                     raise DatasetReadError(f"Unable to read file {blob_name} ({err})") from err
 

opteryx/models/relation_statistics.py

@@ -1,9 +1,16 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# 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 List
 from typing import Optional
 from typing import Tuple
 
+import orjson
+
 
 class RelationStatistics:
     """
@@ -57,3 +64,31 @@ def set_cardinality_estimate(self, column: str, cardinality: int):
         if self.cardinality_estimate is None:
             self.cardinality_estimate = {}
         self.cardinality_estimate[column] = cardinality
+
+    def to_bytes(self) -> bytes:
+        """Serialize the RelationStatistics object to bytes using JSON."""
+        # Convert all attributes to a serializable dict
+        data = {
+            "record_count": self.record_count,
+            "record_count_estimate": self.record_count_estimate,
+            "null_count": self.null_count,
+            "lower_bounds": self.lower_bounds,
+            "upper_bounds": self.upper_bounds,
+            "cardinality_estimate": self.cardinality_estimate,
+            "raw_distribution_data": self.raw_distribution_data,
+        }
+        return orjson.dumps(data, default=str)
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> "RelationStatistics":
+        """Deserialize bytes to a RelationStatistics object using JSON."""
+        obj = cls()
+        loaded = orjson.loads(data)
+        obj.record_count = loaded.get("record_count", 0)
+        obj.record_count_estimate = loaded.get("record_count_estimate", 0)
+        obj.null_count = loaded.get("null_count", None)
+        obj.lower_bounds = loaded.get("lower_bounds", {})
+        obj.upper_bounds = loaded.get("upper_bounds", {})
+        obj.cardinality_estimate = loaded.get("cardinality_estimate", None)
+        obj.raw_distribution_data = loaded.get("raw_distribution_data", [])
+        return obj

opteryx/operators/async_read_node.py

@@ -106,6 +106,12 @@ def execute(self, morsel, **kwargs) -> Generator:
             predicates=self.predicates,
         )
 
+        if self.predicates and hasattr(reader, "prefilter_blobs"):
+            # if we're capturing statistics, we can prefilter the blobs
+            blob_names = reader.prefilter_blobs(
+                blob_names=blob_names, query_statistics=self.statistics, selection=self.predicates
+            )
+
         if len(blob_names) == 0:
             # if we don't have any matching blobs, create an empty dataset
             from orso import DataFrame
@@ -166,6 +172,14 @@ def execute(self, morsel, **kwargs) -> Generator:
                     decoded = decoder(
                         blob_memory_view, projection=self.columns, selection=self.predicates
                     )
+
+                    # We read the statisics from the blob, we can use this for
+                    # prefiltering the files next time we read them.
+                    if hasattr(reader, "read_blob_statistics"):
+                        reader.read_blob_statistics(
+                            blob_name=blob_name, blob_bytes=blob_memory_view, decoder=decoder
+                        )
+
                     self.pool.release(reference)  # release also unlatches the segment
                 except Exception as err:
                     from pyarrow import ArrowInvalid
@@ -197,7 +211,7 @@ def execute(self, morsel, **kwargs) -> Generator:
 
                 yield morsel
             except Exception as err:
-                self.statistics.add_message(f"failed to read {blob_name}")
+                self.statistics.add_message(f"failed to read {blob_name} ({err.__type__.__name__})")
                 self.statistics.failed_reads += 1
                 import warnings
 

opteryx/shared/stats_cache.py

@@ -0,0 +1,114 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# 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.
+
+"""
+Global Blob Statistics Cache.
+
+This module implements a global statistics cache using an LRU-K2 policy to determine
+which items to keep and which to evict. The cache is used to store and retrieve
+table or query statistics efficiently, and is shared across all connections and cursors.
+
+The cache is limited by the number of items (MAX_STATISTICS_CACHE_ITEMS), not by memory volume.
+Eviction occurs when the item count exceeds the configured maximum.
+"""
+
+from typing import Optional
+
+from opteryx.config import MAX_STATISTICS_CACHE_ITEMS
+from opteryx.models import RelationStatistics
+from opteryx.utils.lru_2 import LRU2
+
+
+class _StatsCache:
+    """
+    Implements a statistics cache using an LRU-K2 eviction policy.
+    Stores serialized statistics objects, keyed by bytes.
+    """
+
+    slots = "_lru"
+
+    def __init__(self):
+        self._lru = LRU2()
+
+    def get(self, key: bytes) -> Optional[RelationStatistics]:
+        """
+        Retrieve a statistics object from the cache by key.
+        Returns the deserialized object if found, or None if not present.
+        """
+        cached_stats = self._lru.get(key)
+        if cached_stats is not None:
+            return RelationStatistics.from_bytes(cached_stats)
+        return None
+
+    def delete(self, key: bytes):
+        """
+        Remove a statistics object from the cache by key.
+        """
+        self._lru.delete(key)
+
+    def set(self, key: bytes, value: RelationStatistics) -> Optional[str]:
+        """
+        Store a statistics object in the cache, serializing it to bytes.
+        If the cache exceeds the maximum allowed items, evict the least recently used item.
+
+        Args:
+            key: The key associated with the statistics object.
+            value: The statistics object to store (will be serialized).
+
+        Returns:
+            The key of the evicted item if eviction occurred, otherwise None.
+        """
+        cached_stats = value.to_bytes()
+
+        # Update LRU cache with the new key and memory pool key if commit succeeds
+        self._lru.set(key, cached_stats)
+        if self._lru.size > MAX_STATISTICS_CACHE_ITEMS:
+            # If the cache size exceeds the limit, evict the least recently used item
+            evicted_key = self._lru.evict()
+            if evicted_key:
+                self.delete(evicted_key)
+
+        # Return the evicted key if an eviction occurred, otherwise return None
+        return evicted_key if "evicted_key" in locals() else None
+
+    @property
+    def stats(self) -> tuple:
+        """
+        Return the hit, miss, and eviction statistics for the cache.
+        """
+        return self._lru.stats
+
+    def __del__(self):
+        pass
+        # DEBUG: print(f"Statistics Cache <hits={self.stats[0]}, misses={self.stats[1]}, evictions={self.stats[2]}, inserts={self.stats[3]}>")
+
+
+class StatsCache(_StatsCache):
+    """
+    Singleton wrapper for the _StatsCache class.
+    Ensures only one global statistics cache instance exists.
+    """
+
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = cls._create_instance()
+        return cls._instance
+
+    @classmethod
+    def _create_instance(cls):
+        """
+        Create a new instance of the underlying _StatsCache.
+        """
+        return _StatsCache()
+
+    @classmethod
+    def reset(cls):
+        """
+        Reset the StatsCache singleton instance. This is useful when the configuration changes.
+        """
+        cls._instance = None
+        cls._instance = cls._create_instance()

opteryx/utils/lru_2.py

@@ -48,6 +48,8 @@ def __init__(self, k=2):
         self.evictions = 0
         self.inserts = 0
 
+        self.size = 0
+
     def __len__(self):
         return len(self.slots)
 
@@ -62,6 +64,7 @@ def get(self, key: bytes):
 
     def set(self, key: bytes, value):
         self.inserts += 1
+        self.size += 1
         self.slots[key] = value
         self._update_access_history(key)
         return None
@@ -93,6 +96,7 @@ def evict(self, details=False):
                 continue
             value = self.slots.pop(oldest_key)
             self.access_history.pop(oldest_key)
+            self.size -= 1
             self.evictions += 1
             if details:  # pragma: no cover
                 return oldest_key, value
@@ -110,6 +114,7 @@ def delete(self, key: bytes):
             self.slots.pop(key, None)
             self.access_history.pop(key, None)
             self.evictions += 1
+            self.size -= 1
             return True
         return False