Coverage for oc_meta / run / infodir / _common.py: 74%
248 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-07-25 10:39 +0000
1# SPDX-FileCopyrightText: 2026 Arcangelo Massari <arcangelo.massari@unibo.it>
2#
3# SPDX-License-Identifier: ISC
5from __future__ import annotations
7import multiprocessing
8import os
9import struct
10import zipfile
11from collections.abc import Callable, Iterable, Iterator
12from concurrent.futures import FIRST_COMPLETED, Future, ProcessPoolExecutor, wait
13from dataclasses import dataclass
14from itertools import islice
15from typing import BinaryIO, TypeVar
17import orjson
18from oc_ocdm.support import get_prefix, get_resource_number, get_short_name
20from oc_meta.lib.console import advance_progress, create_progress
21from oc_meta.lib.file_manager import collect_zip_files
23ENTITY_TYPES = frozenset({"ar", "br", "ra", "re", "id"})
24COUNTER_SIZE = 4
25COUNTERS_PER_CHUNK = 100_000
27CounterKey = tuple[str, str]
28ResultType = TypeVar("ResultType")
31class SourceDataError(ValueError):
32 pass
35@dataclass(frozen=True)
36class DataEntity:
37 uri: str
38 resource_number: int
41@dataclass(frozen=True)
42class DataZipResult:
43 path: str
44 entities: dict[CounterKey, list[DataEntity]]
47@dataclass(frozen=True)
48class ProvenanceZipResult:
49 path: str
50 counters: dict[CounterKey, dict[int, int]]
53@dataclass(frozen=True)
54class DataScanResult:
55 zip_files: int
56 entities: int
57 maxima: dict[CounterKey, int]
58 missing_provenance: int
59 missing_examples: list[dict[str, object]]
62@dataclass(frozen=True)
63class ProvenanceScanResult:
64 zip_files: int
65 entities: int
68class SparseCounterStore:
69 def __init__(self, root: str) -> None:
70 self.root = root
71 self.maxima: dict[CounterKey, int] = {}
72 self._files: dict[CounterKey, BinaryIO] = {}
74 def keys(self) -> set[CounterKey]:
75 return set(self.maxima)
77 def maximum(self, key: CounterKey) -> int:
78 return self.maxima[key] if key in self.maxima else 0
80 def write_updates(self, key: CounterKey, updates: dict[int, int]) -> None:
81 if not updates:
82 return
83 ordered = sorted(updates.items())
84 run: list[tuple[int, int]] = []
85 previous_resource = 0
86 for resource_number, counter in ordered:
87 if counter > 2**32 - 1:
88 raise SourceDataError(
89 f"Snapshot counter {counter} exceeds the supported range"
90 )
91 if run and resource_number != previous_resource + 1:
92 self._write_run(key, run)
93 run = []
94 run.append((resource_number, counter))
95 previous_resource = resource_number
96 self._write_run(key, run)
97 self.maxima[key] = max(self.maximum(key), ordered[-1][0])
99 def read_span(self, key: CounterKey, start: int, end: int) -> list[int]:
100 if end < start:
101 return []
102 size = end - start + 1
103 file_handle = self._files[key] if key in self._files else None
104 if file_handle is None:
105 return [0] * size
106 file_handle.seek((start - 1) * COUNTER_SIZE)
107 data = file_handle.read(size * COUNTER_SIZE)
108 if len(data) < size * COUNTER_SIZE:
109 data += b"\0" * (size * COUNTER_SIZE - len(data))
110 return list(struct.unpack(f"<{size}I", data))
112 def iter_chunks(
113 self, key: CounterKey, start: int = 1, end: int | None = None
114 ) -> Iterator[tuple[int, list[int]]]:
115 maximum = self.maximum(key) if end is None else end
116 current = start
117 while current <= maximum:
118 chunk_end = min(current + COUNTERS_PER_CHUNK - 1, maximum)
119 yield current, self.read_span(key, current, chunk_end)
120 current = chunk_end + 1
122 def render(self, key: CounterKey, output_path: str) -> None:
123 os.makedirs(os.path.dirname(output_path), exist_ok=True)
124 rendered: dict[int, bytes] = {0: b"\n"}
125 with open(output_path, "wb") as output_file:
126 for _, counters in self.iter_chunks(key):
127 lines: list[bytes] = []
128 for counter in counters:
129 if counter not in rendered:
130 rendered[counter] = f"{counter}\n".encode()
131 lines.append(rendered[counter])
132 output_file.writelines(lines)
134 def close(self) -> None:
135 for file_handle in self._files.values():
136 file_handle.close()
137 self._files.clear()
139 def _write_run(self, key: CounterKey, run: list[tuple[int, int]]) -> None:
140 start = run[0][0]
141 current = self.read_span(key, start, run[-1][0])
142 merged = [
143 max(persisted, update[1])
144 for persisted, update in zip(current, run, strict=True)
145 ]
146 file_handle = self._file(key)
147 file_handle.seek((start - 1) * COUNTER_SIZE)
148 file_handle.write(struct.pack(f"<{len(merged)}I", *merged))
150 def _file(self, key: CounterKey) -> BinaryIO:
151 if key not in self._files:
152 prefix, short_name = key
153 directory = os.path.join(self.root, prefix)
154 os.makedirs(directory, exist_ok=True)
155 self._files[key] = open(os.path.join(directory, f"{short_name}.bin"), "w+b")
156 return self._files[key]
159def bounded_process_map(
160 paths: Iterable[str],
161 worker: Callable[[str], ResultType],
162 workers: int,
163) -> Iterator[ResultType]:
164 if workers <= 0:
165 raise ValueError("workers must be greater than zero")
166 path_iterator = iter(paths)
167 context = multiprocessing.get_context("forkserver")
168 pending: dict[Future[ResultType], str] = {}
169 with ProcessPoolExecutor(max_workers=workers, mp_context=context) as executor:
170 for path in islice(path_iterator, workers * 2):
171 pending[executor.submit(worker, path)] = path
172 while pending:
173 completed, _ = wait(pending, return_when=FIRST_COMPLETED)
174 for future in completed:
175 pending.pop(future)
176 yield future.result()
177 try:
178 path = next(path_iterator)
179 except StopIteration:
180 continue
181 pending[executor.submit(worker, path)] = path
184def process_data_zip(path: str) -> DataZipResult:
185 entities: dict[CounterKey, list[DataEntity]] = {}
186 for entity in _graph_entities(path):
187 uri = _entity_uri(entity, path)
188 key, resource_number = _counter_identity(uri, path)
189 if key not in entities:
190 entities[key] = []
191 entities[key].append(DataEntity(uri=uri, resource_number=resource_number))
192 for values in entities.values():
193 values.sort(key=lambda entity: entity.resource_number)
194 return DataZipResult(path=path, entities=entities)
197def process_provenance_zip(path: str) -> ProvenanceZipResult:
198 counters: dict[CounterKey, dict[int, int]] = {}
199 for entity in _graph_entities(path):
200 snapshot_uri = _entity_uri(entity, path)
201 parts = snapshot_uri.rsplit("/prov/se/", 1)
202 if len(parts) != 2:
203 raise SourceDataError(
204 f"Invalid provenance entity URI in {path}: {snapshot_uri}"
205 )
206 entity_uri, snapshot_text = parts
207 try:
208 snapshot_number = int(snapshot_text)
209 except ValueError as error:
210 raise SourceDataError(
211 f"Invalid snapshot number in {path}: {snapshot_uri}"
212 ) from error
213 if snapshot_number <= 0:
214 raise SourceDataError(f"Invalid snapshot number in {path}: {snapshot_uri}")
215 key, resource_number = _counter_identity(entity_uri, path)
216 if key not in counters:
217 counters[key] = {}
218 current = (
219 counters[key][resource_number] if resource_number in counters[key] else 0
220 )
221 counters[key][resource_number] = max(current, snapshot_number)
222 return ProvenanceZipResult(path=path, counters=counters)
225def scan_provenance(
226 root: str, store: SparseCounterStore, workers: int
227) -> ProvenanceScanResult:
228 zip_files = 0
229 entities = 0
230 paths = collect_zip_files(root, only_prov=True)
231 with create_progress() as progress:
232 task_id = progress.add_task("Scanning provenance ZIP files", total=len(paths))
233 for result in bounded_process_map(paths, process_provenance_zip, workers):
234 zip_files += 1
235 for key, updates in result.counters.items():
236 store.write_updates(key, updates)
237 entities += len(updates)
238 advance_progress(progress, task_id)
239 return ProvenanceScanResult(zip_files=zip_files, entities=entities)
242def scan_data(
243 root: str,
244 store: SparseCounterStore,
245 workers: int,
246 max_examples: int,
247) -> DataScanResult:
248 zip_files = 0
249 entities = 0
250 maxima: dict[CounterKey, int] = {}
251 missing_provenance = 0
252 missing_examples: list[dict[str, object]] = []
253 paths = collect_zip_files(root, only_data=True)
254 with create_progress() as progress:
255 task_id = progress.add_task("Scanning data ZIP files", total=len(paths))
256 for result in bounded_process_map(paths, process_data_zip, workers):
257 zip_files += 1
258 for key, records in result.entities.items():
259 entities += len(records)
260 maxima[key] = max(
261 maxima[key] if key in maxima else 0,
262 records[-1].resource_number,
263 )
264 start = records[0].resource_number
265 counters = store.read_span(key, start, records[-1].resource_number)
266 for record in records:
267 if counters[record.resource_number - start] != 0:
268 continue
269 missing_provenance += 1
270 if len(missing_examples) < max_examples:
271 missing_examples.append(
272 {
273 "entity_uri": record.uri,
274 "zip_file": result.path,
275 }
276 )
277 advance_progress(progress, task_id)
278 return DataScanResult(
279 zip_files=zip_files,
280 entities=entities,
281 maxima=maxima,
282 missing_provenance=missing_provenance,
283 missing_examples=missing_examples,
284 )
287def write_json(path: str, data: dict[str, object]) -> None:
288 absolute_path = os.path.abspath(path)
289 directory = os.path.dirname(absolute_path)
290 os.makedirs(directory, exist_ok=True)
291 temporary_path = f"{absolute_path}.tmp"
292 with open(temporary_path, "wb") as output_file:
293 output_file.write(orjson.dumps(data, option=orjson.OPT_INDENT_2))
294 output_file.write(b"\n")
295 os.replace(temporary_path, absolute_path)
298def _graph_entities(path: str) -> Iterator[dict[str, object]]:
299 try:
300 with zipfile.ZipFile(path) as archive:
301 json_names = [
302 name
303 for name in archive.namelist()
304 if name.endswith(".json") and not name.endswith("/")
305 ]
306 if len(json_names) != 1:
307 raise SourceDataError(
308 f"Expected one JSON member in {path}, found {len(json_names)}"
309 )
310 try:
311 payload = orjson.loads(archive.read(json_names[0]))
312 except orjson.JSONDecodeError as error:
313 raise SourceDataError(f"Invalid JSON in {path}") from error
314 except zipfile.BadZipFile as error:
315 raise SourceDataError(f"Invalid ZIP file: {path}") from error
317 if not isinstance(payload, list):
318 raise SourceDataError(f"Expected a JSON array in {path}")
319 for graph in payload:
320 if not isinstance(graph, dict) or "@graph" not in graph:
321 raise SourceDataError(f"Invalid JSON-LD graph in {path}")
322 graph_entities = graph["@graph"]
323 if not isinstance(graph_entities, list):
324 raise SourceDataError(f"Invalid JSON-LD graph in {path}")
325 for entity in graph_entities:
326 if not isinstance(entity, dict):
327 raise SourceDataError(f"Invalid JSON-LD entity in {path}")
328 yield entity
331def _entity_uri(entity: dict[str, object], path: str) -> str:
332 if "@id" not in entity or not isinstance(entity["@id"], str):
333 raise SourceDataError(f"JSON-LD entity without a valid @id in {path}")
334 return entity["@id"]
337def _counter_identity(uri: str, path: str) -> tuple[CounterKey, int]:
338 prefix = get_prefix(uri)
339 short_name = get_short_name(uri)
340 resource_number = get_resource_number(uri)
341 if not prefix or short_name not in ENTITY_TYPES or resource_number <= 0:
342 raise SourceDataError(f"Invalid OpenCitations Meta entity URI in {path}: {uri}")
343 return (prefix, short_name), resource_number