"""Narrow adapter around the ArrayRecord Python binding."""
from __future__ import annotations
import importlib
import logging
from collections.abc import Iterable, Sequence
from pathlib import Path
from typing import Any
from signal_dataset.storage.contracts import ShardStore
from signal_dataset.storage.staging import ShardStagingArea, default_staging
logger = logging.getLogger(__name__)
def _module() -> Any:
candidates = ("array_record.python.array_record_module",)
failures: list[str] = []
for candidate in candidates:
try:
return importlib.import_module(candidate)
except ImportError as exc:
failures.append(f"{candidate}: {exc}")
raise ImportError("ArrayRecord native binding is unavailable (" + "; ".join(failures) + ")")
def reader_options(file_reader_buffer_size: int | None) -> dict[str, int]:
if file_reader_buffer_size is None:
return {}
return {"file_reader_buffer_size": file_reader_buffer_size}
[docs]
class ArrayRecordShardStore(ShardStore):
"""Reads indexed records through a staging area.
The binding opens a path and knows one remote scheme, so a shard stored
anywhere else has to be materialized first. Staging is injected rather than
wrapped around this class: ``ShardStore.read_many`` falls back to calling
``read`` per index, so a decorator that staged in both would download the
same object once per record.
"""
# A class attribute, so a subclass that does not call super().__init__()
# still resolves to the default rather than raising AttributeError.
_staging: ShardStagingArea | None = None
def __init__(self, *, staging: ShardStagingArea | None = None) -> None:
self._staging = staging
@property
def staging(self) -> ShardStagingArea:
if self._staging is None:
self._staging = default_staging()
return self._staging
[docs]
def write(self, path: Path, records: Iterable[bytes], *, options: str) -> int:
logger.debug("Writing ArrayRecord shard", extra={"path": str(path)})
writer = _module().ArrayRecordWriter(str(path), options)
count = 0
try:
for record in records:
writer.write(record)
count += 1
finally:
writer.close()
return count
[docs]
def read(
self,
uri: str,
index: int,
*,
generation: int | None = None,
options: str,
file_reader_buffer_size: int | None = None,
) -> bytes:
logger.debug("Reading ArrayRecord entry", extra={"uri": uri, "index": index})
with self.staging.stage(uri, generation=generation) as path:
reader = _module().ArrayRecordReader(
path, options, **reader_options(file_reader_buffer_size)
)
try:
try:
values = reader.read([index])
except RuntimeError as exc:
if "out of bound" in str(exc):
raise IndexError(f"ArrayRecord index out of range: {index}") from exc
raise
finally:
reader.close()
if len(values) != 1:
raise OSError("ArrayRecord did not return exactly one requested record")
return bytes(values[0])
[docs]
def read_many(
self,
uri: str,
indices: Sequence[int],
*,
generation: int | None = None,
options: str,
file_reader_buffer_size: int | None = None,
) -> tuple[bytes, ...]:
"""Read ordered entries with one staging and one reader lifetime."""
if not indices:
return ()
logger.debug(
"Reading ArrayRecord entries",
extra={"uri": uri, "record_count": len(indices)},
)
with self.staging.stage(uri, generation=generation) as path:
reader = _module().ArrayRecordReader(
path, options, **reader_options(file_reader_buffer_size)
)
try:
try:
values = reader.read(list(indices))
except RuntimeError as exc:
if "out of bound" in str(exc):
raise IndexError("ArrayRecord index out of range") from exc
raise
finally:
reader.close()
if len(values) != len(indices):
raise OSError("ArrayRecord did not return every requested record")
return tuple(bytes(value) for value in values)
[docs]
def count(
self,
uri: str,
*,
generation: int | None = None,
options: str,
file_reader_buffer_size: int | None = None,
) -> int:
with self.staging.stage(uri, generation=generation) as path:
reader = _module().ArrayRecordReader(
path, options, **reader_options(file_reader_buffer_size)
)
try:
return int(reader.num_records())
finally:
reader.close()