Реализация кэша LRU в Python с использованием очереди приоритетов

#python #python-3.x #caching #data-structures #lru

#python #python-3.x #кэширование #структуры данных #lru

Вопрос:

У меня есть следующий код для реализации кэша LRU.

 from __future__ import annotations

from time import time

import heapq

from typing import List, Dict, TypeVar, Generic, Optional, Tuple

# LRU Cache
T = TypeVar('T')


class Element:
    def __init__(self, key: str) -> None:
        self.key = key
        self.unixtime = time()

    def __lt__(self, other: Element) -> bool:
        return self.unixtime < other.unixtime

    def __eq__(self, other: Element) -> bool:
        return self.unixtime == other.unixtime

    def __gt__(self, other: Element) -> bool:
        return (not self.unixtime < other.unixtime) and self.unixtime != other.unixtime

    def __repr__(self) -> str:
        return f'({self.key}, {self.unixtime})'


class PriorityQueue(Generic[T], list):
    def __init__(self) -> None:
        self._data: List[Optional[T]] = []
        super().__init__()

    @property
    def is_empty(self) -> bool:
        return not self._data

    def push(self, v: T) -> None:
        heapq.heappush(self._data, v)

    def popq(self) -> Optional[T]:
        if not self.is_empty:
            return heapq.heappop(self._data)
        else:
            return None

    def __repr__(self) -> str:
        return repr(self._data)


class LRUCache:
    def __init__(self, limit: int) -> None:
        self._data: Dict[str, int] = {}
        self.limit = limit
        self._keyqueue: PriorityQueue[Element] = PriorityQueue()

    def put(self, key: str, value: T) -> None:
        if len(self._data) < self.limit:    # there is still room in the cache
            if key not in self._data:
                self._keyqueue.push(Element(key))
            else:
                correct_key = [item for item in self._keyqueue._data if item.key == key][0]
                ind = self._keyqueue._data.index(correct_key)
                self._keyqueue._data[ind].unixtime = time()
            self._data[key] = value
        else:                               # cache is full
            if key not in self._data:
                out_key = self._keyqueue.popq()
                self._data.pop(out_key.key)
                self._keyqueue.push(Element(key))
            else:
                correct_key = [item for item in self._keyqueue._data if item.key == key][0]
                ind = self._keyqueue._data.index(correct_key)
                self._keyqueue._data[ind].unixtime = time()
            self._data[key] = value

    def get(self, key: str) -> Optional[T]:
        if key in self._data:
            correct_key = [item for item in self._keyqueue._data if item.key == key][0]
            ind = self._keyqueue._data.index(correct_key)
            self._keyqueue._data[ind].unixtime = time()
            return self._data[key]
        else:
            raise KeyError('Key not found in cache')

    def __repr__(self) -> str:
        return repr(self._data)

cache = LRUCache(3)
cache.put('owen', 45)
cache.put('john', 32)
cache.put('box', 4556)

cache.get('owen')
cache.get('owen')

cache.put('new', 9)
cache
  

Я использую unixtime атрибут Element класса для определения приоритета. Я использую heapq модуль вместе со списком для реализации очереди приоритетов. Возможно, это не самый эффективный способ реализации кэша LRU в Python, но это то, что я придумал.

Моя проблема в том, что после того, как я дважды обращаюсь к owen ключу с .get() помощью, а затем выдаю cache.put('new', 9) — он должен быть удален john , потому что он используется наименее недавно. Вместо этого он удаляет owen .
Я проверил _keyqueue и owen имеет самый высокий unixtime и john самый низкий, и, как я понимаю, heapq модуль в Python использует min_heap , поэтому john запись должна быть заменена новым значением. Чего мне здесь не хватает?

Ответ №1:

Я, наконец, обнаружил, в чем проблема: всякий раз, когда обновляется время, нам нужно вызывать heapq.heapify() данные кучи после обновления. Я также написал немного более эффективную реализацию, если она кому-то нужна:

 from typing import List, Optional, TypeVar, Tuple, Dict, Generic

from time import time

import heapq

T = TypeVar('T')


class LRUTuple(tuple):
    def __init__(self, key: Tuple[str]) -> None:
        self.key = key
        self.time = time()

    def __lt__(self, other) -> bool:
        return self.time < other.time

    def __gt__(self, other) -> bool:
        return not self.time < other.time


# test class
a = LRUTuple(('owen',))
b = LRUTuple(('aheek',))
assert b > a
assert a < b


class PriorityQueue(Generic[T]):
    def __init__(self) -> None:
        self._data: List[T] = []

    @property
    def is_empty(self) -> bool:
        return not self._data

    def add(self, v: T) -> None:
        heapq.heappush(self._data, v)

    def pop_queue(self) -> Optional[T]:
        if not self.is_empty:
            return heapq.heappop(self._data)
        else:
            print('Empty Queue')
            return None

    def _heapify(self) -> None:
        heapq.heapify(self._data)

    def peek(self) -> Optional[T]:
        if not self.is_empty:
            return self._data[0]
        else:
            print('Empty Queue')
            return None

    def __repr__(self) -> str:
        return repr(self._data)


class LRUCache:
    def __init__(self, limit: int) -> None:
        self._data: Dict[str, T] = {}
        self.limit = limit
        self._keyqueue: PriorityQueue[LRUTuple] = PriorityQueue()

    def _update_key_time(self, key: str) -> None:
        self._keyqueue._data[self._keyqueue._data.index((key,))].time = time()
        self._keyqueue._heapify()

    def put(self, key: str, value: T) -> None:
        if len(self._keyqueue._data) < self.limit:
            if key not in self._data:
                self._data[key] = value
                self._keyqueue.add(LRUTuple((key,)))
            else:
                self._data[key] = value
                self._update_key_time(key)
        else:
            # remove lru key
            poped_key = self._keyqueue.pop_queue()
            self._data.pop(poped_key[0])
            self.put(key, value)

    def get(self, key: str) -> Optional[T]:
        if key in self._data:
            self._update_key_time(key)
            return self._data[key]
        else:
            print('KeyError: key not found')
            return None

    def __repr__(self) -> str:
        return repr([(k[0], k.time) for k in self._keyqueue._data])


# test LRUCache usage
lr = LRUCache(3)
lr.put('owen', 54)
lr.put('arwen', 4)
lr.put('jiji', 98)
lr._keyqueue.peek()
lr.get('owen')
lr._keyqueue.peek()
lr
lr.put('bone', 7)   # should replace arwen!
lr
lr._keyqueue.peek()