Merge pull request #2693 from mabel-dev/#2692

Enable stats and pushdowns on S3 #2692

Author: joocer

opteryx/__version__.py

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

opteryx/utils/lru_2.py opteryx/compiled/structures/lru_k.pyxopteryx/utils/lru_2.py renamed to opteryx/compiled/structures/lru_k.pyx

@@ -1,7 +1,11 @@
-# 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.
+# distutils: language = c++
+# cython: language_level=3
+# cython: nonecheck=False
+# cython: cdivision=True
+# cython: initializedcheck=False
+# cython: infer_types=True
+# cython: wraparound=False
+# cython: boundscheck=False
 
 """
 LRU-K evicts the morsel whose K-th most recent access is furthest in the past. Note, the
@@ -30,13 +34,30 @@
 the BufferPool implements limit to only evict up to 32 items per 'transaction'
 """
 
-import heapq
-import time
+import heapq as py_heapq
+
+from libc.stdint cimport int64_t
 from collections import defaultdict
+from time import monotonic_ns
+
+cdef class LRU_K:
+
+    __slots__ = ("k", "slots", "access_history", "removed", "heap",
+                 "hits", "misses", "evictions", "inserts", "size")
 
+    cdef public int64_t k
+    cdef dict slots
+    cdef object access_history
+    cdef set removed
+    cdef list heap
 
-class LRU2:
-    def __init__(self, k=2):
+    cdef int64_t hits
+    cdef int64_t misses
+    cdef int64_t evictions
+    cdef int64_t inserts
+    cdef public int64_t size
+
+    def __cinit__(self, int64_t k=2):
         self.k = k
         self.slots = {}
         self.access_history = defaultdict(list)
@@ -47,70 +68,72 @@ def __init__(self, k=2):
         self.misses = 0
         self.evictions = 0
         self.inserts = 0
-
         self.size = 0
 
     def __len__(self):
         return len(self.slots)
 
-    def get(self, key: bytes):
-        value = self.slots.get(key)
+    def get(self, bytes key) -> Optional[bytes]:
+        cdef object value = self.slots.get(key)
         if value is not None:
             self.hits += 1
             self._update_access_history(key)
         else:
             self.misses += 1
         return value
 
-    def set(self, key: bytes, value):
+    def set(self, bytes key, bytes value):
         self.inserts += 1
         if key not in self.slots:
             self.size += 1
         self.slots[key] = value
         self._update_access_history(key)
         return None
 
-    def _update_access_history(self, key: bytes):
-        access_time = time.monotonic_ns()
-        if len(self.access_history[key]) == self.k:
-            old_entry = self.access_history[key].pop(0)
+    cdef void _update_access_history(self, bytes key):
+        cdef int64_t access_time = monotonic_ns()
+        cdef list history = self.access_history[key]
+        if len(history) == self.k:
+            old_entry = history.pop(0)
             self.removed.add(old_entry)
-        self.access_history[key].append((access_time, key))
-        heapq.heappush(self.heap, (access_time, key))
-
-    def evict(self, details=False):
+        history.append((access_time, key))
+        py_heapq.heappush(self.heap, (access_time, key))
+
+    def evict(self, bint details=False):
+        cdef int64_t _oldest_access_time
+        cdef bytes oldest_key
+        cdef int64_t new_access_time
+        cdef tuple popped
         while self.heap:
-            oldest_access_time, oldest_key = heapq.heappop(self.heap)
-            if (oldest_access_time, oldest_key) in self.removed:
-                self.removed.remove((oldest_access_time, oldest_key))
+            popped = py_heapq.heappop(self.heap)
+            _oldest_access_time, oldest_key = popped
+            if popped in self.removed:
+                self.removed.remove(popped)
                 continue
 
             if len(self.access_history[oldest_key]) == 1:
                 # Synthetic access to give a grace period
-                new_access_time = time.monotonic_ns()
+                new_access_time = monotonic_ns()
                 self.access_history[oldest_key].append((new_access_time, oldest_key))
-                heapq.heappush(self.heap, (new_access_time, oldest_key))
+                py_heapq.heappush(self.heap, (new_access_time, oldest_key))
                 continue
 
-            # Evict the key with the oldest k-th access
             if oldest_key not in self.slots:
                 continue
+
             value = self.slots.pop(oldest_key)
             self.access_history.pop(oldest_key)
             self.size -= 1
             self.evictions += 1
-            if details:  # pragma: no cover
+            if details:
                 return oldest_key, value
             return oldest_key
 
-        if details:  # pragma: no cover
-            return None, None  # No item was evicted
+        if details:
+            return None, None
         return None
 
-    def delete(self, key: bytes):
-        """
-        Delete an item from the cache.
-        """
+    def delete(self, bytes key):
         if key in self.slots:
             self.slots.pop(key, None)
             self.access_history.pop(key, None)
@@ -127,11 +150,11 @@ def keys(self):
     def stats(self):
         return self.hits, self.misses, self.evictions, self.inserts
 
-    def reset(self, reset_stats=False):
-        self.slots = {}
+    def reset(self, bint reset_stats=False):
+        self.slots.clear()
         self.access_history.clear()
         self.removed.clear()
-        self.heap = []
+        self.heap.clear()
         if reset_stats:
             self.hits = 0
             self.misses = 0

opteryx/connectors/aws_s3_connector.py

@@ -9,16 +9,20 @@
 
 import asyncio
 import os
+from typing import Dict
 from typing import List
 
 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 Asynchronous
 from opteryx.connectors.capabilities import Cacheable
 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 MissingDependencyError
@@ -31,10 +35,31 @@
 OS_SEP = os.sep
 
 
-class AwsS3Connector(BaseConnector, Cacheable, Partitionable, Asynchronous):
+class AwsS3Connector(
+    BaseConnector, Cacheable, Partitionable, PredicatePushable, Asynchronous, Statistics
+):
     __mode__ = "Blob"
     __type__ = "S3"
 
+    PUSHABLE_OPS: Dict[str, bool] = {
+        "Eq": True,
+        "NotEq": True,
+        "Gt": True,
+        "GtEq": True,
+        "Lt": True,
+        "LtEq": True,
+    }
+
+    PUSHABLE_TYPES = {
+        OrsoTypes.BLOB,
+        OrsoTypes.BOOLEAN,
+        OrsoTypes.DOUBLE,
+        OrsoTypes.INTEGER,
+        OrsoTypes.VARCHAR,
+        OrsoTypes.TIMESTAMP,
+        OrsoTypes.DATE,
+    }
+
     def __init__(self, credentials=None, **kwargs):
         try:
             from minio import Minio  # type:ignore
@@ -44,7 +69,9 @@ def __init__(self, credentials=None, **kwargs):
         BaseConnector.__init__(self, **kwargs)
         Partitionable.__init__(self, **kwargs)
         Cacheable.__init__(self, **kwargs)
+        PredicatePushable.__init__(self, **kwargs)
         Asynchronous.__init__(self, **kwargs)
+        Statistics.__init__(self, **kwargs)
 
         # fmt:off
         end_point = kwargs.get("S3_END_POINT", os.environ.get("MINIO_END_POINT"))
@@ -89,6 +116,13 @@ def read_dataset(
                 blob_bytes = self.read_blob(blob_name=blob_name, statistics=self.statistics)
                 try:
                     decoded = decoder(blob_bytes, projection=columns, 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 = stats
+
                 except Exception as err:
                     raise DataError(f"Unable to read file {blob_name} ({err})") from err
                 if not just_schema:

opteryx/connectors/capabilities/cacheable.py

@@ -11,7 +11,7 @@
 from opteryx.config import MAX_CACHEABLE_ITEM_SIZE
 from opteryx.third_party.cyan4973.xxhash import hash_bytes
 
-__all__ = ("Cacheable", "async_read_thru_cache")
+__all__ = ("Cacheable", "async_read_thru_cache", "read_thru_cache")
 
 SOURCE_NOT_FOUND = 0
 SOURCE_BUFFER_POOL = 1
@@ -58,6 +58,75 @@ def purge_blob(self, blob_name: str):
             return None
 
 
+def read_thru_cache(func):
+    """
+    Decorator to implement a read-thru cache.
+
+    It intercepts requests to read blobs and first looks them up in the in-memory
+    cache (BufferPool) and optionally in a secondary cache (like MemcacheD or Redis).
+    """
+
+    # Capture the max_evictions value at decoration time
+    from opteryx import get_cache_manager
+    from opteryx.managers.cache import NullCache
+    from opteryx.shared import BufferPool
+
+    cache_manager = get_cache_manager()
+    max_evictions = MAX_CACHE_EVICTIONS_PER_QUERY
+    remote_cache = cache_manager.cache_backend
+    if not remote_cache:
+        # rather than make decisions - just use a dummy
+        remote_cache = NullCache()
+
+    buffer_pool = BufferPool()
+
+    my_keys = set()
+
+    @wraps(func)
+    def wrapper(blob_name, statistics, **kwargs):
+        nonlocal max_evictions
+
+        key = hex(hash_bytes(blob_name.encode())).encode()
+        my_keys.add(key)
+
+        # try the buffer pool first
+        result = buffer_pool.get(key)
+        if result is not None:
+            statistics.bufferpool_hits += 1
+            return result
+
+        # try the remote cache next
+        result = remote_cache.get(key)
+        if result is not None:
+            statistics.remote_cache_hits += 1
+            return result
+
+        # Key is not in cache, execute the function and store the result in cache
+        result = func(blob_name=blob_name, **kwargs)
+
+        # Write the result to caches
+        if max_evictions:
+            # we set a per-query eviction limit
+            if len(result) < MAX_CACHEABLE_ITEM_SIZE:
+                evicted = buffer_pool.set(key, result)
+                remote_cache.set(key, result)
+                if evicted:
+                    # if we're evicting items we're putting into the cache
+                    if evicted in my_keys:
+                        max_evictions = 0
+                    else:
+                        max_evictions -= 1
+                    statistics.cache_evictions += 1
+            else:
+                statistics.cache_oversize += 1
+
+        statistics.cache_misses += 1
+
+        return result
+
+    return wrapper
+
+
 def async_read_thru_cache(func):
     """
     This is added to the reader by the binder.