381 lines
12 KiB
Python
381 lines
12 KiB
Python
"""
|
|
Lazy GDSII readers and writers backed by native Arrow scan/materialize paths.
|
|
|
|
This module is intentionally separate from `gdsii.arrow` so the eager read path
|
|
keeps its current behavior and performance profile.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import IO, TYPE_CHECKING, Any
|
|
import gzip
|
|
import logging
|
|
import mmap
|
|
import pathlib
|
|
|
|
import numpy
|
|
|
|
from . import arrow
|
|
from ..utils import is_gzipped
|
|
from ...library import (
|
|
ILibraryView,
|
|
IMaterializable,
|
|
LibraryView,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterator, Sequence
|
|
|
|
from numpy.typing import NDArray
|
|
import pyarrow
|
|
|
|
from ...pattern import Pattern
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _StructRange:
|
|
start: int
|
|
end: int
|
|
|
|
|
|
@dataclass
|
|
class _SourceBuffer:
|
|
path: pathlib.Path
|
|
data: bytes | mmap.mmap
|
|
handle: IO[bytes] | None = None
|
|
|
|
def raw_slice(self, start: int, end: int) -> bytes:
|
|
return self.data[start:end]
|
|
|
|
|
|
@dataclass
|
|
class _ScanRefs:
|
|
offsets: NDArray[numpy.integer[Any]]
|
|
targets: NDArray[numpy.integer[Any]]
|
|
xy: NDArray[numpy.int32]
|
|
xy0: NDArray[numpy.int32]
|
|
xy1: NDArray[numpy.int32]
|
|
counts: NDArray[numpy.int64]
|
|
invert_y: NDArray[numpy.bool_ | numpy.bool]
|
|
angle_rad: NDArray[numpy.floating[Any]]
|
|
scale: NDArray[numpy.floating[Any]]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _CellScan:
|
|
cell_id: int
|
|
struct_range: _StructRange
|
|
ref_start: int
|
|
ref_stop: int
|
|
children: set[str]
|
|
|
|
|
|
@dataclass
|
|
class _ScanPayload:
|
|
libarr: pyarrow.StructScalar
|
|
library_info: dict[str, Any]
|
|
cell_names: list[str]
|
|
cell_order: list[str]
|
|
cells: dict[str, _CellScan]
|
|
refs: _ScanRefs
|
|
|
|
def _open_source_buffer(path: pathlib.Path) -> _SourceBuffer:
|
|
if is_gzipped(path):
|
|
with gzip.open(path, mode='rb') as stream:
|
|
data = stream.read()
|
|
return _SourceBuffer(path=path, data=data)
|
|
|
|
handle = path.open(mode='rb', buffering=0)
|
|
mapped = mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ)
|
|
return _SourceBuffer(path=path, data=mapped, handle=handle)
|
|
|
|
|
|
def _extract_scan_payload(libarr: pyarrow.StructScalar) -> _ScanPayload:
|
|
library_info = arrow._read_header(libarr)
|
|
cell_names = libarr['cell_names'].as_py()
|
|
|
|
cells = libarr['cells']
|
|
cell_values = cells.values
|
|
cell_ids = cell_values.field('id').to_numpy()
|
|
struct_starts = cell_values.field('struct_start_offset').to_numpy()
|
|
struct_ends = cell_values.field('struct_end_offset').to_numpy()
|
|
|
|
refs = cell_values.field('refs')
|
|
ref_values = refs.values
|
|
ref_offsets = refs.offsets.to_numpy()
|
|
targets = ref_values.field('target').to_numpy()
|
|
xy = arrow._packed_xy_u64_to_pairs(ref_values.field('xy').to_numpy())
|
|
xy0 = arrow._packed_xy_u64_to_pairs(ref_values.field('xy0').to_numpy())
|
|
xy1 = arrow._packed_xy_u64_to_pairs(ref_values.field('xy1').to_numpy())
|
|
counts = arrow._packed_counts_u32_to_pairs(ref_values.field('counts').to_numpy())
|
|
invert_y = ref_values.field('invert_y').to_numpy(zero_copy_only=False)
|
|
angle_rad = ref_values.field('angle_rad').to_numpy()
|
|
scale = ref_values.field('scale').to_numpy()
|
|
|
|
ref_payload = _ScanRefs(
|
|
offsets=ref_offsets,
|
|
targets=targets,
|
|
xy=xy,
|
|
xy0=xy0,
|
|
xy1=xy1,
|
|
counts=counts,
|
|
invert_y=invert_y,
|
|
angle_rad=angle_rad,
|
|
scale=scale,
|
|
)
|
|
|
|
cell_order = [cell_names[int(cell_id)] for cell_id in cell_ids]
|
|
cell_scan: dict[str, _CellScan] = {}
|
|
for cc, name in enumerate(cell_order):
|
|
ref_start = int(ref_offsets[cc])
|
|
ref_stop = int(ref_offsets[cc + 1])
|
|
children = {
|
|
cell_names[int(target)]
|
|
for target in targets[ref_start:ref_stop]
|
|
}
|
|
cell_scan[name] = _CellScan(
|
|
cell_id=int(cell_ids[cc]),
|
|
struct_range=_StructRange(int(struct_starts[cc]), int(struct_ends[cc])),
|
|
ref_start=ref_start,
|
|
ref_stop=ref_stop,
|
|
children=children,
|
|
)
|
|
|
|
return _ScanPayload(
|
|
libarr=libarr,
|
|
library_info=library_info,
|
|
cell_names=cell_names,
|
|
cell_order=cell_order,
|
|
cells=cell_scan,
|
|
refs=ref_payload,
|
|
)
|
|
|
|
def _make_ref_rows(
|
|
xy: NDArray[numpy.integer[Any]],
|
|
angle_rad: NDArray[numpy.floating[Any]],
|
|
invert_y: NDArray[numpy.bool_ | numpy.bool],
|
|
scale: NDArray[numpy.floating[Any]],
|
|
) -> NDArray[numpy.float64]:
|
|
rows = numpy.empty((len(xy), 5), dtype=float)
|
|
rows[:, :2] = xy
|
|
rows[:, 2] = angle_rad
|
|
rows[:, 3] = invert_y.astype(float)
|
|
rows[:, 4] = scale
|
|
return rows
|
|
|
|
|
|
def _expand_aref_row(
|
|
xy: NDArray[numpy.integer[Any]],
|
|
xy0: NDArray[numpy.integer[Any]],
|
|
xy1: NDArray[numpy.integer[Any]],
|
|
counts: NDArray[numpy.integer[Any]],
|
|
angle_rad: float,
|
|
invert_y: bool,
|
|
scale: float,
|
|
) -> NDArray[numpy.float64]:
|
|
a_count = int(counts[0])
|
|
b_count = int(counts[1])
|
|
aa, bb = numpy.meshgrid(numpy.arange(a_count), numpy.arange(b_count), indexing='ij')
|
|
displacements = aa.reshape(-1, 1) * xy0[None, :] + bb.reshape(-1, 1) * xy1[None, :]
|
|
rows = numpy.empty((displacements.shape[0], 5), dtype=float)
|
|
rows[:, :2] = xy + displacements
|
|
rows[:, 2] = angle_rad
|
|
rows[:, 3] = float(invert_y)
|
|
rows[:, 4] = scale
|
|
return rows
|
|
|
|
|
|
class ArrowLibrary(ILibraryView, IMaterializable):
|
|
"""
|
|
Read-only library backed by the native lazy Arrow scan schema.
|
|
|
|
Materializing a cell via `__getitem__` caches a real `Pattern` for that cell.
|
|
Cached cells are treated as edited for future writes from this module.
|
|
"""
|
|
|
|
path: pathlib.Path
|
|
library_info: dict[str, Any]
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
path: pathlib.Path,
|
|
payload: _ScanPayload,
|
|
source: _SourceBuffer,
|
|
) -> None:
|
|
self.path = path
|
|
self.library_info = payload.library_info
|
|
self._payload = payload
|
|
self._source = source
|
|
self._cache: dict[str, Pattern] = {}
|
|
|
|
@classmethod
|
|
def from_file(cls, filename: str | pathlib.Path) -> ArrowLibrary:
|
|
path = pathlib.Path(filename).expanduser().resolve()
|
|
source = _open_source_buffer(path)
|
|
scan_arr = arrow._scan_buffer_to_arrow(source.data)
|
|
assert len(scan_arr) == 1
|
|
payload = _extract_scan_payload(scan_arr[0])
|
|
return cls(path=path, payload=payload, source=source)
|
|
|
|
def __getitem__(self, key: str) -> Pattern:
|
|
return self.materialize(key, persist=True)
|
|
|
|
def __iter__(self) -> Iterator[str]:
|
|
return iter(self._payload.cell_order)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._payload.cell_order)
|
|
|
|
def __contains__(self, key: object) -> bool:
|
|
return key in self._payload.cells
|
|
|
|
def source_order(self) -> tuple[str, ...]:
|
|
return tuple(self._payload.cell_order)
|
|
|
|
def raw_struct_bytes(self, name: str) -> bytes:
|
|
struct_range = self._payload.cells[name].struct_range
|
|
return self._source.raw_slice(struct_range.start, struct_range.end)
|
|
|
|
def can_copy_raw_struct(self, name: str) -> bool:
|
|
return name not in self._cache
|
|
|
|
def materialize_many(
|
|
self,
|
|
names: Sequence[str],
|
|
*,
|
|
persist: bool = True,
|
|
) -> LibraryView:
|
|
mats = self._materialize_patterns(names, persist=persist, detached=False)
|
|
return LibraryView(mats)
|
|
|
|
def materialize_many_detached(
|
|
self,
|
|
names: Sequence[str],
|
|
) -> LibraryView:
|
|
mats = self._materialize_patterns(names, persist=False, detached=True)
|
|
return LibraryView(mats)
|
|
|
|
def _materialize_patterns(
|
|
self,
|
|
names: Sequence[str],
|
|
*,
|
|
persist: bool,
|
|
detached: bool,
|
|
) -> dict[str, Pattern]:
|
|
ordered_names = list(dict.fromkeys(names))
|
|
missing = [name for name in ordered_names if name not in self._payload.cells]
|
|
if missing:
|
|
raise KeyError(missing[0])
|
|
|
|
materialized: dict[str, Pattern] = {}
|
|
uncached = [name for name in ordered_names if name not in self._cache]
|
|
if uncached:
|
|
ranges = numpy.asarray(
|
|
[
|
|
[
|
|
self._payload.cells[name].struct_range.start,
|
|
self._payload.cells[name].struct_range.end,
|
|
]
|
|
for name in uncached
|
|
],
|
|
dtype=numpy.uint64,
|
|
)
|
|
arrow_arr = arrow._read_selected_cells_to_arrow(self._source.data, ranges)
|
|
assert len(arrow_arr) == 1
|
|
selected_lib, _info = arrow.read_arrow(arrow_arr[0])
|
|
for name in uncached:
|
|
pat = selected_lib[name]
|
|
materialized[name] = pat
|
|
if persist:
|
|
self._cache[name] = pat
|
|
|
|
for name in ordered_names:
|
|
if name not in materialized:
|
|
cached = self._cache[name]
|
|
materialized[name] = cached.deepcopy() if detached else cached
|
|
return materialized
|
|
|
|
def materialize(self, name: str, *, persist: bool = True) -> Pattern:
|
|
return self._materialize_patterns((name,), persist=persist, detached=False)[name]
|
|
|
|
def materialize_detached(self, name: str) -> Pattern:
|
|
return self._materialize_patterns((name,), persist=False, detached=True)[name]
|
|
|
|
def _raw_children(self, name: str) -> set[str]:
|
|
if name in self._cache:
|
|
return super()._raw_children(name)
|
|
return set(self._payload.cells[name].children)
|
|
|
|
def _collect_raw_transforms(self, cell: _CellScan, target_id: int) -> list[NDArray[numpy.float64]]:
|
|
refs = self._payload.refs
|
|
start = cell.ref_start
|
|
stop = cell.ref_stop
|
|
if stop <= start:
|
|
return []
|
|
|
|
targets = refs.targets[start:stop]
|
|
mask = targets == target_id
|
|
if not mask.any():
|
|
return []
|
|
|
|
rows: list[NDArray[numpy.float64]] = []
|
|
counts = refs.counts[start:stop]
|
|
unit_mask = mask & (counts[:, 0] == 1) & (counts[:, 1] == 1)
|
|
if unit_mask.any():
|
|
rows.append(_make_ref_rows(
|
|
refs.xy[start:stop][unit_mask],
|
|
refs.angle_rad[start:stop][unit_mask],
|
|
refs.invert_y[start:stop][unit_mask],
|
|
refs.scale[start:stop][unit_mask],
|
|
))
|
|
|
|
aref_indices = numpy.nonzero(mask & ~unit_mask)[0]
|
|
for idx in aref_indices:
|
|
abs_idx = start + int(idx)
|
|
rows.append(_expand_aref_row(
|
|
xy=refs.xy[abs_idx],
|
|
xy0=refs.xy0[abs_idx],
|
|
xy1=refs.xy1[abs_idx],
|
|
counts=refs.counts[abs_idx],
|
|
angle_rad=float(refs.angle_rad[abs_idx]),
|
|
invert_y=bool(refs.invert_y[abs_idx]),
|
|
scale=float(refs.scale[abs_idx]),
|
|
))
|
|
return rows
|
|
|
|
def close(self) -> None:
|
|
data = self._source.data
|
|
if isinstance(data, mmap.mmap):
|
|
data.close()
|
|
if self._source.handle is not None:
|
|
self._source.handle.close()
|
|
self._source.handle = None
|
|
|
|
def __enter__(self) -> ArrowLibrary:
|
|
return self
|
|
|
|
def __exit__(self, *_args: object) -> None:
|
|
self.close()
|
|
|
|
def _raw_ref_transforms(
|
|
self,
|
|
parent: str,
|
|
target: str,
|
|
) -> list[NDArray[numpy.float64]]:
|
|
if parent in self._cache:
|
|
return super()._raw_ref_transforms(parent, target)
|
|
target_cell = self._payload.cells.get(target)
|
|
if target_cell is None or parent not in self._payload.cells:
|
|
return []
|
|
return self._collect_raw_transforms(self._payload.cells[parent], target_cell.cell_id)
|
|
|
|
|
|
def readfile(
|
|
filename: str | pathlib.Path,
|
|
) -> tuple[ArrowLibrary, dict[str, Any]]:
|
|
lib = ArrowLibrary.from_file(filename)
|
|
return lib, lib.library_info
|