Back to report index

iolabs-common 639d755: AI3D-382 Drive the point-record contract from a schema registry

Miroslav Simko <ms@iolabs.ch> 2026-09-01T14:54:43+02:00

Commit #3 ยท 12 snippets

 src/iolabs/common/segment_points_io.py | 398 ++++++++++++++++++++++++---------
 tests/test_segment_points_io.py        | 205 +++++++++++++++++
 2 files changed, 496 insertions(+), 107 deletions(-)

Largest structural change. Replaces per-key special-casing in common with a POINT_RECORD_SCHEMA registry (key, dtype, required/optional, fill) and adds generic mask_record/concat_records helpers driven by it. Later producer/consumer commits iterate this registry instead of hardcoding names. Big diff but mostly mechanical; the registry table and the load/save loops are the parts worth reading.

Importance #1: src/iolabs/common/segment_points_io.py @@ -25,41 +34,115 @@

The POINT_RECORD_SCHEMA registry. One entry per NPZ member; this is now the single source of truth for keys, dtypes and fill values.

25and consumers working in the local frame ignore it. This module never34and consumers working in the local frame ignore it. This module never
26applies, negates, or bakes in a sign.35applies, negates, or bakes in a sign.
27"""36"""
2837
38import functools
29import json39import json
30import logging40import logging
31import zipfile41import zipfile
32from collections.abc import Iterator, Mapping, Sequence42from collections.abc import Callable, Iterator, Mapping, Sequence
33from dataclasses import dataclass43from dataclasses import dataclass
34from fnmatch import fnmatch44from fnmatch import fnmatch
35from pathlib import Path45from pathlib import Path
46from types import MappingProxyType
3647
37import numpy as np48import numpy as np
3849
39logger = logging.getLogger(__name__)50logger = logging.getLogger(__name__)
4051
52#: Builds a member for a record that predates its key, given the point count.
53FillFactory = Callable[[int], np.ndarray]
54
55#: The point-record key whose row count defines ``N`` for every other member.
56POINTS_KEY = "points"
57
41#: Per-point LAS return count; see the module docstring for the 0 = unknown rule.58#: Per-point LAS return count; see the module docstring for the 0 = unknown rule.
42NUMBER_OF_RETURNS_KEY = "number_of_returns"59NUMBER_OF_RETURNS_KEY = "number_of_returns"
4360
44#: Storage dtype of :data:`NUMBER_OF_RETURNS_KEY` (LAS carries 3 bits, values 1-7).61#: Storage dtype of :data:`NUMBER_OF_RETURNS_KEY` (LAS carries 3 bits, values 1-7).
45NUMBER_OF_RETURNS_DTYPE = np.uint862NUMBER_OF_RETURNS_DTYPE = np.uint8
4663
47#: Keys every record must already carry; a file missing one of these is corrupt.64
48REQUIRED_POINT_RECORD_KEYS: tuple[str, ...] = (65@dataclass(frozen=True)
49 "points",66class PointFieldSpec:
50 "red",67 """How one member of the point record is validated, stored and synthesised.
51 "green",68
52 "blue",69 Attributes:
53 "intensity",70 required: Whether :func:`load_points_npz` rejects a file that lacks the
54 "scan_angle",71 key. Optional keys are ones added after datasets were already on
72 disk; they are synthesised via :attr:`fill` instead.
73 storage_dtype: The dtype the member is coerced to -- and validated
74 against -- on both save and load, so one contract governs the bytes
75 on disk and the array handed to callers. ``None`` accepts whatever
76 dtype the producer used (``red``/``green``/``blue``/``intensity``/
77 ``scan_angle``, whose historical widths vary across datasets).
78 fill: Factory called with the point count to build the member for a
79 record that does not carry it. Required for optional keys, unused
80 for required ones.
81 columns: ``None`` for a 1-D ``(N,)`` member; an int for a 2-D
82 ``(N, columns)`` one (``points`` is the only such key today).
83 noun: Singular name of one stored value, used in dtype error messages
84 ("a boolean is a mask, not a return count").
85 """
86
87 required: bool
88 storage_dtype: np.dtype | None = None
89 fill: FillFactory | None = None
90 columns: int | None = None
91 noun: str = "value"
92
93 def __post_init__(self) -> None:
94 """Reject a spec an optional key could not be synthesised from.
95
96 Raises:
97 ValueError: The key is optional but declares no fill factory, which
98 would make a record that predates it unloadable.
99 """
100 if not self.required and self.fill is None:
101 raise ValueError(
102 "An optional point-record field needs a fill factory: records written "
103 "before the key existed have nothing to load for it"
104 )
105
106
107#: The point-record contract: ordered key -> spec. Adding a field to the NPZ
108#: schema means adding an entry here (plus teaching producers to emit it); the
109#: load/save/merge machinery below is entirely registry-driven.
110POINT_RECORD_SCHEMA: Mapping[str, PointFieldSpec] = MappingProxyType(
111 {
112 POINTS_KEY: PointFieldSpec(required=True, columns=3, noun="coordinate"),
113 "red": PointFieldSpec(required=True),
114 "green": PointFieldSpec(required=True),
115 "blue": PointFieldSpec(required=True),
116 "intensity": PointFieldSpec(required=True),
117 "scan_angle": PointFieldSpec(required=True),
118 NUMBER_OF_RETURNS_KEY: PointFieldSpec(
119 required=False,
120 storage_dtype=np.dtype(NUMBER_OF_RETURNS_DTYPE),
121 fill=functools.partial(np.zeros, dtype=NUMBER_OF_RETURNS_DTYPE),
122 noun="return count",
123 ),
124 }
55)125)
56126
57#: Keys added after datasets were already on disk: zero-filled when absent on load.127
58OPTIONAL_POINT_RECORD_KEYS: tuple[str, ...] = (NUMBER_OF_RETURNS_KEY,)128def _schema_keys(
129 schema: Mapping[str, PointFieldSpec], *, required: bool | None = None
130) -> tuple[str, ...]:
131 """List the schema's keys in registry order, optionally by required-ness."""
132 return tuple(
133 key for key, spec in schema.items() if required is None or spec.required is required
134 )
135
136
137#: Keys every record must already carry; a file missing one of these is corrupt.
138REQUIRED_POINT_RECORD_KEYS: tuple[str, ...] = _schema_keys(POINT_RECORD_SCHEMA, required=True)
139
140#: Keys added after datasets were already on disk: filled when absent on load.
141OPTIONAL_POINT_RECORD_KEYS: tuple[str, ...] = _schema_keys(POINT_RECORD_SCHEMA, required=False)
59142
60#: Every key a record carries once loaded, and every key that is written.143#: Every key a record carries once loaded, and every key that is written.
61POINT_RECORD_KEYS: tuple[str, ...] = REQUIRED_POINT_RECORD_KEYS + OPTIONAL_POINT_RECORD_KEYS144POINT_RECORD_KEYS: tuple[str, ...] = _schema_keys(POINT_RECORD_SCHEMA)
62145
63PointRecord = dict[str, np.ndarray]146PointRecord = dict[str, np.ndarray]
64147
65GEOSHIFT_NAME = "run3_geoshift.json"148GEOSHIFT_NAME = "run3_geoshift.json"
Importance #2: src/iolabs/common/segment_points_io.py @@ -67,143 +150,248 @@

Load/save rewritten as loops over the registry; also introduces mask_record/concat_records.

67RUN3_POINTS_SUFFIX = "_run3_points.npz"150RUN3_POINTS_SUFFIX = "_run3_points.npz"
68RUN3_POINTS_GLOB = f"*{RUN3_POINTS_SUFFIX}"151RUN3_POINTS_GLOB = f"*{RUN3_POINTS_SUFFIX}"
69152
70153
71def _coerce_number_of_returns(values: np.ndarray, *, source: str) -> np.ndarray:154def _coerce_storage_dtype(
72 """Cast a return-count array to the contract's uint8 storage dtype.155 values: np.ndarray, *, key: str, spec: PointFieldSpec, source: str
156) -> np.ndarray:
157 """Cast one member to its declared storage dtype, rejecting lossy inputs.
158
159 Keys whose spec declares no storage dtype are passed through untouched. For
160 an integer storage dtype the input must itself be integral -- a bool array
161 is a mask, and a float array is not a count -- and must fit the target
162 range, because ``astype`` would otherwise wrap silently.
73163
74 Args:164 Args:
75 values: Array-like of per-point return counts.165 values: Array-like member as the producer supplied or the file stored it.
166 key: Member name, used in error messages.
167 spec: The key's registry entry.
76 source: Label used in error messages (typically the target path).168 source: Label used in error messages (typically the target path).
77169
78 Returns:170 Returns:
79 The values as uint8, unchanged when they already are.171 The values as the declared dtype, unchanged when they already are it or
172 when the spec declares none.
80173
81 Raises:174 Raises:
82 ValueError: The array is boolean or of non-integer dtype, or holds a175 ValueError: The declared dtype is integral and the values are boolean,
83 value outside the uint8 range (which would wrap silently on cast).176 non-integral, or outside the target dtype's range.
84 """177 """
85 array = np.asarray(values)178 array = np.asarray(values)
86 if array.dtype == NUMBER_OF_RETURNS_DTYPE:179 dtype = spec.storage_dtype
180 if dtype is None or array.dtype == dtype:
87 return array181 return array
88 if array.dtype.kind == "b":182 if dtype.kind in "ui":
89 raise ValueError(183 if array.dtype.kind == "b":
90 f"{source}: '{NUMBER_OF_RETURNS_KEY}' must be an integer array, got a bool "184 raise ValueError(
91 "array; a boolean is a mask, not a return count (casting it would fabricate "185 f"{source}: '{key}' must be an integer array, got a bool array; a boolean "
92 "counts of 0 and 1)"186 f"is a mask, not a {spec.noun} (casting it would fabricate {spec.noun}s of "
93 )187 "0 and 1)"
94 if array.dtype.kind not in "ui":188 )
95 raise ValueError(189 if array.dtype.kind not in "ui":
96 f"{source}: '{NUMBER_OF_RETURNS_KEY}' must be an integer array, "190 raise ValueError(
97 f"got dtype {array.dtype}"191 f"{source}: '{key}' must be an integer array, got dtype {array.dtype}"
98 )192 )
99 if array.size and (int(array.min()) < 0 or int(array.max()) > 255):193 info = np.iinfo(dtype)
194 if array.size and (int(array.min()) < info.min or int(array.max()) > info.max):
195 raise ValueError(
196 f"{source}: '{key}' values must fit in {dtype.name}, got range "
197 f"[{int(array.min())}, {int(array.max())}]"
198 )
199 return array.astype(dtype)
200
201
202def _shape_text(spec: PointFieldSpec) -> str:
203 """Render a spec's expected shape for an error message."""
204 return "(N,)" if spec.columns is None else f"(N, {spec.columns})"
205
206
207def _validate_record_shapes(
208 record: Mapping[str, np.ndarray],
209 schema: Mapping[str, PointFieldSpec],
210 *,
211 source: str,
212) -> int:
213 """Check every present member against its spec's shape and return ``N``.
214
215 Args:
216 record: Members to check; keys must all be in *schema*.
217 schema: The registry the record is validated against.
218 source: Label used in error messages (typically the record's path).
219
220 Returns:
221 The point count taken from :data:`POINTS_KEY`.
222
223 Raises:
224 ValueError: ``points`` is not ``(N, 3)``, or a member's shape disagrees
225 with its spec.
226 """
227 anchor_spec = schema[POINTS_KEY]
228 anchor = record[POINTS_KEY]
229 if anchor.ndim != 2 or anchor.shape[1] != anchor_spec.columns:
100 raise ValueError(230 raise ValueError(
101 f"{source}: '{NUMBER_OF_RETURNS_KEY}' values must fit in uint8, got range "231 f"{source}: '{POINTS_KEY}' must have shape {_shape_text(anchor_spec)}, "
102 f"[{int(array.min())}, {int(array.max())}]"232 f"got {anchor.shape}"
103 )233 )
104 return array.astype(NUMBER_OF_RETURNS_DTYPE)234 n_points = int(anchor.shape[0])
235 for key, array in record.items():
236 spec = schema[key]
237 expected = (n_points,) if spec.columns is None else (n_points, spec.columns)
238 if array.shape != expected:
239 raise ValueError(
240 f"{source}: '{key}' must have shape {_shape_text(spec)}, got {array.shape}"
241 )
242 return n_points
105243
106244
107def load_points_npz(path: str | Path) -> PointRecord:245def load_points_npz(path: str | Path) -> PointRecord:
108 """Load one ``*_run3_points.npz``-style file and validate the schema.246 """Load one ``*_run3_points.npz``-style file and validate the schema.
109247
110 Requires all :data:`REQUIRED_POINT_RECORD_KEYS`. ``points`` must be shape248 Every required key of :data:`POINT_RECORD_SCHEMA` must be stored; ``points``
111 ``(N, 3)``; every ancillary array (``red``, ``green``, ``blue``,249 must be shape ``(N, 3)`` and every ancillary array (``red``, ``green``,
112 ``intensity``, ``scan_angle``, ``number_of_returns``) must be 1-D with250 ``blue``, ``intensity``, ``scan_angle``, ``number_of_returns``) 1-D with
113 shape ``(N,)``.251 shape ``(N,)``.
114252
115 :data:`OPTIONAL_POINT_RECORD_KEYS` -- today just ``number_of_returns`` --253 :data:`OPTIONAL_POINT_RECORD_KEYS` -- today just ``number_of_returns`` --
116 may be absent: records written before AI3D-382 predate the key, and are254 may be absent: records written before AI3D-382 predate the key, and are
117 filled with a zeros uint8 array (0 = unknown, see the module docstring)255 built from the spec's fill factory (zeros uint8, 0 = unknown, see the
118 rather than rejected. When the key *is* present it is validated and cast256 module docstring) rather than rejected. When such a key *is* present it is
119 to uint8 under the same rules :func:`save_points_npz` applies, so a257 validated and cast to its declared storage dtype under the same rules
120 hand-rolled producer's wider (or boolean) dtype cannot leak downstream and258 :func:`save_points_npz` applies, so a hand-rolled producer's wider (or
121 break a mixed-vintage concatenation. The returned mapping always holds259 boolean) dtype cannot leak downstream and break a mixed-vintage
122 every key in :data:`POINT_RECORD_KEYS`, with ``number_of_returns`` always260 concatenation. The returned mapping always holds every key in
123 uint8.261 :data:`POINT_RECORD_KEYS`, in registry order.
124262
125 Raises:263 Raises:
126 ValueError: A required key is missing, an array has the wrong shape, or264 ValueError: A required key is missing, an array has the wrong shape, or
127 ``number_of_returns`` is stored with a non-integer/bool dtype or265 a member with a declared storage dtype is stored with a
128 values outside the uint8 range.266 non-integer/bool dtype or values outside that dtype's range.
129 """267 """
130 npz_path = Path(path)268 npz_path = Path(path)
269 schema = POINT_RECORD_SCHEMA
270 required_keys = _schema_keys(schema, required=True)
131 with np.load(npz_path) as data:271 with np.load(npz_path) as data:
132 missing = [key for key in REQUIRED_POINT_RECORD_KEYS if key not in data.files]272 missing = [key for key in required_keys if key not in data.files]
133 if missing:273 if missing:
134 raise ValueError(274 raise ValueError(
135 f"{npz_path}: missing required key(s) {missing}; "275 f"{npz_path}: missing required key(s) {missing}; "
136 f"expected {list(REQUIRED_POINT_RECORD_KEYS)}"276 f"expected {list(required_keys)}"
137 )277 )
138 record = {key: np.asarray(data[key]) for key in REQUIRED_POINT_RECORD_KEYS}278 record = {key: np.asarray(data[key]) for key in schema if key in data.files}
139 present_optional = [key for key in OPTIONAL_POINT_RECORD_KEYS if key in data.files]
140 record.update({key: np.asarray(data[key]) for key in present_optional})
141279
142 points = record["points"]280 n_points = _validate_record_shapes(record, schema, source=str(npz_path))
143 if points.ndim != 2 or points.shape[1] != 3:281 for key, spec in schema.items():
144 raise ValueError(282 if key in record:
145 f"{npz_path}: 'points' must have shape (N, 3), got {points.shape}"283 record[key] = _coerce_storage_dtype(
146 )284 record[key], key=key, spec=spec, source=str(npz_path)
147 n_points = int(points.shape[0])
148 for key, arr in record.items():
149 if key == "points":
150 continue
151 if arr.shape != (n_points,):
152 raise ValueError(
153 f"{npz_path}: '{key}' must have shape (N,), got {arr.shape}"
154 )285 )
155 if NUMBER_OF_RETURNS_KEY in record:286 continue
156 record[NUMBER_OF_RETURNS_KEY] = _coerce_number_of_returns(287 if spec.fill is None: # pragma: no cover - required keys are checked above
157 record[NUMBER_OF_RETURNS_KEY], source=str(npz_path)288 raise ValueError(f"{npz_path}: missing required key(s) ['{key}']")
158 )
159 else:
160 logger.debug(289 logger.debug(
161 "%s: no '%s' member (pre-AI3D-382 record); filling %d zeros (unknown)",290 "%s: no '%s' member (record predates the key); synthesising %d values",
162 npz_path.name,291 npz_path.name,
163 NUMBER_OF_RETURNS_KEY,292 key,
164 n_points,293 n_points,
165 )294 )
166 record[NUMBER_OF_RETURNS_KEY] = np.zeros(n_points, dtype=NUMBER_OF_RETURNS_DTYPE)295 record[key] = spec.fill(n_points)
167 return {key: record[key] for key in POINT_RECORD_KEYS}296 return {key: record[key] for key in schema}
168297
169298
170def save_points_npz(path: str | Path, record: Mapping[str, np.ndarray]) -> Path:299def save_points_npz(path: str | Path, record: Mapping[str, np.ndarray]) -> Path:
171 """Write a point record as a compressed NPZ (inverse of :func:`load_points_npz`).300 """Write a point record as a compressed NPZ (inverse of :func:`load_points_npz`).
172301
173 Every key in :data:`POINT_RECORD_KEYS` must be present, ``number_of_returns``302 Every key in :data:`POINT_RECORD_KEYS` must be present, ``number_of_returns``
174 included: it is optional on load only, to read datasets that predate it.303 included: it is optional on load only, to read datasets that predate it.
175 A producer that omits it here is not "unknown", it is out of date, so this304 A producer that omits it here is not "unknown", it is out of date, so this
176 raises instead of zero-filling. ``number_of_returns`` is stored as uint8.305 raises instead of zero-filling. Members whose spec declares a storage dtype
306 (``number_of_returns``: uint8) are coerced to it before the write; the rest
307 are stored with the dtype the producer supplied.
308
309 Raises:
310 ValueError: A key is missing, an array has the wrong shape, or a member
311 with a declared storage dtype cannot be cast to it losslessly.
177 """312 """
178 npz_path = Path(path)313 npz_path = Path(path)
179 missing = [key for key in POINT_RECORD_KEYS if key not in record]314 schema = POINT_RECORD_SCHEMA
315 keys = _schema_keys(schema)
316 missing = [key for key in keys if key not in record]
180 if missing:317 if missing:
181 raise ValueError(318 raise ValueError(
182 f"Cannot save {npz_path}: missing required key(s) {missing}; "319 f"Cannot save {npz_path}: missing required key(s) {missing}; "
183 f"expected {list(POINT_RECORD_KEYS)}"320 f"expected {list(keys)}"
184 )321 )
185 payload = {key: np.asarray(record[key]) for key in POINT_RECORD_KEYS}322 source = f"Cannot save {npz_path}"
186 payload[NUMBER_OF_RETURNS_KEY] = _coerce_number_of_returns(323 payload = {
187 payload[NUMBER_OF_RETURNS_KEY], source=f"Cannot save {npz_path}"324 key: _coerce_storage_dtype(record[key], key=key, spec=schema[key], source=source)
188 )325 for key in keys
189 points = payload["points"]326 }
190 if points.ndim != 2 or points.shape[1] != 3:327 _validate_record_shapes(payload, schema, source=source)
328 npz_path.parent.mkdir(parents=True, exist_ok=True)
329 np.savez_compressed(npz_path, **payload)
330 return npz_path
331
332
333def mask_record(record: Mapping[str, np.ndarray], mask: np.ndarray) -> PointRecord:
334 """Select the same points from every member of a point record.
335
336 Schema-agnostic: whatever keys the record carries are all indexed with
337 *mask* along axis 0, so ``points`` keeps its ``(n, 3)`` rows while the
338 ancillary members stay 1-D. Adding a field to the contract needs no change
339 here.
340
341 Args:
342 record: Point record whose members all share a leading dimension.
343 mask: Boolean mask of shape ``(N,)`` or an integer index array; any
344 NumPy row indexer works.
345
346 Returns:
347 A new record with the same keys, each member indexed by *mask*.
348
349 Raises:
350 ValueError: The record is empty or its members disagree on point count.
351 """
352 if not record:
353 raise ValueError("Cannot mask an empty point record")
354 arrays = {key: np.asarray(value) for key, value in record.items()}
355 counts = {key: array.shape[0] if array.ndim else -1 for key, array in arrays.items()}
356 if len(set(counts.values())) > 1:
191 raise ValueError(357 raise ValueError(
192 f"Cannot save {npz_path}: 'points' must have shape (N, 3), got {points.shape}"358 f"Cannot mask a point record whose members disagree on point count: {counts}"
193 )359 )
194 n_points = int(points.shape[0])360 return {key: array[mask] for key, array in arrays.items()}
195 for key in POINT_RECORD_KEYS:361
196 if key == "points":362
197 continue363def concat_records(records: Sequence[Mapping[str, np.ndarray]]) -> PointRecord:
198 arr = payload[key]364 """Concatenate point records member-by-member along axis 0.
199 if arr.shape != (n_points,):365
366 Schema-agnostic: every key of the first record is concatenated across all
367 records, so ``points`` grows by rows and the ancillary members by elements.
368 Dtypes are not checked here -- :func:`concat_points_npz` owns that policy.
369
370 Args:
371 records: Records to join, in the intended order. All must carry the
372 same key set.
373
374 Returns:
375 One record holding the concatenated members, keyed in the first
376 record's order.
377
378 Raises:
379 ValueError: *records* is empty or the records' key sets differ.
380 """
381 if not records:
382 raise ValueError("Cannot concatenate an empty sequence of point records")
383 keys = list(records[0])
384 for index, other in enumerate(records[1:], start=1):
385 if set(other) != set(keys):
200 raise ValueError(386 raise ValueError(
201 f"Cannot save {npz_path}: '{key}' must have shape (N,), got {arr.shape}"387 f"Cannot concatenate point records with different keys: record 0 has "
388 f"{sorted(keys)}, record {index} has {sorted(other)}"
202 )389 )
203 npz_path.parent.mkdir(parents=True, exist_ok=True)390 return {
204 np.savez_compressed(npz_path, **payload)391 key: np.concatenate([np.asarray(record[key]) for record in records], axis=0)
205 return npz_path392 for key in keys
393 }
206394
207395
208def load_segment_points(396def load_segment_points(
209 files: Sequence[str | Path],397 files: Sequence[str | Path],
Importance #3: src/iolabs/common/segment_points_io.py @@ -5,8 +5,17 @@
5segment's run3 records, chunked streaming of the ``points`` member,5segment's run3 records, chunked streaming of the ``points`` member,
6``run3_geoshift.json`` lookup/loading, and per-segment fnmatch blacklist6``run3_geoshift.json`` lookup/loading, and per-segment fnmatch blacklist
7filtering. Pure NumPy only.7filtering. Pure NumPy only.
88
9The record schema itself lives in one place: :data:`POINT_RECORD_SCHEMA`, an
10ordered ``key -> PointFieldSpec`` registry that declares whether a key is
11required on load, the storage dtype it is coerced to and validated against on
12save *and* load, its column count, and the factory that synthesises it for a
13record written before the key existed. Every public key tuple is derived from
14that registry and every validation/coercion path is driven by it, so adding a
15field to the contract is a single registry entry plus its producers -- no
16per-key branches in the load/save/merge code.
17
9``number_of_returns`` (AI3D-382) is the LAS per-point return count, stored as18``number_of_returns`` (AI3D-382) is the LAS per-point return count, stored as
10uint8. It is **always written** by :func:`save_points_npz` -- a record that19uint8. It is **always written** by :func:`save_points_npz` -- a record that
11does not carry it is rejected rather than zero-filled, so a producer that20does not carry it is rejected rather than zero-filled, so a producer that
12still has the data cannot silently downgrade it. On the read side the key is21still has the data cannot silently downgrade it. On the read side the key is
Importance #4: src/iolabs/common/segment_points_io.py @@ -217,23 +405,20 @@
217 npz_files = [Path(path) for path in files]405 npz_files = [Path(path) for path in files]
218 if not npz_files:406 if not npz_files:
219 raise FileNotFoundError("No *_points.npz files provided for segment load")407 raise FileNotFoundError("No *_points.npz files provided for segment load")
220408
221 buckets: dict[str, list[np.ndarray]] = {key: [] for key in POINT_RECORD_KEYS}409 records: list[PointRecord] = []
222 point_file_ids_parts: list[np.ndarray] = []410 point_file_ids_parts: list[np.ndarray] = []
223 file_stems: list[str] = []411 file_stems: list[str] = []
224412
225 for file_idx, npz_path in enumerate(npz_files):413 for file_idx, npz_path in enumerate(npz_files):
226 record = load_points_npz(npz_path)414 record = load_points_npz(npz_path)
227 point_count = int(record["points"].shape[0])415 point_count = int(record[POINTS_KEY].shape[0])
228 for key in POINT_RECORD_KEYS:416 records.append(record)
229 buckets[key].append(record[key])
230 point_file_ids_parts.append(np.full(point_count, file_idx, dtype=np.int32))417 point_file_ids_parts.append(np.full(point_count, file_idx, dtype=np.int32))
231 file_stems.append(npz_path.stem)418 file_stems.append(npz_path.stem)
232419
233 merged: PointRecord = {420 merged = concat_records(records)
234 key: np.concatenate(parts, axis=0) for key, parts in buckets.items()
235 }
236 point_file_ids = (421 point_file_ids = (
237 np.concatenate(point_file_ids_parts, axis=0)422 np.concatenate(point_file_ids_parts, axis=0)
238 if point_file_ids_parts423 if point_file_ids_parts
239 else np.zeros((0,), dtype=np.int32)424 else np.zeros((0,), dtype=np.int32)
Importance #5: src/iolabs/common/segment_points_io.py @@ -679,9 +864,10 @@
679864
680 casts: dict[str, np.dtype] = {865 casts: dict[str, np.dtype] = {
681 key: np.dtype(value) for key, value in (target_dtypes or {}).items()866 key: np.dtype(value) for key, value in (target_dtypes or {}).items()
682 }867 }
683 buckets: dict[str, list[np.ndarray]] = {key: [] for key in POINT_RECORD_KEYS}868 keys = _schema_keys(POINT_RECORD_SCHEMA)
869 records: list[PointRecord] = []
684 spans: list[RecordSpan] = []870 spans: list[RecordSpan] = []
685 dtypes: dict[str, np.dtype] = {}871 dtypes: dict[str, np.dtype] = {}
686 offset = 0872 offset = 0
687873
Importance #6: src/iolabs/common/segment_points_io.py @@ -690,27 +876,25 @@
690 for key, dtype in casts.items():876 for key, dtype in casts.items():
691 if key in record:877 if key in record:
692 record[key] = np.asarray(record[key], dtype=dtype)878 record[key] = np.asarray(record[key], dtype=dtype)
693 if not dtypes:879 if not dtypes:
694 dtypes = {key: record[key].dtype for key in POINT_RECORD_KEYS}880 dtypes = {key: record[key].dtype for key in keys}
695 else:881 else:
696 for key in POINT_RECORD_KEYS:882 for key in keys:
697 if key in casts:883 if key in casts:
698 continue884 continue
699 if record[key].dtype != dtypes[key]:885 if record[key].dtype != dtypes[key]:
700 raise ValueError(886 raise ValueError(
701 f"{npz_path}: '{key}' dtype {record[key].dtype} does not match "887 f"{npz_path}: '{key}' dtype {record[key].dtype} does not match "
702 f"{dtypes[key]} from {npz_files[0].name}; refusing to concatenate "888 f"{dtypes[key]} from {npz_files[0].name}; refusing to concatenate "
703 "records with mixed dtypes (pass target_dtypes to opt into casting)"889 "records with mixed dtypes (pass target_dtypes to opt into casting)"
704 )890 )
705 count = int(record["points"].shape[0])891 count = int(record[POINTS_KEY].shape[0])
706 for key in POINT_RECORD_KEYS:892 records.append(record)
707 buckets[key].append(record[key])
708 spans.append(RecordSpan(name=npz_path.name, offset=offset, count=count))893 spans.append(RecordSpan(name=npz_path.name, offset=offset, count=count))
709 offset += count894 offset += count
710895
711 merged: PointRecord = {key: np.concatenate(parts, axis=0) for key, parts in buckets.items()}896 return concat_records(records), spans
712 return merged, spans
713897
714898
715def load_run3_segment(899def load_run3_segment(
716 segment_dir: str | Path,900 segment_dir: str | Path,
Importance #7: tests/test_segment_points_io.py @@ -1,20 +1,27 @@
1"""Tests for the per-segment point-NPZ contract helpers."""1"""Tests for the per-segment point-NPZ contract helpers."""
22
3import functools
3import json4import json
4import logging5import logging
5from pathlib import Path6from pathlib import Path
7from types import MappingProxyType
68
7import numpy as np9import numpy as np
8import pytest10import pytest
911
12from iolabs.common import segment_points_io
10from iolabs.common.segment_points_io import (13from iolabs.common.segment_points_io import (
11 GEOSHIFT_NAME,14 GEOSHIFT_NAME,
12 NUMBER_OF_RETURNS_KEY,15 NUMBER_OF_RETURNS_KEY,
16 OPTIONAL_POINT_RECORD_KEYS,
13 POINT_RECORD_KEYS,17 POINT_RECORD_KEYS,
18 POINT_RECORD_SCHEMA,
14 REQUIRED_POINT_RECORD_KEYS,19 REQUIRED_POINT_RECORD_KEYS,
20 PointFieldSpec,
15 RecordSpan,21 RecordSpan,
16 concat_points_npz,22 concat_points_npz,
23 concat_records,
17 discover_run3_files,24 discover_run3_files,
18 filter_segment_files,25 filter_segment_files,
19 find_geoshift,26 find_geoshift,
20 find_geoshift_or_none,27 find_geoshift_or_none,
Importance #8: tests/test_segment_points_io.py @@ -24,8 +31,9 @@
24 load_geoshift,31 load_geoshift,
25 load_points_npz,32 load_points_npz,
26 load_run3_segment,33 load_run3_segment,
27 load_segment_points,34 load_segment_points,
35 mask_record,
28 normalize_segment_file_blacklist,36 normalize_segment_file_blacklist,
29 parse_segment_key,37 parse_segment_key,
30 read_points_header,38 read_points_header,
31 save_points_npz,39 save_points_npz,
Importance #9: tests/test_segment_points_io.py @@ -724,4 +732,201 @@
724 shift = find_geoshift(seg_dir)732 shift = find_geoshift(seg_dir)
725733
726 np.testing.assert_array_equal(shift, np.array([1.0, 1.0, 1.0]))734 np.testing.assert_array_equal(shift, np.array([1.0, 1.0, 1.0]))
727 assert any("Multiple" in message for message in caplog.messages)735 assert any("Multiple" in message for message in caplog.messages)
736
737
738def test_public_key_tuples_are_derived_from_the_schema() -> None:
739 assert POINT_RECORD_KEYS == tuple(POINT_RECORD_SCHEMA)
740 assert REQUIRED_POINT_RECORD_KEYS == (
741 "points",
742 "red",
743 "green",
744 "blue",
745 "intensity",
746 "scan_angle",
747 )
748 assert OPTIONAL_POINT_RECORD_KEYS == (NUMBER_OF_RETURNS_KEY,)
749 assert all(POINT_RECORD_SCHEMA[key].required for key in REQUIRED_POINT_RECORD_KEYS)
750 assert POINT_RECORD_SCHEMA[NUMBER_OF_RETURNS_KEY].storage_dtype == np.dtype(np.uint8)
751 assert POINT_RECORD_SCHEMA["points"].columns == 3
752
753
754def test_point_field_spec_rejects_an_optional_key_without_a_fill() -> None:
755 """An optional key with no fill would make pre-existing records unloadable."""
756 with pytest.raises(ValueError, match="needs a fill factory"):
757 PointFieldSpec(required=False)
758
759
760def test_mask_record_masks_points_rows_and_ancillary_elements() -> None:
761 record = _make_record(5, seed=80)
762 mask = np.array([True, False, True, False, True])
763
764 masked = mask_record(record, mask)
765
766 assert list(masked) == list(record)
767 assert masked["points"].shape == (3, 3)
768 np.testing.assert_array_equal(masked["points"], record["points"][mask])
769 for key in POINT_RECORD_KEYS:
770 np.testing.assert_array_equal(masked[key], record[key][mask])
771
772
773def test_mask_record_accepts_an_integer_index_array() -> None:
774 record = _make_record(4, seed=81)
775 index = np.array([3, 0])
776
777 masked = mask_record(record, index)
778
779 np.testing.assert_array_equal(masked["points"], record["points"][index])
780 np.testing.assert_array_equal(masked["intensity"], record["intensity"][index])
781
782
783def test_mask_record_rejects_misaligned_members() -> None:
784 record = _make_record(4, seed=82)
785 record["red"] = record["red"][:2]
786 with pytest.raises(ValueError, match="disagree on point count"):
787 mask_record(record, np.array([True, False, True, False]))
788
789
790def test_mask_record_rejects_an_empty_record() -> None:
791 with pytest.raises(ValueError, match="empty point record"):
792 mask_record({}, np.array([True]))
793
794
795def test_concat_records_joins_every_member_on_axis_zero() -> None:
796 first = _make_record(3, seed=83)
797 second = _make_record(2, seed=84)
798
799 merged = concat_records([first, second])
800
801 assert list(merged) == list(POINT_RECORD_KEYS)
802 assert merged["points"].shape == (5, 3)
803 for key in POINT_RECORD_KEYS:
804 np.testing.assert_array_equal(
805 merged[key], np.concatenate([first[key], second[key]], axis=0)
806 )
807
808
809def test_concat_records_rejects_mismatched_key_sets() -> None:
810 with pytest.raises(ValueError, match="different keys"):
811 concat_records([_make_record(2, seed=85), _make_legacy_record(2, seed=86)])
812
813
814def test_concat_records_rejects_an_empty_sequence() -> None:
815 with pytest.raises(ValueError, match="empty sequence"):
816 concat_records([])
817
818
819# --- Adding a field to the contract must be a registry entry and nothing else ---
820
821EXTRA_KEY = "point_source_id"
822
823EXTRA_SPEC = PointFieldSpec(
824 required=False,
825 storage_dtype=np.dtype(np.int16),
826 fill=functools.partial(np.full, fill_value=-1, dtype=np.int16),
827 noun="source id",
828)
829
830
831@pytest.fixture
832def extended_schema(monkeypatch: pytest.MonkeyPatch) -> tuple[str, ...]:
833 """Register one extra optional key, exactly as a real schema addition would."""
834 monkeypatch.setattr(
835 segment_points_io,
836 "POINT_RECORD_SCHEMA",
837 MappingProxyType({**POINT_RECORD_SCHEMA, EXTRA_KEY: EXTRA_SPEC}),
838 )
839 return (*POINT_RECORD_KEYS, EXTRA_KEY)
840
841
842def _make_extended_record(n: int, *, seed: int) -> dict[str, np.ndarray]:
843 record = _make_record(n, seed=seed)
844 record[EXTRA_KEY] = np.arange(n, dtype=np.int16)
845 return record
846
847
848def test_registry_entry_alone_round_trips_a_new_field(
849 tmp_path: Path, extended_schema: tuple[str, ...]
850) -> None:
851 record = _make_extended_record(4, seed=90)
852
853 path = save_points_npz(tmp_path / "extended_run3_points.npz", record)
854
855 with np.load(path) as data:
856 assert EXTRA_KEY in data.files
857 loaded = load_points_npz(path)
858 assert list(loaded) == list(extended_schema)
859 assert loaded[EXTRA_KEY].dtype == np.int16
860 np.testing.assert_array_equal(loaded[EXTRA_KEY], record[EXTRA_KEY])
861
862
863def test_registry_entry_alone_coerces_and_validates_a_new_field(
864 tmp_path: Path, extended_schema: tuple[str, ...]
865) -> None:
866 record = _make_extended_record(3, seed=91)
867 record[EXTRA_KEY] = record[EXTRA_KEY].astype(np.int64)
868 loaded = load_points_npz(save_points_npz(tmp_path / "cast_run3_points.npz", record))
869 assert loaded[EXTRA_KEY].dtype == np.int16
870
871 out_of_range = _make_extended_record(3, seed=92)
872 out_of_range[EXTRA_KEY] = np.array([1, 40_000, 3], dtype=np.int32)
873 with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*fit in int16"):
874 save_points_npz(tmp_path / "range_run3_points.npz", out_of_range)
875
876 boolean = _make_extended_record(3, seed=93)
877 boolean[EXTRA_KEY] = np.array([True, False, True])
878 with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*bool"):
879 save_points_npz(tmp_path / "bool_run3_points.npz", boolean)
880
881 misshaped = _make_extended_record(3, seed=94)
882 misshaped[EXTRA_KEY] = misshaped[EXTRA_KEY].reshape(3, 1)
883 with pytest.raises(ValueError, match=rf"'{EXTRA_KEY}' must have shape \(N,\)"):
884 save_points_npz(tmp_path / "shape_run3_points.npz", misshaped)
885
886
887def test_registry_entry_alone_fills_and_requires_a_new_field(
888 tmp_path: Path, extended_schema: tuple[str, ...]
889) -> None:
890 older = _make_record(5, seed=95)
891 path = tmp_path / "older_run3_points.npz"
892 np.savez_compressed(path, **older)
893
894 loaded = load_points_npz(path)
895
896 assert list(loaded) == list(extended_schema)
897 np.testing.assert_array_equal(loaded[EXTRA_KEY], np.full(5, -1, dtype=np.int16))
898 with pytest.raises(ValueError, match=EXTRA_KEY):
899 save_points_npz(tmp_path / "incomplete_run3_points.npz", older)
900
901
902def test_registry_entry_alone_flows_through_merge_mask_and_concat(
903 tmp_path: Path, extended_schema: tuple[str, ...]
904) -> None:
905 seg_dir = tmp_path / "segment_042"
906 seg_dir.mkdir()
907 first = _make_extended_record(3, seed=96)
908 second = _make_extended_record(2, seed=97)
909 save_points_npz(seg_dir / "a_run3_points.npz", first)
910 save_points_npz(seg_dir / "b_run3_points.npz", second)
911
912 merged, spans = load_run3_segment(seg_dir)
913 segment_merged, point_file_ids, _ = load_segment_points(
914 [seg_dir / "a_run3_points.npz", seg_dir / "b_run3_points.npz"]
915 )
916 masked = mask_record(merged, np.array([True, False, True, False, True]))
917 joined = concat_records([first, second])
918
919 expected = np.concatenate([first[EXTRA_KEY], second[EXTRA_KEY]], axis=0)
920 assert [span.count for span in spans] == [3, 2]
921 assert point_file_ids.shape == (5,)
922 np.testing.assert_array_equal(merged[EXTRA_KEY], expected)
923 np.testing.assert_array_equal(segment_merged[EXTRA_KEY], expected)
924 np.testing.assert_array_equal(masked[EXTRA_KEY], expected[[0, 2, 4]])
925 np.testing.assert_array_equal(joined[EXTRA_KEY], expected)
926
927
928def test_extended_schema_does_not_leak_into_the_real_contract() -> None:
929 """The monkeypatched registry above must not pollute the shipped schema."""
930 assert EXTRA_KEY not in POINT_RECORD_SCHEMA
931 assert EXTRA_KEY not in segment_points_io.POINT_RECORD_SCHEMA
932 assert tuple(segment_points_io.POINT_RECORD_SCHEMA) == POINT_RECORD_KEYS
Importance #10: tests/test_segment_points_io.py @@ -1,20 +1,27 @@
1"""Tests for the per-segment point-NPZ contract helpers."""1"""Tests for the per-segment point-NPZ contract helpers."""
22
3import functools
3import json4import json
4import logging5import logging
5from pathlib import Path6from pathlib import Path
7from types import MappingProxyType
68
7import numpy as np9import numpy as np
8import pytest10import pytest
911
12from iolabs.common import segment_points_io
10from iolabs.common.segment_points_io import (13from iolabs.common.segment_points_io import (
11 GEOSHIFT_NAME,14 GEOSHIFT_NAME,
12 NUMBER_OF_RETURNS_KEY,15 NUMBER_OF_RETURNS_KEY,
16 OPTIONAL_POINT_RECORD_KEYS,
13 POINT_RECORD_KEYS,17 POINT_RECORD_KEYS,
18 POINT_RECORD_SCHEMA,
14 REQUIRED_POINT_RECORD_KEYS,19 REQUIRED_POINT_RECORD_KEYS,
20 PointFieldSpec,
15 RecordSpan,21 RecordSpan,
16 concat_points_npz,22 concat_points_npz,
23 concat_records,
17 discover_run3_files,24 discover_run3_files,
18 filter_segment_files,25 filter_segment_files,
19 find_geoshift,26 find_geoshift,
20 find_geoshift_or_none,27 find_geoshift_or_none,
Importance #11: tests/test_segment_points_io.py @@ -24,8 +31,9 @@
24 load_geoshift,31 load_geoshift,
25 load_points_npz,32 load_points_npz,
26 load_run3_segment,33 load_run3_segment,
27 load_segment_points,34 load_segment_points,
35 mask_record,
28 normalize_segment_file_blacklist,36 normalize_segment_file_blacklist,
29 parse_segment_key,37 parse_segment_key,
30 read_points_header,38 read_points_header,
31 save_points_npz,39 save_points_npz,
Importance #12: tests/test_segment_points_io.py @@ -724,4 +732,201 @@
724 shift = find_geoshift(seg_dir)732 shift = find_geoshift(seg_dir)
725733
726 np.testing.assert_array_equal(shift, np.array([1.0, 1.0, 1.0]))734 np.testing.assert_array_equal(shift, np.array([1.0, 1.0, 1.0]))
727 assert any("Multiple" in message for message in caplog.messages)735 assert any("Multiple" in message for message in caplog.messages)
736
737
738def test_public_key_tuples_are_derived_from_the_schema() -> None:
739 assert POINT_RECORD_KEYS == tuple(POINT_RECORD_SCHEMA)
740 assert REQUIRED_POINT_RECORD_KEYS == (
741 "points",
742 "red",
743 "green",
744 "blue",
745 "intensity",
746 "scan_angle",
747 )
748 assert OPTIONAL_POINT_RECORD_KEYS == (NUMBER_OF_RETURNS_KEY,)
749 assert all(POINT_RECORD_SCHEMA[key].required for key in REQUIRED_POINT_RECORD_KEYS)
750 assert POINT_RECORD_SCHEMA[NUMBER_OF_RETURNS_KEY].storage_dtype == np.dtype(np.uint8)
751 assert POINT_RECORD_SCHEMA["points"].columns == 3
752
753
754def test_point_field_spec_rejects_an_optional_key_without_a_fill() -> None:
755 """An optional key with no fill would make pre-existing records unloadable."""
756 with pytest.raises(ValueError, match="needs a fill factory"):
757 PointFieldSpec(required=False)
758
759
760def test_mask_record_masks_points_rows_and_ancillary_elements() -> None:
761 record = _make_record(5, seed=80)
762 mask = np.array([True, False, True, False, True])
763
764 masked = mask_record(record, mask)
765
766 assert list(masked) == list(record)
767 assert masked["points"].shape == (3, 3)
768 np.testing.assert_array_equal(masked["points"], record["points"][mask])
769 for key in POINT_RECORD_KEYS:
770 np.testing.assert_array_equal(masked[key], record[key][mask])
771
772
773def test_mask_record_accepts_an_integer_index_array() -> None:
774 record = _make_record(4, seed=81)
775 index = np.array([3, 0])
776
777 masked = mask_record(record, index)
778
779 np.testing.assert_array_equal(masked["points"], record["points"][index])
780 np.testing.assert_array_equal(masked["intensity"], record["intensity"][index])
781
782
783def test_mask_record_rejects_misaligned_members() -> None:
784 record = _make_record(4, seed=82)
785 record["red"] = record["red"][:2]
786 with pytest.raises(ValueError, match="disagree on point count"):
787 mask_record(record, np.array([True, False, True, False]))
788
789
790def test_mask_record_rejects_an_empty_record() -> None:
791 with pytest.raises(ValueError, match="empty point record"):
792 mask_record({}, np.array([True]))
793
794
795def test_concat_records_joins_every_member_on_axis_zero() -> None:
796 first = _make_record(3, seed=83)
797 second = _make_record(2, seed=84)
798
799 merged = concat_records([first, second])
800
801 assert list(merged) == list(POINT_RECORD_KEYS)
802 assert merged["points"].shape == (5, 3)
803 for key in POINT_RECORD_KEYS:
804 np.testing.assert_array_equal(
805 merged[key], np.concatenate([first[key], second[key]], axis=0)
806 )
807
808
809def test_concat_records_rejects_mismatched_key_sets() -> None:
810 with pytest.raises(ValueError, match="different keys"):
811 concat_records([_make_record(2, seed=85), _make_legacy_record(2, seed=86)])
812
813
814def test_concat_records_rejects_an_empty_sequence() -> None:
815 with pytest.raises(ValueError, match="empty sequence"):
816 concat_records([])
817
818
819# --- Adding a field to the contract must be a registry entry and nothing else ---
820
821EXTRA_KEY = "point_source_id"
822
823EXTRA_SPEC = PointFieldSpec(
824 required=False,
825 storage_dtype=np.dtype(np.int16),
826 fill=functools.partial(np.full, fill_value=-1, dtype=np.int16),
827 noun="source id",
828)
829
830
831@pytest.fixture
832def extended_schema(monkeypatch: pytest.MonkeyPatch) -> tuple[str, ...]:
833 """Register one extra optional key, exactly as a real schema addition would."""
834 monkeypatch.setattr(
835 segment_points_io,
836 "POINT_RECORD_SCHEMA",
837 MappingProxyType({**POINT_RECORD_SCHEMA, EXTRA_KEY: EXTRA_SPEC}),
838 )
839 return (*POINT_RECORD_KEYS, EXTRA_KEY)
840
841
842def _make_extended_record(n: int, *, seed: int) -> dict[str, np.ndarray]:
843 record = _make_record(n, seed=seed)
844 record[EXTRA_KEY] = np.arange(n, dtype=np.int16)
845 return record
846
847
848def test_registry_entry_alone_round_trips_a_new_field(
849 tmp_path: Path, extended_schema: tuple[str, ...]
850) -> None:
851 record = _make_extended_record(4, seed=90)
852
853 path = save_points_npz(tmp_path / "extended_run3_points.npz", record)
854
855 with np.load(path) as data:
856 assert EXTRA_KEY in data.files
857 loaded = load_points_npz(path)
858 assert list(loaded) == list(extended_schema)
859 assert loaded[EXTRA_KEY].dtype == np.int16
860 np.testing.assert_array_equal(loaded[EXTRA_KEY], record[EXTRA_KEY])
861
862
863def test_registry_entry_alone_coerces_and_validates_a_new_field(
864 tmp_path: Path, extended_schema: tuple[str, ...]
865) -> None:
866 record = _make_extended_record(3, seed=91)
867 record[EXTRA_KEY] = record[EXTRA_KEY].astype(np.int64)
868 loaded = load_points_npz(save_points_npz(tmp_path / "cast_run3_points.npz", record))
869 assert loaded[EXTRA_KEY].dtype == np.int16
870
871 out_of_range = _make_extended_record(3, seed=92)
872 out_of_range[EXTRA_KEY] = np.array([1, 40_000, 3], dtype=np.int32)
873 with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*fit in int16"):
874 save_points_npz(tmp_path / "range_run3_points.npz", out_of_range)
875
876 boolean = _make_extended_record(3, seed=93)
877 boolean[EXTRA_KEY] = np.array([True, False, True])
878 with pytest.raises(ValueError, match=f"{EXTRA_KEY}.*bool"):
879 save_points_npz(tmp_path / "bool_run3_points.npz", boolean)
880
881 misshaped = _make_extended_record(3, seed=94)
882 misshaped[EXTRA_KEY] = misshaped[EXTRA_KEY].reshape(3, 1)
883 with pytest.raises(ValueError, match=rf"'{EXTRA_KEY}' must have shape \(N,\)"):
884 save_points_npz(tmp_path / "shape_run3_points.npz", misshaped)
885
886
887def test_registry_entry_alone_fills_and_requires_a_new_field(
888 tmp_path: Path, extended_schema: tuple[str, ...]
889) -> None:
890 older = _make_record(5, seed=95)
891 path = tmp_path / "older_run3_points.npz"
892 np.savez_compressed(path, **older)
893
894 loaded = load_points_npz(path)
895
896 assert list(loaded) == list(extended_schema)
897 np.testing.assert_array_equal(loaded[EXTRA_KEY], np.full(5, -1, dtype=np.int16))
898 with pytest.raises(ValueError, match=EXTRA_KEY):
899 save_points_npz(tmp_path / "incomplete_run3_points.npz", older)
900
901
902def test_registry_entry_alone_flows_through_merge_mask_and_concat(
903 tmp_path: Path, extended_schema: tuple[str, ...]
904) -> None:
905 seg_dir = tmp_path / "segment_042"
906 seg_dir.mkdir()
907 first = _make_extended_record(3, seed=96)
908 second = _make_extended_record(2, seed=97)
909 save_points_npz(seg_dir / "a_run3_points.npz", first)
910 save_points_npz(seg_dir / "b_run3_points.npz", second)
911
912 merged, spans = load_run3_segment(seg_dir)
913 segment_merged, point_file_ids, _ = load_segment_points(
914 [seg_dir / "a_run3_points.npz", seg_dir / "b_run3_points.npz"]
915 )
916 masked = mask_record(merged, np.array([True, False, True, False, True]))
917 joined = concat_records([first, second])
918
919 expected = np.concatenate([first[EXTRA_KEY], second[EXTRA_KEY]], axis=0)
920 assert [span.count for span in spans] == [3, 2]
921 assert point_file_ids.shape == (5,)
922 np.testing.assert_array_equal(merged[EXTRA_KEY], expected)
923 np.testing.assert_array_equal(segment_merged[EXTRA_KEY], expected)
924 np.testing.assert_array_equal(masked[EXTRA_KEY], expected[[0, 2, 4]])
925 np.testing.assert_array_equal(joined[EXTRA_KEY], expected)
926
927
928def test_extended_schema_does_not_leak_into_the_real_contract() -> None:
929 """The monkeypatched registry above must not pollute the shipped schema."""
930 assert EXTRA_KEY not in POINT_RECORD_SCHEMA
931 assert EXTRA_KEY not in segment_points_io.POINT_RECORD_SCHEMA
932 assert tuple(segment_points_io.POINT_RECORD_SCHEMA) == POINT_RECORD_KEYS