Skip to content

Commit

Permalink
Cache the in-memory FITS buffer
Browse files Browse the repository at this point in the history
In some cases a person asking for a component might be asking
for two components.
  • Loading branch information
timj committed Feb 27, 2025
1 parent fa9e900 commit f3fe595
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions python/lsst/obs/base/formatters/fitsExposure.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
"standardizeAmplifierParameters",
)

import uuid
import warnings
from abc import abstractmethod
from collections.abc import Set
Expand Down Expand Up @@ -523,6 +524,7 @@ class FitsExposureFormatter(FitsMaskedImageFormatter):

can_read_from_uri = True
ReaderClass = ExposureFitsReader
_cached_fits: tuple[uuid.UUID, MemFileManager] = (None, None)

def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any:
# For now only support small non-pixel components. In future
Expand All @@ -534,10 +536,17 @@ def read_from_uri(self, uri: ResourcePath, component: str | None = None, expecte
if not component or component in pixel_components:
# For pixel access download the whole file.
return NotImplemented

cached_id, cached_mem = type(self)._cached_fits
if self.dataset_ref.id == cached_id:
self._reader = self.ReaderClass(cached_mem)
return self.readComponent(component)

try:
fs, fspath = uri.to_fsspec()
hdul = []
with fs.open(fspath) as f, astropy.io.fits.open(f) as fits_obj:
# Read all non-pixel components and cache.
for hdu in fits_obj:
hdr = hdu.header
extname = hdr.get("EXTNAME")
Expand Down Expand Up @@ -566,6 +575,7 @@ def read_from_uri(self, uri: ResourcePath, component: str | None = None, expecte
fits_data = buffer.getvalue()
mem = MemFileManager(len(fits_data))
mem.setData(fits_data, len(fits_data))
type(self)._cached_fits = (self.dataset_ref.id, mem)
self._reader = self.ReaderClass(mem)
return self.readComponent(component)

Expand Down

0 comments on commit f3fe595

Please sign in to comment.