Coverage for python/lsst/images/fits/_input_archive.py: 72%
238 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-22 01:54 -0700
1# This file is part of lsst-images.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
12from __future__ import annotations
14__all__ = (
15 "DEFAULT_PAGE_SIZE",
16 "READ_CACHE_MAX_BYTES",
17 "FitsInputArchive",
18 "FitsOpaqueMetadata",
19)
21import io
22import os
23from collections.abc import Callable, Iterator
24from contextlib import contextmanager
25from functools import cached_property
26from types import EllipsisType
27from typing import IO, Any, Self
29import astropy.io.fits
30import astropy.table
31import fsspec
32import numpy as np
34from lsst.resources import ResourcePath, ResourcePathExpression
36from .._transforms import FrameSet
37from ..serialization import (
38 ArchiveInfo,
39 ArchiveReadError,
40 ArchiveTree,
41 ArrayReferenceModel,
42 InlineArrayModel,
43 InputArchive,
44 TableModel,
45 no_header_updates,
46 parameterize_tree,
47 tree_class_for_info,
48)
49from ..serialization._common import _check_format_version
50from ._common import (
51 JSON_COLUMN,
52 JSON_EXTNAME,
53 ExtensionHDU,
54 ExtensionKey,
55 FitsOpaqueMetadata,
56 InvalidFitsArchiveError,
57 PointerModel,
58)
60_FITS_FORMAT_VERSION = 1
61"""Container layout version this release of `FitsInputArchive` understands."""
63DEFAULT_PAGE_SIZE = 2880 * 800
64"""Default fsspec read-block size for partial (remote) reads, in bytes.
66This is the single place to tune the block size for remote-store performance.
67On a buffered remote filesystem (e.g. GCS) each cache miss is one range
68request, so the block size trades round trips against over-fetch: a component
69read that touches scattered compressed tiles pulls one block per cluster of
70nearby tiles, rounded up to this size.
72The optimum depends on the access pattern. Larger blocks favor reads that
73touch most of the file (full planes, large cutouts); smaller blocks reduce
74wasted bytes for small scattered cutouts. ``2880 * 800`` (~2.3 MB, and a
75multiple of the 2880-byte FITS block) is a robust middle: across cutout sizes
76and full reads it stays within ~1.5x of the per-pattern optimum, whereas the
77historical 144 KB default was several times slower for all but the tiniest
78cutout. Raise it (e.g. ``2880 * 1600``) when whole-file or large-cutout reads
79dominate; lower it when many tiny cutouts across many files dominate.
81Local filesystems ignore this (their opener does no buffering), so it only
82affects remote stores.
83"""
85_READ_CACHE_TYPE = "blockcache"
86"""fsspec cache strategy for partial reads.
88``blockcache`` keeps a bounded set of fixed-size blocks (so memory stays
89capped) and reuses them across the multiple components of one file -- image,
90mask, variance and so on often share blocks -- unlike the default unbounded
91single-block ``readahead``.
92"""
94READ_CACHE_MAX_BYTES = 64 * 1024 * 1024
95"""Approximate memory budget for the partial-read block cache, per open file.
97The fsspec block cache evicts least-recently-used blocks once it holds more
98than ``maxblocks``; we derive ``maxblocks`` from this budget and the block
99size (`DEFAULT_PAGE_SIZE`) so the memory cap is expressed in bytes and stays
100fixed even when the block size is retuned. Measured benefit saturates at two
101cached blocks for the access patterns we care about, so this budget is purely
102headroom plus a guard against unbounded growth; it is far below fsspec's
103implicit default of ``32 * block_size``.
104"""
107class FitsInputArchive(InputArchive[PointerModel]):
108 """An implementation of the `.serialization.InputArchive` interface that
109 reads from FITS files.
111 Instances of this class should only be constructed via the `open`
112 context manager.
113 """
115 @classmethod
116 def get_basic_info(cls, path: ResourcePathExpression) -> ArchiveInfo:
117 """Read ``DATAMODL`` (schema URL) and ``FMTVER`` (container version)
118 from the primary header.
120 Every FITS file written by this package records the schema URL in
121 the ``DATAMODL`` card, so the schema can be identified without
122 reading the (potentially large) JSON tree HDU.
123 """
124 with ResourcePath(path).open("rb") as stream:
125 primary = astropy.io.fits.PrimaryHDU.readfrom(stream)
126 header = primary.header
127 format_version = int(header.get("FMTVER", 1))
128 schema_url = header.get("DATAMODL")
129 if not schema_url: 129 ↛ 130line 129 didn't jump to line 130 because the condition on line 129 was never true
130 raise ArchiveReadError(f"{path!r} is not an lsst.images FITS archive (no DATAMODL card).")
131 return ArchiveInfo.from_schema_url(schema_url, format_version=format_version)
133 @classmethod
134 @contextmanager
135 def open_tree(
136 cls,
137 path: ResourcePathExpression,
138 *,
139 partial: bool = True,
140 **backend_kwargs: Any,
141 ) -> Iterator[tuple[Self, ArchiveTree, ArchiveInfo]]:
142 """Open the FITS file and yield ``(archive, tree, info)``.
144 Parameters
145 ----------
146 path
147 The file resource to open.
148 partial
149 If `True` the file is opened without reading it all into memory.
150 **backend_kwargs
151 Optional parameters for this backend. Currently supports
152 ``page_size`` which can be used to override the default
153 page size (which can be overridden globally by modifying
154 `DEFAULT_PAGE_SIZE`).
155 """
156 page_size = backend_kwargs.pop("page_size", DEFAULT_PAGE_SIZE)
157 with cls.open(path, page_size=page_size, partial=partial) as archive:
158 info = archive.info
159 tree_cls = tree_class_for_info(info, path)
160 parameterized = parameterize_tree(tree_cls, PointerModel)
161 tree = archive.get_tree(parameterized)
162 yield archive, tree, info
164 def __init__(self, stream: IO[bytes]):
165 self._primary_hdu = astropy.io.fits.PrimaryHDU.readfrom(stream)
166 on_disk_fmtver: int = self._primary_hdu.header.pop("FMTVER", 1)
167 # DATAMODL is informational only on read; the JSON tree's
168 # schema_version / min_read_version drive data-model checks. We
169 # capture it here as ArchiveInfo so callers (e.g. open_tree) can
170 # identify the schema from this open rather than reopening the file.
171 # A schema-less file can still be opened directly; only callers that
172 # need the schema (via the `info` property) require DATAMODL.
173 schema_url = self._primary_hdu.header.pop("DATAMODL", None)
174 self._info = (
175 ArchiveInfo.from_schema_url(schema_url, format_version=on_disk_fmtver) if schema_url else None
176 )
177 _check_format_version("fits", on_disk_fmtver, _FITS_FORMAT_VERSION)
178 # TODO: do some basic checks that the file format conforms to our
179 # expectations (e.g. primary HDU should have no data).
180 #
181 # Read and strip the addresses and sizes from the headers. We don't
182 # actually need the index address because we always want to read the
183 # JSON HDU, too, and the index HDU is always the next one (but this
184 # could change in the future).
185 json_address: int = self._primary_hdu.header.pop("JSONADDR")
186 json_size: int = self._primary_hdu.header.pop("JSONSIZE")
187 del self._primary_hdu.header["INDXADDR"]
188 index_size: int = self._primary_hdu.header.pop("INDXSIZE")
189 # Save the remaining primary header keys so we can propagate them on
190 # rewrite.
191 self._opaque_metadata = FitsOpaqueMetadata()
192 self._opaque_metadata.add_header(self._primary_hdu.header.copy(strip=True), name="", ver=1)
193 # Read the JSON and index HDUs from the end.
194 stream.seek(json_address)
195 tail_data = stream.read(json_size + index_size)
196 index_hdu = astropy.io.fits.BinTableHDU.fromstring(tail_data[json_size:])
197 # Initialize lazy readers for all of the regular HDUs and the JSON HDU.
198 self._readers = {
199 ExtensionKey.from_index_row(row): _ExtensionReader.from_index_row(row, stream)
200 for row in index_hdu.data
201 }
202 self._readers[ExtensionKey(JSON_COLUMN)] = _ExtensionReader.from_bytes(
203 astropy.io.fits.BinTableHDU, tail_data[:json_size]
204 )
205 # Make any empty dictionary to cache deserialized objects. Keys are
206 # the zero-indexed row in the JSON table.
207 self._deserialized_pointer_cache: dict[int, Any] = {}
209 @classmethod
210 @contextmanager
211 def open(
212 cls,
213 path: ResourcePathExpression,
214 *,
215 page_size: int = DEFAULT_PAGE_SIZE,
216 partial: bool = False,
217 ) -> Iterator[Self]:
218 """Create an output archive that writes to the given file.
220 Parameters
221 ----------
222 path
223 File to read; convertible to `lsst.resources.ResourcePath`.
224 page_size
225 Size of the fsspec read block for partial (remote) reads, in
226 bytes; a multiple of the FITS block size (2880) is recommended.
227 Defaults to `DEFAULT_PAGE_SIZE`; see it for the tuning tradeoff.
228 partial
229 Whether we will be reading only some of the archive, or if memory
230 pressure forces us to read it only a little at a time. If `False`
231 (default), the entire raw file may be read into memory up front.
233 Returns
234 -------
235 `contextlib.AbstractContextManager` [`FitsInputArchive`]
236 A context manager that returns a `FitsInputArchive` when entered.
237 """
238 path = ResourcePath(path)
239 stream: IO[bytes]
240 if not partial:
241 stream = io.BytesIO(path.read())
242 yield cls(stream)
243 else:
244 fs: fsspec.AbstractFileSystem
245 fs, fp = path.to_fsspec()
246 # Cap cached blocks from the byte budget so memory stays bounded as
247 # the block size is retuned; keep at least two so the shared
248 # header/index block survives between a file's components.
249 maxblocks = max(2, READ_CACHE_MAX_BYTES // page_size)
250 with fs.open(
251 fp,
252 block_size=page_size,
253 cache_type=_READ_CACHE_TYPE,
254 cache_options={"maxblocks": maxblocks},
255 ) as stream:
256 yield cls(stream)
258 @property
259 def info(self) -> ArchiveInfo:
260 """Schema/format info read from the primary header on open.
261 (`.serialization.ArchiveInfo`)
262 """
263 if self._info is None: 263 ↛ 264line 263 didn't jump to line 264 because the condition on line 263 was never true
264 raise ArchiveReadError("This is not an lsst.images FITS archive (no DATAMODL card).")
265 return self._info
267 def get_tree[T: ArchiveTree](self, model_type: type[T]) -> T:
268 """Read the JSON tree from the archive.
270 Parameters
271 ----------
272 model_type
273 A Pydantic model type to use to validate the JSON.
275 Returns
276 -------
277 T
278 The validated Pydantic model.
279 """
280 json_bytes = self._readers[ExtensionKey(JSON_EXTNAME)].data[0][JSON_COLUMN].tobytes()
281 return model_type.model_validate_json(json_bytes)
283 def deserialize_pointer[U: ArchiveTree, V](
284 self,
285 pointer: PointerModel,
286 model_type: type[U],
287 deserializer: Callable[[U, InputArchive[PointerModel]], V],
288 ) -> V:
289 # Docstring inherited.
290 if (cached := self._deserialized_pointer_cache.get(pointer.row)) is not None:
291 return cached
292 if not isinstance(pointer.column.data, ArrayReferenceModel):
293 raise ArchiveReadError(f"Invalid pointer with inline array:\n{pointer.model_dump_json(indent=2)}")
294 _, reader = self._get_source_reader(pointer.column.data.source, is_table=True)
295 try:
296 json_bytes = reader.data[pointer.row][JSON_COLUMN].tobytes()
297 except Exception as err:
298 raise InvalidFitsArchiveError(
299 f"Failed to access the table cell referenced by {pointer.model_dump_json()}."
300 ) from err
301 result = deserializer(model_type.model_validate_json(json_bytes), self)
302 self._deserialized_pointer_cache[pointer.row] = result
303 return result
305 def get_frame_set(self, ref: PointerModel) -> FrameSet:
306 try:
307 result = self._deserialized_pointer_cache[ref.row]
308 except KeyError:
309 raise AssertionError(
310 f"Frame set at {ref.model_dump_json(indent=2)} must be deserialized "
311 "before any dependent transform can be."
312 ) from None
313 if not isinstance(result, FrameSet):
314 raise InvalidFitsArchiveError(f"Expected a FrameSet instance at {ref.model_dump_json(indent=2)}.")
315 return result
317 def get_array(
318 self,
319 model: ArrayReferenceModel | InlineArrayModel,
320 *,
321 slices: tuple[slice, ...] | EllipsisType = ...,
322 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
323 ) -> np.ndarray:
324 if not isinstance(model, ArrayReferenceModel): 324 ↛ 325line 324 didn't jump to line 325 because the condition on line 324 was never true
325 raise ArchiveReadError("Inline array found where a reference array was expected.")
326 key, reader = self._get_source_reader(model.source, is_table=False)
327 if slices is not ...:
328 array = reader.section[slices]
329 else:
330 array = reader.data
331 if key not in self._opaque_metadata.headers: 331 ↛ 335line 331 didn't jump to line 335 because the condition on line 331 was always true
332 opaque_header = reader.header.copy(strip=True)
333 strip_header(opaque_header)
334 self._opaque_metadata.add_header(opaque_header, key=key)
335 return array
337 def get_table(
338 self,
339 model: TableModel,
340 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
341 ) -> astropy.table.Table:
342 # Docstring inherited.
343 array = self.get_structured_array(model, strip_header)
344 table = astropy.table.Table(array)
345 for c in model.columns:
346 c.update_table(table)
347 return table
349 def get_structured_array(
350 self,
351 model: TableModel,
352 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
353 ) -> np.ndarray:
354 # Docstring inherited.
355 if not isinstance(model.columns[0].data, ArrayReferenceModel):
356 raise ArchiveReadError("Inline array found where a reference array was expected.")
357 # All columns should have the same data.source; just use the first.
358 key, reader = self._get_source_reader(model.columns[0].data.source, is_table=True)
359 if key not in self._opaque_metadata.headers:
360 opaque_header = reader.header.copy(strip=True)
361 strip_header(opaque_header)
362 self._opaque_metadata.add_header(opaque_header, key=key)
363 return reader.hdu.data
365 def get_opaque_metadata(self) -> FitsOpaqueMetadata:
366 # Docstring inherited.
367 return self._opaque_metadata
369 def _get_source_reader(self, source: str | int, is_table: bool) -> tuple[ExtensionKey, _ExtensionReader]:
370 """Get a reader for the extension referenced by a serialiation model's
371 ``source`` field.
373 Parameters
374 ----------
375 source
376 A ``source`` field of the form ``fits:${hdu}`` or
377 ``fits:${hdu}[${col}]``.
378 is_table
379 Whether the source should be for a table HDU.
381 Returns
382 -------
383 key
384 Identifier pair for the HDU (EXTNAME, EXTVER).
385 reader
386 A reader object for the extension.
387 """
388 if not isinstance(source, str): 388 ↛ 389line 388 didn't jump to line 389 because the condition on line 388 was never true
389 raise InvalidFitsArchiveError(f"Reference with source={source!r} is not a string.")
390 if not source.startswith("fits:"): 390 ↛ 391line 390 didn't jump to line 391 because the condition on line 390 was never true
391 raise InvalidFitsArchiveError(f"Reference with source={source!r} does not start with 'fits:'.")
392 key = ExtensionKey.from_str(source)
393 try:
394 reader = self._readers[key]
395 except KeyError:
396 raise InvalidFitsArchiveError(f"Unrecognized source value {key}.") from None
397 if is_table and not reader.is_table: 397 ↛ 398line 397 didn't jump to line 398 because the condition on line 397 was never true
398 raise InvalidFitsArchiveError(
399 f"Extension with source={key} was expected to be be a binary table, not an image."
400 )
401 elif not is_table and reader.is_table: 401 ↛ 402line 401 didn't jump to line 402 because the condition on line 401 was never true
402 raise InvalidFitsArchiveError(
403 f"Extension with source={key} was expected to be be an image, not a binary table."
404 )
405 return key, reader
408class _ExtensionReader:
409 """A lazy-load reader for a single extension HDU.
411 Parameters
412 ----------
413 hdu_cls
414 The type of the astropy HDU instance to construct.
415 stream
416 The file-like object to read from.
417 """
419 def __init__(self, hdu_cls: type[ExtensionHDU], stream: IO[bytes]):
420 self._hdu_cls = hdu_cls
421 self._stream = stream
423 @classmethod
424 def from_index_row(cls, index_row: np.void, stream: IO[bytes]) -> _ExtensionReader:
425 """Construct from a row of the binary table index HDU.
427 Parameters
428 ----------
429 index_row
430 A record array row from the index HDU.
431 stream
432 The file-like object being used to read the full FITS file.
434 Returns
435 -------
436 reader
437 A reader object for the extension.
438 """
439 match index_row["XTENSION"].strip():
440 case "IMAGE":
441 hdu_cls = astropy.io.fits.ImageHDU
442 case "BINTABLE": 442 ↛ 447line 442 didn't jump to line 447 because the pattern on line 442 always matched
443 if index_row["ZIMAGE"]:
444 hdu_cls = astropy.io.fits.CompImageHDU
445 else:
446 hdu_cls = astropy.io.fits.BinTableHDU
447 case other:
448 raise AssertionError(f"Unsupported HDU type {other!r}.")
449 return _ExtensionReader(
450 hdu_cls,
451 _RangeStreamProxy(
452 stream,
453 start=int(index_row["HDRADDR"]),
454 ),
455 )
457 @classmethod
458 def from_bytes(cls, hdu_cls: type[ExtensionHDU], data: bytes) -> _ExtensionReader:
459 """Construct from already-read `bytes`
461 Parameters
462 ----------
463 hdu_cls
464 The HDU type to instantiate.
465 data
466 Raw data for the HDU.
468 Returns
469 -------
470 reader
471 A reader object for extension.
472 """
473 return _ExtensionReader(hdu_cls, io.BytesIO(data))
475 @property
476 def is_table(self) -> bool:
477 """Whether this is logically a table HDU.
479 This is `False` for compressed image HDUs, even though they are
480 represented in FITS as a binary table.
481 """
482 return issubclass(self._hdu_cls, astropy.io.fits.BinTableHDU) and not issubclass(
483 self._hdu_cls, astropy.io.fits.CompImageHDU
484 )
486 @cached_property
487 def hdu(self) -> ExtensionHDU:
488 """The Astropy HDU object."""
489 self._stream.seek(0)
490 if self._hdu_cls is astropy.io.fits.CompImageHDU:
491 # CompImageHDU.readfrom doesn't work; we need to make a minimal
492 # example and report it upstream. Happily this workaround does
493 # work.
494 bintable_hdu = astropy.io.fits.BinTableHDU.readfrom(self._stream, memmap=False, cache=False)
495 return self._hdu_cls(bintable=bintable_hdu)
496 else:
497 return self._hdu_cls.readfrom(self._stream, memmap=False, cache=False, uint=True)
499 @property
500 def header(self) -> astropy.io.fits.Header:
501 """The header of the HDU."""
502 return self.hdu.header
504 @property
505 def data(self) -> np.ndarray:
506 """The data for the HDU."""
507 return self.hdu.data
509 @property
510 def section(self) -> astropy.io.fits.Section | astropy.io.fits.CompImageSection:
511 """An Astropy expression object that reads a subset of the data when
512 sliced.
513 """
514 return self.hdu.section
517class _RangeStreamProxy(IO[bytes]):
518 """A readable IO proxy object that makes the beginning of the file appear
519 at a custom position.
521 Parameters
522 ----------
523 base
524 Underlying readable, seekable buffer to proxy.
525 start
526 Offset into the base stream that will be considered the start of the
527 proxy stream.
529 Notes
530 -----
531 This class exists because Astropy doesn't seem to provide a way to read a
532 single HDU that starts at the current seek position of a file-like object.
533 It does provide ``readfrom`` methods on its HDU objects that take a
534 file-like object, but these assume (possibly unintentionally; it only
535 happens when Astropy is trying to see whether the file was opened for
536 appending) that ``seek(0)`` will set the file-like object to the start of
537 the HDU.
538 """
540 def __init__(self, base: IO[bytes], start: int):
541 self._base = base
542 self._start = start
544 @property
545 def mode(self) -> str:
546 return "rb"
548 def __enter__(self) -> Self:
549 raise AssertionError("This proxy should not be used as a context manager.")
551 def __exit__(self, type: Any, value: Any, traceback: Any) -> None:
552 raise AssertionError("This proxy should not be used as a context manager.")
554 def __iter__(self) -> Iterator[bytes]:
555 return self._base.__iter__()
557 def __next__(self) -> bytes:
558 return self._base.__next__()
560 def close(self) -> None:
561 raise AssertionError("This proxy should not ever be closed.")
563 @property
564 def closed(self) -> bool:
565 return False
567 def fileno(self) -> int:
568 raise OSError()
570 def flush(self) -> None:
571 pass
573 def isatty(self) -> bool:
574 return False
576 def read(self, n: int = -1, /) -> bytes:
577 result = self._base.read(n)
578 return result
580 def readable(self) -> bool:
581 return True
583 def readline(self, limit: int = -1, /) -> bytes:
584 return self._base.readline(limit)
586 def readlines(self, hint: int = -1, /) -> list[bytes]:
587 return self._base.readlines(hint)
589 def seek(self, offset: int, whence: int = 0) -> int:
590 match whence:
591 case os.SEEK_SET:
592 return self._base.seek(offset + self._start, os.SEEK_SET) - self._start
593 case os.SEEK_CUR: 593 ↛ 594line 593 didn't jump to line 594 because the pattern on line 593 never matched
594 return self._base.seek(offset, os.SEEK_CUR) - self._start
595 case os.SEEK_END: 595 ↛ 597line 595 didn't jump to line 597 because the pattern on line 595 always matched
596 return self._base.seek(offset, os.SEEK_END) - self._start
597 raise TypeError(f"Invalid value for 'whence': {whence}.")
599 def seekable(self) -> bool:
600 return True
602 def tell(self) -> int:
603 return self._base.tell() - self._start
605 def truncate(self, size: int | None = None, /) -> int:
606 raise OSError()
608 def writable(self) -> bool:
609 return False
611 def write(self, arg: Any, /) -> int:
612 raise OSError()
614 def writelines(self, arg: Any, /) -> None:
615 raise OSError()