Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Use cachetools's LRUCache to cache manifest list #1187

Merged
merged 7 commits into from
Sep 24, 2024
Merged

Use cachetools's LRUCache to cache manifest list #1187

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
3a5846d
use cachetools
kevinjqliu Sep 20, 2024
39c56f6
use LRU cache
kevinjqliu Sep 20, 2024
ab6de0d
return tuple
kevinjqliu Sep 20, 2024
177dadc
comment
kevinjqliu Sep 20, 2024
70960ca
clear global cache for tests
kevinjqliu Sep 20, 2024
9d2d853
move _manifests to manifest.py
kevinjqliu Sep 20, 2024
ade0fc5
rebase poetry.lock
kevinjqliu Sep 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 15 additions & 2 deletions poetry.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

10 changes: 10 additions & 0 deletions pyiceberg/manifest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,12 @@
List,
Literal,
Optional,
Tuple,
Type,
)

from cachetools import LRUCache, cached
from cachetools.keys import hashkey
from pydantic_core import to_json

from pyiceberg.avro.file import AvroFile, AvroOutputFile
Expand Down Expand Up @@ -620,6 +623,13 @@ def fetch_manifest_entry(self, io: FileIO, discard_deleted: bool = True) -> List
]


@cached(cache=LRUCache(maxsize=128), key=lambda io, manifest_list: hashkey(manifest_list))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did the use of cachetools LRUCache, instead of functools.lru_cache solve the issues with the FileSystem on M1 Macs we were observing before?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yep! the difference is that functools.lru_cache uses all of the function's args as the cache key, including the io arg.
Using cachetools, I can specify which argument to use as the cache key. In this case, only using manifest_list as the cache key

key=lambda io, manifest_list: hashkey(manifest_list)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's so clever - so the issue was with the function arguments being 'held up' by the cache preventing them from being GCed 🤯 🧠

This is next level @kevinjqliu 💯

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's my best theory, haha. Regardless of the underlying cause, we should not be using io as the cache key!

def _manifests(io: FileIO, manifest_list: str) -> Tuple[ManifestFile, ...]:
"""Read and cache manifests from the given manifest list, returning a tuple to prevent modification."""
file = io.new_input(manifest_list)
return tuple(read_manifest_list(file))


def read_manifest_list(input_file: InputFile) -> Iterator[ManifestFile]:
"""
Read the manifests from the manifest list.
Expand Down
12 changes: 2 additions & 10 deletions pyiceberg/table/snapshots.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,12 @@
import time
from collections import defaultdict
from enum import Enum
from functools import lru_cache
from typing import TYPE_CHECKING, Any, DefaultDict, Dict, Iterable, List, Mapping, Optional

from pydantic import Field, PrivateAttr, model_serializer

from pyiceberg.io import FileIO
from pyiceberg.manifest import DataFile, DataFileContent, ManifestFile, read_manifest_list
from pyiceberg.manifest import DataFile, DataFileContent, ManifestFile, _manifests
from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec
from pyiceberg.schema import Schema

Expand Down Expand Up @@ -231,13 +230,6 @@ def __eq__(self, other: Any) -> bool:
)


@lru_cache
def _manifests(io: FileIO, manifest_list: str) -> List[ManifestFile]:
"""Return the manifests from the manifest list."""
file = io.new_input(manifest_list)
return list(read_manifest_list(file))


class Snapshot(IcebergBaseModel):
snapshot_id: int = Field(alias="snapshot-id")
parent_snapshot_id: Optional[int] = Field(alias="parent-snapshot-id", default=None)
Expand All @@ -260,7 +252,7 @@ def __str__(self) -> str:
def manifests(self, io: FileIO) -> List[ManifestFile]:
"""Return the manifests for the given snapshot."""
if self.manifest_list:
return _manifests(io, self.manifest_list)
return list(_manifests(io, self.manifest_list))
return []


Expand Down
Loading