Coverage for oc_meta / run / meta / preprocess_input.py: 99%
163 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-07-25 10:39 +0000
1#!/usr/bin/python
3# SPDX-FileCopyrightText: 2025-2026 Arcangelo Massari <arcangelo.massari@unibo.it>
4#
5# SPDX-License-Identifier: ISC
7from __future__ import annotations
9import argparse
10import multiprocessing
11import os
12from concurrent.futures import ProcessPoolExecutor, as_completed
13from dataclasses import dataclass
14from typing import Callable, List
16import redis
17import yaml
18from rich.table import Table
19from rich_argparse import RichHelpFormatter
21from oc_meta.constants import QLEVER_BATCH_SIZE, QLEVER_MAX_WORKERS
22from oc_meta.lib.console import console, create_progress
23from oc_meta.lib.file_manager import get_csv_data, normalize_path, write_csv
24from oc_meta.lib.sparql import run_queries_parallel
25from oc_meta.run.meta.merge_csv import resolve_output_path
27DATACITE_PREFIX = "http://purl.org/spar/datacite/"
30@dataclass
31class ProcessingStats:
32 total_rows: int = 0
33 duplicate_rows: int = 0
34 existing_ids_rows: int = 0
35 processed_rows: int = 0
38@dataclass
39class FileResult:
40 file_path: str
41 rows: list[tuple[tuple[tuple[str, str], ...], dict[str, str]]]
42 stats: ProcessingStats
45def create_redis_connection(host: str, port: int, db: int = 10) -> redis.Redis:
46 return redis.Redis(host=host, port=port, db=db, decode_responses=True)
49def check_ids_existence_batch(
50 rows: list[dict[str, str]], redis_client: redis.Redis
51) -> list[bool]:
52 row_id_lists: list[list[str]] = []
53 for row in rows:
54 ids_str = row["id"]
55 row_id_lists.append(ids_str.split() if ids_str else [])
57 pipe = redis_client.pipeline()
58 for id_list in row_id_lists:
59 for id_str in id_list:
60 pipe.exists(id_str)
62 results = pipe.execute()
64 row_results: list[bool] = []
65 idx = 0
66 for id_list in row_id_lists:
67 if not id_list:
68 row_results.append(False)
69 else:
70 all_exist = True
71 for _ in id_list:
72 if not results[idx]:
73 all_exist = False
74 idx += 1
75 row_results.append(all_exist)
77 return row_results
80def check_ids_sparql(
81 identifiers: set[str],
82 endpoint_url: str,
83 workers: int = QLEVER_MAX_WORKERS,
84 progress_callback: Callable[[int], None] | None = None,
85) -> set[str]:
86 if not identifiers:
87 return set()
89 id_list = sorted(identifiers)
90 batch_queries: list[str] = []
91 batch_sizes: list[int] = []
93 for i in range(0, len(id_list), QLEVER_BATCH_SIZE):
94 batch = id_list[i : i + QLEVER_BATCH_SIZE]
95 values_entries = []
96 for id_str in batch:
97 schema, value = id_str.split(":", 1)
98 escaped_value = value.replace("\\", "\\\\").replace('"', '\\"')
99 values_entries.append(
100 '("{}"^^xsd:string datacite:{})'.format(escaped_value, schema)
101 )
103 query = (
104 "PREFIX datacite: <http://purl.org/spar/datacite/>\n"
105 "PREFIX literal: <http://www.essepuntato.it/2010/06/literalreification/>\n"
106 "PREFIX xsd: <http://www.w3.org/2001/XMLSchema#>\n"
107 "SELECT ?val ?scheme WHERE {{\n"
108 " VALUES (?val ?scheme) {{ {} }}\n"
109 " ?id literal:hasLiteralValue ?val ;\n"
110 " datacite:usesIdentifierScheme ?scheme .\n"
111 "}}"
112 ).format(" ".join(values_entries))
113 batch_queries.append(query)
114 batch_sizes.append(len(batch))
116 all_bindings = run_queries_parallel(
117 endpoint_url, batch_queries, batch_sizes, workers, progress_callback
118 )
120 found: set[str] = set()
121 for bindings in all_bindings:
122 for result in bindings:
123 val = result["val"]["value"]
124 scheme_uri = result["scheme"]["value"]
125 scheme = (
126 scheme_uri[len(DATACITE_PREFIX) :]
127 if scheme_uri.startswith(DATACITE_PREFIX)
128 else scheme_uri
129 )
130 found.add("{}:{}".format(scheme, val))
132 return found
135def get_csv_files(directory: str) -> List[str]:
136 if not os.path.isdir(directory):
137 raise ValueError("The specified path '{}' is not a directory".format(directory))
139 return [
140 os.path.join(directory, f)
141 for f in os.listdir(directory)
142 if f.endswith(".csv") and os.path.isfile(os.path.join(directory, f))
143 ]
146def collect_rows_from_file(file_path: str) -> FileResult:
147 data = get_csv_data(file_path, clean_data=False)
148 stats = ProcessingStats()
149 stats.total_rows = len(data)
150 valid_rows: list[tuple[tuple[tuple[str, str], ...], dict[str, str]]] = []
151 for row in data:
152 row_hash = tuple(sorted(row.items()))
153 valid_rows.append((row_hash, row))
154 return FileResult(file_path=file_path, rows=valid_rows, stats=stats)
157def filter_existing_ids_from_file(
158 file_path: str, redis_host: str, redis_port: int, redis_db: int
159) -> FileResult:
160 redis_client = create_redis_connection(redis_host, redis_port, redis_db)
161 data = get_csv_data(file_path, clean_data=False)
163 stats = ProcessingStats()
164 stats.total_rows = len(data)
166 existence_results = check_ids_existence_batch(data, redis_client)
168 valid_rows: list[tuple[tuple[tuple[str, str], ...], dict[str, str]]] = []
169 for row, exists in zip(data, existence_results):
170 if exists:
171 stats.existing_ids_rows += 1
172 else:
173 row_hash = tuple(sorted(row.items()))
174 valid_rows.append((row_hash, row))
176 return FileResult(file_path=file_path, rows=valid_rows, stats=stats)
179def filter_sparql_results(
180 results: list[FileResult],
181 found_ids: set[str],
182) -> None:
183 for result in results:
184 filtered: list[tuple[tuple[tuple[str, str], ...], dict[str, str]]] = []
185 for row_hash, row in result.rows:
186 ids_str = row["id"]
187 if ids_str:
188 row_ids = ids_str.split()
189 if row_ids and all(id_str in found_ids for id_str in row_ids):
190 result.stats.existing_ids_rows += 1
191 continue
192 filtered.append((row_hash, row))
193 result.rows = filtered
196def deduplicate_and_write(
197 results: list[FileResult],
198 output_path: str,
199 rows_per_file: int | None = None,
200 subfolders: bool = False,
201) -> ProcessingStats:
202 def chunk_file(num: int) -> str:
203 if subfolders:
204 return os.path.join(output_path, str(num), "{}.csv".format(num))
205 return os.path.join(output_path, "{}.csv".format(num))
207 seen_rows: set[tuple[tuple[str, str], ...]] = set()
208 rows_to_write: list[dict[str, str]] = []
209 file_num = 0
211 total_stats = ProcessingStats()
213 with create_progress() as progress:
214 task = progress.add_task("Deduplicating and writing", total=len(results))
216 for result in results:
217 total_stats.total_rows += result.stats.total_rows
218 total_stats.existing_ids_rows += result.stats.existing_ids_rows
220 for row_hash, row in result.rows:
221 if row_hash in seen_rows:
222 total_stats.duplicate_rows += 1
223 continue
225 seen_rows.add(row_hash)
226 total_stats.processed_rows += 1
227 rows_to_write.append(row)
229 if rows_per_file and len(rows_to_write) >= rows_per_file:
230 write_csv(chunk_file(file_num), rows_to_write)
231 file_num += 1
232 rows_to_write = []
234 progress.advance(task)
236 if rows_to_write:
237 output_file = (
238 chunk_file(file_num) if rows_per_file else resolve_output_path(output_path)
239 )
240 write_csv(output_file, rows_to_write)
242 return total_stats
245def print_processing_report(stats: ProcessingStats, num_files: int) -> None:
246 table = Table(title="Processing Report")
247 table.add_column("Metric", style="cyan")
248 table.add_column("Value", style="green")
250 table.add_row("Total input files processed", str(num_files))
251 table.add_row("Total input rows", str(stats.total_rows))
252 table.add_row("Rows discarded (duplicates)", str(stats.duplicate_rows))
253 table.add_row("Rows discarded (existing IDs)", str(stats.existing_ids_rows))
254 table.add_row("Rows written to output", str(stats.processed_rows))
256 if stats.total_rows > 0:
257 duplicate_percent = (stats.duplicate_rows / stats.total_rows) * 100
258 existing_percent = (stats.existing_ids_rows / stats.total_rows) * 100
259 processed_percent = (stats.processed_rows / stats.total_rows) * 100
261 table.add_row("", "")
262 table.add_row("Duplicate rows %", "{:.1f}%".format(duplicate_percent))
263 table.add_row("Existing IDs %", "{:.1f}%".format(existing_percent))
264 table.add_row("Processed rows %", "{:.1f}%".format(processed_percent))
266 console.print(table)
269def main(): # pragma: no cover
270 parser = argparse.ArgumentParser(
271 description=(
272 "Split the input CSVs declared in a meta_config.yaml into chunks. "
273 "Input directory and SPARQL endpoint are read from the config; the worker "
274 "count is set with --workers. "
275 "When rdf_files_only is True each chunk is written to its own subfolder "
276 "(<output>/<n>/<n>.csv), so it can be processed and re-indexed one at a time "
277 "(meta_process does not upload, so cross-chunk de-duplication needs the "
278 "triplestore refreshed between chunks). When rdf_files_only is False all chunks "
279 "share one folder, since meta_process uploads inline and resolves cross-chunk "
280 "duplicates within a single run."
281 ),
282 formatter_class=RichHelpFormatter,
283 )
284 parser.add_argument("config", help="Path to meta_config.yaml")
285 parser.add_argument("output", help="Directory where chunk files are written")
287 output_group = parser.add_mutually_exclusive_group()
288 output_group.add_argument(
289 "--rows-per-file",
290 type=int,
291 default=None,
292 help="Split output into files of N rows each (default: 3000)",
293 )
294 output_group.add_argument(
295 "--single-file",
296 action="store_true",
297 help="Write all output rows to a single CSV file",
298 )
300 filter_group = parser.add_mutually_exclusive_group()
301 filter_group.add_argument(
302 "--sparql",
303 action="store_true",
304 help="Drop rows whose IDs all already exist on Meta, queried via the config triplestore_url",
305 )
306 filter_group.add_argument(
307 "--redis-port",
308 type=int,
309 help="Drop rows whose IDs all already exist, checked against Redis on this port",
310 )
311 parser.add_argument(
312 "--redis-host", default="localhost", help="Redis host (default: localhost)"
313 )
314 parser.add_argument(
315 "--redis-db", type=int, default=10, help="Redis database number (default: 10)"
316 )
317 parser.add_argument(
318 "--workers", type=int, default=4, help="Number of parallel workers (default: 4)"
319 )
320 args = parser.parse_args()
322 with open(args.config, encoding="utf-8") as f:
323 settings = yaml.full_load(f)
325 input_dir = normalize_path(settings["input_csv_dir"])
326 sparql_endpoint = settings["triplestore_url"]
327 rdf_files_only = settings["rdf_files_only"]
329 if args.single_file:
330 rows_per_file = None
331 elif args.rows_per_file is not None:
332 rows_per_file = args.rows_per_file
333 else:
334 rows_per_file = 3000
336 if rows_per_file:
337 os.makedirs(args.output, exist_ok=True)
339 csv_files = get_csv_files(input_dir)
340 if not csv_files:
341 console.print(
342 "[red]No CSV files found in directory: {}[/red]".format(input_dir)
343 )
344 return 1
346 subfolders = rows_per_file is not None and rdf_files_only
347 layout = "one subfolder per chunk" if subfolders else "single folder"
348 if args.redis_port is not None:
349 mode = "redis"
350 elif args.sparql:
351 mode = "sparql"
352 else:
353 mode = "split-only"
355 console.print(
356 "Found [green]{}[/green] CSV files; [green]{}[/green] workers; mode [green]{}[/green]; layout [green]{}[/green]".format(
357 len(csv_files), args.workers, mode, layout
358 )
359 )
361 file_order = {f: i for i, f in enumerate(csv_files)}
363 if args.redis_port is not None:
364 results: list[FileResult] = []
365 with create_progress() as progress:
366 task = progress.add_task("Filtering existing IDs", total=len(csv_files))
367 with ProcessPoolExecutor(
368 max_workers=args.workers,
369 mp_context=multiprocessing.get_context("forkserver"),
370 ) as executor:
371 futures = {
372 executor.submit(
373 filter_existing_ids_from_file,
374 csv_file,
375 args.redis_host,
376 args.redis_port,
377 args.redis_db,
378 ): csv_file
379 for csv_file in csv_files
380 }
381 for future in as_completed(futures):
382 results.append(future.result())
383 progress.advance(task)
384 else:
385 results = []
386 with create_progress() as progress:
387 task = progress.add_task("Reading CSV files", total=len(csv_files))
388 with ProcessPoolExecutor(
389 max_workers=args.workers,
390 mp_context=multiprocessing.get_context("forkserver"),
391 ) as executor:
392 futures = {
393 executor.submit(collect_rows_from_file, f): f for f in csv_files
394 }
395 for future in as_completed(futures):
396 results.append(future.result())
397 progress.advance(task)
399 results.sort(key=lambda r: file_order[r.file_path])
401 if args.sparql:
402 all_ids: set[str] = set()
403 for result in results:
404 for _hash, row in result.rows:
405 ids_str = row["id"]
406 if ids_str:
407 all_ids.update(ids_str.split())
409 if all_ids:
410 console.print(
411 "Checking [green]{}[/green] unique identifiers against SPARQL endpoint".format(
412 len(all_ids)
413 )
414 )
415 with create_progress() as progress:
416 task = progress.add_task("Querying SPARQL", total=len(all_ids))
418 def on_batch(batch_size: int) -> None:
419 progress.advance(task, batch_size)
421 found_ids = check_ids_sparql(
422 all_ids, sparql_endpoint, args.workers, on_batch
423 )
424 else:
425 found_ids = set()
427 filter_sparql_results(results, found_ids)
429 total_stats = deduplicate_and_write(
430 results, args.output, rows_per_file, subfolders=subfolders
431 )
433 print_processing_report(total_stats, len(csv_files))
435 return 0
438if __name__ == "__main__": # pragma: no cover
439 main()