diff --git a/README.md b/README.md index d48b7ee..317f1ee 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,14 @@ Exactly one input option and exactly one output option are required. | `--output-dataset` | All records to a TIMDEX dataset. | | `--output-parquet` | All records to parquet files. | +Source records are stored in the output as they were read: the bytes of each XML `` element (as serialized by lxml), each JSONLines line, or the `source_record` of each dataset or parquet row. + +#### Workers and performance + +`--workers N` transforms records in `N` worker processes. The main process reads source records and writes output; workers parse and transform. Records are handled in input order, so output is identical for any number of workers. The default, `1`, transforms in the main process. + +Workers help most for sources with slow, CPU-bound transforms. Each worker uses roughly 0.5 GB of memory. JSON sources are already fast, and may not benefit. + ```text Usage: transform [OPTIONS] @@ -124,6 +132,11 @@ Options: allows a single run_timestamp to be associated with all outputs for single run_id. + -w, --workers INTEGER RANGE Number of worker processes to transform + records with. With 1, records are + transformed in the main process. Output is + the same for any number of workers. + [default: 1; x>=1] -v, --verbose Pass to log at debug level instead of info --help Show this message and exit. ``` \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 7f7457a..8ea24ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,6 @@ dependencies = [ "beautifulsoup4==4.12.3", "boto3", "click", - "jsonlines", "lxml", "pandas", "pyarrow", diff --git a/tests/conftest.py b/tests/conftest.py index 082f16b..575241d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -74,7 +74,7 @@ def runner(): @pytest.fixture def generic_transformer(): class GenericTransformer(Transformer): - def parse_source_file(self): + def read_source_file(self): pass def parse_source_record(self): diff --git a/tests/sources/test_transformer.py b/tests/sources/test_transformer.py index e112adb..dfb0f17 100644 --- a/tests/sources/test_transformer.py +++ b/tests/sources/test_transformer.py @@ -14,6 +14,7 @@ import transmogrifier.models as timdex from transmogrifier.exceptions import DeletedRecordEvent, SkippedRecordEvent from transmogrifier.sources import transformer as transformer_module +from transmogrifier.sources.json.aardvark import MITAardvark from transmogrifier.sources.jsontransformer import JSONTransformer from transmogrifier.sources.transformer import PARQUET_FILES_SCHEMA, Transformer from transmogrifier.sources.xml.datacite import Datacite @@ -33,7 +34,7 @@ def test_load_exclusion_list(source_transformer, mock_s3_exclusion_list): def test_next_iter_sets_action_skip_when_record_is_excluded(tmp_path): class ExcludingTransformer(Transformer): @classmethod - def parse_source_file(cls, _source_file: str): + def read_source_file(cls, _source_file: str): return iter(()) @classmethod @@ -352,7 +353,7 @@ def test_jsontransformer_parse_source_record_non_dict_raises_error(): JSONTransformer.parse_source_record(b"[1, 2, 3]") -def test_transformer_parse_source_records_from_dataset_yields_parsed_records(): +def test_transformer_read_source_records_from_dataset_yields_bytes(): with mock.patch( "transmogrifier.sources.transformer.TIMDEXDataset" ) as mocked_dataset_class: @@ -364,12 +365,12 @@ def test_transformer_parse_source_records_from_dataset_yields_parsed_records(): ] ) records = list( - JSONTransformer.parse_source_records_from_dataset( + JSONTransformer.read_source_records_from_dataset( "cool-repo", "s3://bucket/dataset" ) ) - assert records == [{"id": "1"}, {"id": "2"}] + assert records == [b'{"id": "1"}', b'{"id": "2"}'] mocked_dataset_class.assert_called_once_with("s3://bucket/dataset") mocked_records.read_dicts_iter.assert_called_once_with( table="current_records", @@ -404,7 +405,7 @@ def test_transformer_load_from_source_parquet(tmp_path, source_parquet_location) assert next(transformer).action == "index" -def test_transformer_parse_source_records_from_parquet_filters_source_and_action( +def test_transformer_read_source_records_from_parquet_filters_source_and_action( tmp_path, ): pq.write_table( @@ -431,14 +432,14 @@ def test_transformer_parse_source_records_from_parquet_filters_source_and_action tmp_path / "records.parquet", ) records = list( - JSONTransformer.parse_source_records_from_parquet( + JSONTransformer.read_source_records_from_parquet( "cool-repo", str(tmp_path / "records.parquet") ) ) - assert records == [{"id": "1"}, {"id": "4"}] + assert records == [b'{"id": "1"}', b'{"id": "4"}'] -def test_transformer_parse_source_records_from_parquet_directory_in_path_order( +def test_transformer_read_source_records_from_parquet_directory_in_path_order( tmp_path, ): data_dir = tmp_path / "export" / "data" @@ -458,11 +459,11 @@ def test_transformer_parse_source_records_from_parquet_directory_in_path_order( data_dir / f"part-{i:05d}.parquet", ) records = list( - JSONTransformer.parse_source_records_from_parquet( + JSONTransformer.read_source_records_from_parquet( "cool-repo", str(tmp_path / "export") ) ) - assert records == [{"id": "0"}, {"id": "1"}] + assert records == [b'{"id": "0"}', b'{"id": "1"}'] def test_transformer_write_to_parquet_files_matches_timmy_export_layout( @@ -544,3 +545,62 @@ def test_transformer_write_to_file_writes_only_indexed_records( assert count == len(lines) == source_transformer.transformed_record_count == 3 # noqa: PLR2004 assert len(source_transformer.deleted_records) == 1 assert all(json.loads(line)["timdex_record_id"] for line in lines) + + +def test_transformer_load_reads_source_records_as_bytes(source_input_file): + transformer = Transformer.load("cool-repo", source_input_file) + assert isinstance(next(transformer.source_records), bytes) + + +def test_transformer_next_iter_parses_bytes_source_record(source_input_file): + transformer = Transformer.load("cool-repo", source_input_file) + with mock.patch.object( + transformer, "parse_source_record", wraps=transformer.parse_source_record + ) as mocked_parse: + dataset_record = next(transformer) + mocked_parse.assert_called_once() + assert isinstance(mocked_parse.call_args.args[0], bytes) + assert dataset_record.action == "index" + + +def test_transformer_next_iter_unparseable_source_record_sets_action_error(): + transformer = MITAardvark("gismit", iter([b"not json"])) + dataset_record = next(transformer) + assert dataset_record.action == "error" + assert dataset_record.source_record == b"not json" + assert transformer.error_record_count == 1 + + +def test_xmltransformer_read_source_file_yields_record_bytes(): + records = list(XMLTransformer.read_source_file("tests/fixtures/oai_pmh_records.xml")) + assert all(isinstance(record, bytes) for record in records) + assert XMLTransformer.parse_source_record(records[0]).find("record") + + +def test_jsontransformer_read_source_file_yields_line_bytes(tmp_path): + source_file = tmp_path / "records.jsonl" + source_file.write_text('{"id": "1"}\n\n{"id": "2"}\r\n') + assert list(JSONTransformer.read_source_file(str(source_file))) == [ + b'{"id": "1"}', + b'{"id": "2"}', + ] + + +def test_transformer_next_iter_stores_original_source_record_bytes(source_input_file): + original_records = list(XMLTransformer.read_source_file(source_input_file)) + transformer = Transformer.load("cool-repo", source_input_file) + stored_records = [dataset_record.source_record for dataset_record in transformer] + assert stored_records == original_records + assert not stored_records[0].startswith(b" 1}, + *[str(code.string) for code in code_elements if code.get("status")][:25], + *all_codes[::50], + "wrong", + ] + + uncached = [] + for code in codes: + with mock.patch.dict(marc._loc_crosswalk_code_cache, clear=True): + uncached.append( + Marc.loc_crosswalk_code_to_name(code, crosswalk, "record-01", "code") + ) + cached = [ + Marc.loc_crosswalk_code_to_name(code, crosswalk, "record-01", "code") + for code in codes + ] + + assert cached == uncached + + +def test_loc_crosswalk_code_to_name_cached_duplicate_code_uses_first_match( + caplog, loc_country_crosswalk +): + # 'ai' is an obsolete code for Anguilla before it is the current code for Armenia + caplog.set_level(logging.DEBUG) + for record_id in ("record-01", "record-02"): + assert ( + Marc.loc_crosswalk_code_to_name( + "ai", loc_country_crosswalk, record_id, "country" + ) + == "Anguilla" + ) + assert f"Record #{record_id} uses an obsolete country code: ai" in caplog.text + + def test_get_main_titles_record_with_title(): marc_xml_records = Marc.parse_source_file( "tests/fixtures/marc/marc_record_all_fields.xml" diff --git a/tests/sources/xml/test_researchdatabases.py b/tests/sources/xml/test_researchdatabases.py index 3cf646a..c433301 100644 --- a/tests/sources/xml/test_researchdatabases.py +++ b/tests/sources/xml/test_researchdatabases.py @@ -62,14 +62,19 @@ def researchdatabases_transformer(): def test_researchdatabases_yields_oai_records_pass(researchdatabases_transformer): - oai_records = list( - researchdatabases_transformer._yield_oai_xml_records_for_indexing( - researchdatabases_transformer.source_file + with patch.object( + ResearchDatabases, + "_yield_api_records_for_deleting", + return_value=iter([]), + ): + oai_records = list( + researchdatabases_transformer.read_source_file( + researchdatabases_transformer.source_file + ) ) - ) assert len(oai_records) == 3 - assert isinstance(oai_records[0], Tag) + assert isinstance(oai_records[0], bytes) def test_researchdatabases_yields_deleted_records_success( diff --git a/tests/test_cli.py b/tests/test_cli.py index e865724..b8ca274 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -365,3 +365,41 @@ def test_transform_output_file(caplog, runner, source_input_file, tmp_path): assert result.exit_code == 0 assert "deleted records: 1" in caplog.text assert f"File write complete - records: 3, file: {output_file}" in caplog.text + + +def test_transform_workers_option(caplog, runner, source_parquet_location, tmp_path): + caplog.set_level("INFO") + result = runner.invoke( + main, + [ + "-s", + "researchdatabases", + "--input-parquet", + source_parquet_location, + "--output-parquet", + str(tmp_path / "output"), + "--workers", + "2", + ], + ) + assert result.exit_code == 0 + assert "Transforming records with 2 worker processes" in caplog.text + assert "Completed transform, total records processed: 3" in caplog.text + + +def test_transform_workers_option_must_be_positive(runner, tmp_path): + result = runner.invoke( + main, + [ + "-s", + "alma", + "--input-file", + "fake-input-file", + "--output-dataset", + str(tmp_path / "dataset"), + "--workers", + "0", + ], + ) + assert result.exit_code == 2 # noqa: PLR2004 + assert "Invalid value for '-w' / '--workers'" in result.output diff --git a/tests/test_parallel.py b/tests/test_parallel.py new file mode 100644 index 0000000..020fba1 --- /dev/null +++ b/tests/test_parallel.py @@ -0,0 +1,103 @@ +import os +from concurrent.futures.process import BrokenProcessPool + +import pytest + +from transmogrifier import parallel +from transmogrifier.exceptions import CriticalError +from transmogrifier.parallel import iter_transformed_records +from transmogrifier.sources.transformer import Transformer +from transmogrifier.sources.xml.datacite import Datacite +from transmogrifier.sources.xmltransformer import XMLTransformer + +COUNT_ATTRIBUTES = ( + "processed_record_count", + "transformed_record_count", + "skipped_record_count", + "error_record_count", + "deleted_records", +) + + +class CriticalErrorDatacite(Datacite): + """Module level, so worker processes can import it.""" + + def get_timdex_record_id(self, _source_record): + raise CriticalError("Catastrophic failure in a worker!") + + +class DyingWorkerDatacite(Datacite): + """Module level, so worker processes can import it.""" + + def get_timdex_record_id(self, _source_record): + os._exit(1) + + +@pytest.fixture +def load_transformer(source_input_file, run_id): + def _load(): + return Transformer.load( + "jpal", + source_input_file, + run_id=run_id, + run_timestamp="2024-06-03T15:30:45", + ) + + return _load + + +@pytest.mark.parametrize( + ("chunk_size", "max_pending_chunks_per_worker"), + [(1, 1), (2, 1), (100, 4)], +) +def test_iter_transformed_records_matches_single_process( + monkeypatch, load_transformer, chunk_size, max_pending_chunks_per_worker +): + monkeypatch.setattr(parallel, "WORKER_CHUNK_SIZE", chunk_size) + monkeypatch.setattr( + parallel, "WORKER_MAX_PENDING_CHUNKS_PER_WORKER", max_pending_chunks_per_worker + ) + single_process = load_transformer() + single_process_records = list(single_process) + + with_workers = load_transformer() + with_workers_records = list(iter_transformed_records(with_workers, workers=2)) + + assert with_workers_records == single_process_records + assert len(with_workers.deleted_records) == 1 + for attribute in COUNT_ATTRIBUTES: + assert getattr(with_workers, attribute) == getattr(single_process, attribute) + + +def test_iter_transformed_records_critical_error_stops_run(source_input_file): + transformer = CriticalErrorDatacite( + "jpal", + XMLTransformer.read_source_file(source_input_file), + source_file=source_input_file, + ) + with pytest.raises(CriticalError, match="Catastrophic failure in a worker!"): + list(iter_transformed_records(transformer, workers=2)) + + +def test_iter_transformed_records_worker_process_death_stops_run(source_input_file): + transformer = DyingWorkerDatacite( + "jpal", + XMLTransformer.read_source_file(source_input_file), + source_file=source_input_file, + ) + with pytest.raises(BrokenProcessPool): + list(iter_transformed_records(transformer, workers=2)) + + +def test_transformer_iterates_provided_transformed_records(load_transformer): + transformer = load_transformer() + provided_records = iter(["a", "b"]) + transformer._transformed_records = provided_records # noqa: SLF001 + + assert iter(transformer) is provided_records + assert list(transformer) == ["a", "b"] + + +def test_transformer_iterates_itself_by_default(load_transformer): + transformer = load_transformer() + assert iter(transformer) is transformer diff --git a/transmogrifier/cli.py b/transmogrifier/cli.py index 5324d04..411c69a 100644 --- a/transmogrifier/cli.py +++ b/transmogrifier/cli.py @@ -5,6 +5,7 @@ import click from transmogrifier.config import SOURCES, configure_logger, configure_sentry +from transmogrifier.parallel import iter_transformed_records from transmogrifier.sources.transformer import Transformer logger = logging.getLogger(__name__) @@ -90,6 +91,15 @@ "possible for the TIMDEX StepFunction to invoke Transmogrifier multiple times, this " "allows a single run_timestamp to be associated with all outputs for single run_id.", ) +@click.option( + "-w", + "--workers", + type=click.IntRange(min=1), + default=1, + show_default=True, + help="Number of worker processes to transform records with. With 1, records are " + "transformed in the main process. Output is the same for any number of workers.", +) @click.option( "-v", "--verbose", is_flag=True, help="Pass to log at debug level instead of info" ) @@ -104,6 +114,7 @@ def main( exclusion_list_path: str, run_id: str, run_timestamp: str, + workers: int, verbose: bool, # noqa: FBT001 ) -> None: start_time = perf_counter() @@ -134,6 +145,11 @@ def main( run_id=run_id, run_timestamp=run_timestamp, ) + if workers > 1: + transformer._transformed_records = iter_transformed_records( # noqa: SLF001 + transformer, workers + ) + if output_file: transformer.write_to_file(output_file) elif output_dataset: diff --git a/transmogrifier/parallel.py b/transmogrifier/parallel.py new file mode 100644 index 0000000..274ea8c --- /dev/null +++ b/transmogrifier/parallel.py @@ -0,0 +1,163 @@ +"""Transform records with a pool of worker processes. + +Transforming records can be CPU heavy, and a Python process effectively uses one CPU core +at a time. When the CLI is run with more than one worker, this module spreads that work +across multiple processes, and keeps multiprocessing concerns out of the Transformer +classes. + +Transformers are used almost untouched: + - Each worker process has its own transformer, which iterates over source records as + usual, but is given one chunk of them at a time. + - The main process transformer reads the source records and writes the output as + usual, but its __iter__() returns the records coming back from the workers, by way of + _transformed_records, instead of transforming them itself. +""" + +from __future__ import annotations + +import itertools +import logging +import multiprocessing +from collections import deque +from concurrent.futures import ProcessPoolExecutor +from typing import TYPE_CHECKING + +from transmogrifier.config import configure_logger, configure_sentry + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + from concurrent.futures import Future + + from bs4 import Tag + from timdex_dataset_api import DatasetRecord + + from transmogrifier.sources.transformer import JSON, Transformer + +logger = logging.getLogger(__name__) + +WORKER_CHUNK_SIZE = 100 +WORKER_MAX_PENDING_CHUNKS_PER_WORKER = 4 + +# set once per worker process by _init_worker(), and used by _transform_chunk() +_worker_transformer: Transformer | None = None + + +def iter_transformed_records( + transformer: Transformer, workers: int +) -> Iterator[DatasetRecord]: + """Yield DatasetRecords for a transformer's source records, using worker processes. + + The main process reads chunks of the transformer's source records and sends them to + workers. Each worker has its own transformer that iterates over one chunk at a time, + as it would for a whole run, and returns the resulting DatasetRecords. + + Records are yielded in input order, and only a limited number of chunks are in flight + at a time. Chunks are only read as records are consumed, which limits memory use. + + The record counts and deleted records of the passed transformer are updated as records + are yielded. A CriticalError in any worker, or a worker process dying, stops the run. + + Args: + transformer: The main process transformer, whose source records are read. + workers: The number of worker processes to use. + """ + logger.info(f"Transforming records with {workers} worker processes") + # "spawn" starts each worker as a fresh Python process, instead of a copy of this one, + # so each worker runs _init_worker() once at startup to set itself up + executor = ProcessPoolExecutor( + max_workers=workers, + mp_context=multiprocessing.get_context("spawn"), + initializer=_init_worker, + initargs=( + type(transformer), + transformer.source, + transformer.exclusion_list_path, + transformer.run_data, + logging.getLogger().getEffectiveLevel() == logging.DEBUG, + ), + ) + + # the upper limit of chunks in flight, which are chunks sent to workers but not yet + # yielded here; this keeps a large input from piling up in memory + max_pending_chunks = workers * WORKER_MAX_PENDING_CHUNKS_PER_WORKER + + pending: deque[Future[list[DatasetRecord]]] = deque() + start_offset = 0 + try: + for chunk in itertools.batched( + transformer.source_records, WORKER_CHUNK_SIZE, strict=False + ): + pending.append(executor.submit(_transform_chunk, start_offset, list(chunk))) + start_offset += len(chunk) + + # at the limit, stop reading source records until the oldest chunk is done + # - workers can finish chunks in any order, but always taking the oldest + # chunk means records are yielded in input order + # - result() blocks until the chunk is done, and re-raises any exception + # from the worker + if len(pending) >= max_pending_chunks: + yield from _tally_run_counts(transformer, pending.popleft().result()) + + while pending: + yield from _tally_run_counts(transformer, pending.popleft().result()) + finally: + executor.shutdown(wait=True, cancel_futures=True) + + +def _tally_run_counts( + transformer: Transformer, records: Iterable[DatasetRecord] +) -> Iterator[DatasetRecord]: + """Update a transformer's record counts and deleted records, and yield each record.""" + for record in records: + transformer.processed_record_count += 1 + match record.action: + case "index": + transformer.transformed_record_count += 1 + case "delete": + transformer.deleted_records.append(record.timdex_record_id) + case "skip": + transformer.skipped_record_count += 1 + case "error": + transformer.error_record_count += 1 + yield record + + +def _init_worker( + transformer_class: type[Transformer], + source: str, + exclusion_list_path: str | None, + run_data: dict, + verbose: bool, # noqa: FBT001 +) -> None: + """Configure this worker process and build its transformer. + + The transformer starts with no source records, iter(()), because _transform_chunk() + sets them for each chunk. + """ + global _worker_transformer # noqa: PLW0603 + configure_logger(logging.getLogger(), verbose=verbose) + configure_sentry() + _worker_transformer = transformer_class( + source, iter(()), exclusion_list_path=exclusion_list_path + ) + _worker_transformer.run_data = run_data + + +def _transform_chunk( + start_offset: int, source_records: list[bytes | dict[str, JSON] | Tag] +) -> list[DatasetRecord]: + """Transform one chunk of source records in a worker process. + + Setting the processed record count means each record gets its position in the whole + run as its run_record_offset. The worker's other counts and deleted records are not + used; the main process tallies them from the returned records. + + Args: + start_offset: The run_record_offset of the first record in the chunk. + source_records: The chunk of source records, in run order. + """ + if _worker_transformer is None: + raise RuntimeError("Worker transformer was not initialized") + _worker_transformer.source_records = iter(source_records) + _worker_transformer.processed_record_count = start_offset + return list(_worker_transformer) diff --git a/transmogrifier/sources/jsontransformer.py b/transmogrifier/sources/jsontransformer.py index 1e47677..0303da2 100644 --- a/transmogrifier/sources/jsontransformer.py +++ b/transmogrifier/sources/jsontransformer.py @@ -4,7 +4,6 @@ from abc import abstractmethod from typing import TYPE_CHECKING, final -import jsonlines import smart_open # type: ignore[import-untyped] from transmogrifier.sources.transformer import JSON, Transformer @@ -18,19 +17,19 @@ class JSONTransformer(Transformer): @final @classmethod - def parse_source_file(cls, source_file: str) -> Iterator[dict[str, JSON]]: + def read_source_file(cls, source_file: str) -> Iterator[bytes]: """ - Parse JSON file and return source records as JSON objects via an iterator. + Read JSONLines file and yield each non-blank line as bytes. May not be overridden. - Validates that records in the file are dicts for proper processing. - Args: source_file: A file containing source records to be transformed. """ - with jsonlines.Reader(smart_open.open(source_file, "r")) as records: - yield from records.iter(type=dict) + with smart_open.open(source_file, "rb") as file: + for line in file: + if line.strip(): + yield line.rstrip(b"\r\n") @final @classmethod diff --git a/transmogrifier/sources/transformer.py b/transmogrifier/sources/transformer.py index c980fb5..471d627 100644 --- a/transmogrifier/sources/transformer.py +++ b/transmogrifier/sources/transformer.py @@ -18,7 +18,6 @@ import pyarrow as pa import pyarrow.compute as pc import pyarrow.dataset as ds -import pyarrow.parquet as pq import smart_open from bs4 import Tag from pyarrow import fs @@ -78,7 +77,7 @@ class Transformer(ABC): def __init__( self, source: str, - source_records: Iterator[dict[str, JSON] | Tag], + source_records: Iterator[bytes | dict[str, JSON] | Tag], exclusion_list_path: str | None = None, source_file: str | None = None, source_dataset: str | None = None, @@ -93,7 +92,8 @@ def __init__( source: Source repository label. Must match a source key from config.SOURCES. exclusion_list_path: S3 or local filepath to exclusion list CSV file. exclsion_list: The exclusion list for this particular source. - source_records: A set of source records to be processed. + source_records: A set of source records to be processed, as bytes to be + parsed during transformation, or already parsed records. source_file: Filepath of the input source file. source_dataset: Location of the TIMDEX dataset with input source records. source_parquet: Location of parquet file(s) with input source records. @@ -105,7 +105,7 @@ def __init__( self._exclusion_list: list[str] | None = None self.source_base_url: str = SOURCES[source]["base-url"] self.source_name = SOURCES[source]["name"] - self.source_records: Iterator[JSON | Tag] = source_records + self.source_records: Iterator[bytes | dict[str, JSON] | Tag] = source_records self.processed_record_count: int = 0 self.transformed_record_count: int = 0 self.skipped_record_count: int = 0 @@ -115,6 +115,11 @@ def __init__( self.source_dataset = source_dataset self.source_parquet = source_parquet + # transformed records are yielded from this private object, by default this + # instance; with multiple workers, the CLI replaces this with records coming + # back from the workers + self._transformed_records: Iterator[DatasetRecord] = self + self.run_data = self.get_run_data( source=source, source_file=source_file, @@ -166,19 +171,25 @@ def load_exclusion_list(self) -> list[str]: @final def __iter__(self) -> Iterator[DatasetRecord]: """Iterate over transformed records.""" - return self + return self._transformed_records @final def __next__(self) -> DatasetRecord: """Return next transformed record.""" while True: + source_record = None transformed_record = None timdex_record_id = None - source_record = next(self.source_records) + record = next(self.source_records) self.processed_record_count += 1 try: + source_record = ( + self.parse_source_record(record) + if isinstance(record, bytes) + else record + ) transformed_record = self.transform(source_record) timdex_record_id = transformed_record.timdex_record_id transformed_record.timdex_provenance = timdex.TimdexProvenance( @@ -210,7 +221,11 @@ def __next__(self) -> DatasetRecord: return DatasetRecord( timdex_record_id=timdex_record_id, - source_record=self.serialize_source_record(source_record), + source_record=( + record + if isinstance(record, bytes) + else self.serialize_source_record(record) + ), transformed_record=( json.dumps(transformed_record.asdict()).encode() if transformed_record @@ -266,14 +281,16 @@ def load( if len([i for i in (source_file, source_dataset, source_parquet) if i]) != 1: raise ValueError(INPUT_SOURCE_ERROR_MESSAGE) + # source records are read as bytes, and parsed one at a time during transform + source_records: Iterator[bytes] if source_file: - source_records = transformer_class.parse_source_file(source_file) + source_records = transformer_class.read_source_file(source_file) elif source_dataset: - source_records = transformer_class.parse_source_records_from_dataset( + source_records = transformer_class.read_source_records_from_dataset( source, source_dataset ) elif source_parquet: - source_records = transformer_class.parse_source_records_from_parquet( + source_records = transformer_class.read_source_records_from_parquet( source, source_parquet ) @@ -401,6 +418,11 @@ def _parse_run_data_from_source_file(source_file: str) -> dict: return run_data def serialize_source_record(self, source_record: Tag | dict) -> bytes | None: + """Serialize an already parsed source record, for storing in the dataset. + + Source records read as bytes are stored as-is; this is only used when a + transformer is given already parsed records. + """ if isinstance(source_record, Tag): return source_record.encode() if isinstance(source_record, dict): @@ -565,9 +587,9 @@ def get_valid_title(self, source_record: dict[str, JSON] | Tag) -> str: @classmethod @abstractmethod - def parse_source_file(cls, source_file: str) -> Iterator[dict[str, JSON] | Tag]: + def read_source_file(cls, source_file: str) -> Iterator[bytes]: """ - Parse source file and return source records via an iterator. + Read source file and yield each source record as bytes, without parsing. Must be overridden by format subclasses. @@ -575,6 +597,17 @@ def parse_source_file(cls, source_file: str) -> Iterator[dict[str, JSON] | Tag]: source_file: A file containing source records to be transformed. """ + @classmethod + def parse_source_file(cls, source_file: str) -> Iterator[dict[str, JSON] | Tag]: + """ + Read source file and yield each source record parsed. + + Args: + source_file: A file containing source records to be transformed. + """ + for source_record in cls.read_source_file(source_file): + yield cls.parse_source_record(source_record) + @classmethod @abstractmethod def parse_source_record(cls, source_record: bytes) -> dict[str, JSON] | Tag: @@ -588,15 +621,13 @@ def parse_source_record(cls, source_record: bytes) -> dict[str, JSON] | Tag: """ @classmethod - def parse_source_records_from_dataset( + def read_source_records_from_dataset( cls, source: str, source_dataset: str - ) -> Iterator[dict[str, JSON] | Tag]: - """Read current source records from a TIMDEX dataset and parse them. + ) -> Iterator[bytes]: + """Read current source records from a TIMDEX dataset, as bytes. - Each record is parsed by the format-specific `parse_source_record` method. - Currently, format-specific sub-classes XMLTransformer and JSONTransformer - define this for parsing their respective file record formats. - Source-specific transformers will generally just inhereit. + Records are parsed during transformation by the format-specific + `parse_source_record` method. Args: source_dataset: Location of TIMDEX dataset with source records. @@ -609,44 +640,38 @@ def parse_source_records_from_dataset( source=source, action="index", ): - yield cls.parse_source_record(record["source_record"]) + yield record["source_record"] @classmethod - def parse_source_records_from_parquet( + def read_source_records_from_parquet( cls, source: str, source_parquet: str - ) -> Iterator[dict[str, JSON] | Tag]: - """Read source records from parquet file(s) and parse them. + ) -> Iterator[bytes]: + """Read source records from parquet file(s), as bytes. The parquet file(s) must have the TIMDEX records columns 'source_record', - 'source', and 'action', e.g. a Timmy export. A directory is read recursively, - in sorted path order, ignoring non-parquet files like a Timmy 'export.json'. - Only rows for the given source with action 'index' are parsed, in file order. + 'source', and 'action'. A directory is read recursively, in sorted path order, + ignoring non-parquet files. Only rows for the given source with action 'index' + are read, in file order. Args: source: Source label to filter rows by. source_parquet: Local or S3 location of a parquet file or directory. """ filesystem, path = _parse_location(source_parquet) - if filesystem.get_file_info(path).type == fs.FileType.Directory: - filepaths = _list_parquet_files(filesystem, path) - else: - filepaths = [path] - logger.info(f"Reading source records from {len(filepaths)} parquet file(s)") - - for filepath in filepaths: - with filesystem.open_input_file(filepath) as file: - for batch in pq.ParquetFile(file).iter_batches( - columns=["source_record", "source", "action"], - batch_size=PARQUET_DATASET_BATCH_SIZE, - ): - mask = pc.and_( - pc.equal(batch.column("source"), source), - pc.equal(batch.column("action"), "index"), - ) - for source_record in ( - batch.filter(mask).column("source_record").to_pylist() - ): - yield cls.parse_source_record(source_record) + dataset = ds.dataset( + path, + filesystem=filesystem, + format="parquet", + exclude_invalid_files=True, + ) + logger.info(f"Reading source records from {len(dataset.files)} parquet file(s)") + + for batch in dataset.to_batches( + columns=["source_record"], + filter=(pc.field("source") == source) & (pc.field("action") == "index"), + batch_size=PARQUET_DATASET_BATCH_SIZE, + ): + yield from batch.column("source_record").to_pylist() @classmethod @abstractmethod @@ -819,14 +844,3 @@ def _parse_location(location: str) -> tuple[fs.FileSystem, str]: if location.startswith("s3://"): return fs.S3FileSystem(), location.removeprefix("s3://") return fs.LocalFileSystem(), os.path.abspath(location) - - -def _list_parquet_files(filesystem: fs.FileSystem, path: str) -> list[str]: - """Return sorted paths of parquet files under a directory, if it exists.""" - if filesystem.get_file_info(path).type != fs.FileType.Directory: - return [] - return sorted( - file_info.path - for file_info in filesystem.get_file_info(fs.FileSelector(path, recursive=True)) - if file_info.path.endswith(".parquet") - ) diff --git a/transmogrifier/sources/xml/marc.py b/transmogrifier/sources/xml/marc.py index ccfec5f..279d2e2 100644 --- a/transmogrifier/sources/xml/marc.py +++ b/transmogrifier/sources/xml/marc.py @@ -2,7 +2,7 @@ from collections import defaultdict from collections.abc import Iterator -from bs4 import Tag # type: ignore[import-untyped] +from bs4 import Tag import transmogrifier.models as timdex from transmogrifier.config import load_external_config @@ -12,6 +12,9 @@ logger = logging.getLogger(__name__) +# cache of LOC crosswalk code element lookups, keyed by (code type, code) +_loc_crosswalk_code_cache: dict[tuple[str, str], Tag | None] = {} + class Marc(XMLTransformer): """Marc transformer.""" @@ -192,7 +195,12 @@ def loc_crosswalk_code_to_name( str | None: If a mapping for the code is found in the crosswalk, the mapped value is returned; else None is returned. """ - code_element = crosswalk.find("code", string=code) + # searching the full crosswalk is expensive and repeated for every record, so + # cache the code element found for each code + cache_key = (code_type, code) + if cache_key not in _loc_crosswalk_code_cache: + _loc_crosswalk_code_cache[cache_key] = crosswalk.find("code", string=code) + code_element = _loc_crosswalk_code_cache[cache_key] if code_element is None: logger.debug( "Record #%s uses an invalid %s code: %s", record_id, code_type, code diff --git a/transmogrifier/sources/xml/researchdatabases.py b/transmogrifier/sources/xml/researchdatabases.py index cf52100..1ea621a 100644 --- a/transmogrifier/sources/xml/researchdatabases.py +++ b/transmogrifier/sources/xml/researchdatabases.py @@ -2,9 +2,7 @@ from collections.abc import Iterator from datetime import UTC, datetime -import smart_open # type: ignore[import-untyped] from bs4 import BeautifulSoup, Tag -from lxml import etree from timdex_dataset_api import TIMDEXDataset # type: ignore[import-untyped] from transmogrifier import config @@ -20,7 +18,7 @@ class ResearchDatabases(SpringshareOaiDc): @classmethod - def parse_source_file(cls, source_file: str) -> Iterator[Tag]: + def read_source_file(cls, source_file: str) -> Iterator[bytes]: """Yield records from harvested OAI + API gathered records for possible delete. If the following env vars are not set, detecting deleted records will be skipped: @@ -28,27 +26,9 @@ def parse_source_file(cls, source_file: str) -> Iterator[Tag]: - LIBGUIDES_CLIENT_ID - TIMDEX_DATASET_LOCATION """ - yield from cls._yield_oai_xml_records_for_indexing(source_file) - yield from cls._yield_api_records_for_deleting() - - @classmethod - def _yield_oai_xml_records_for_indexing(cls, source_file: str) -> Iterator[Tag]: - """Yield OAI records from extracted XML file. - - This functionality is a direct port from XMLTransformer, but allows for a custom - self.parse_source_file in this transformer class. - """ - with smart_open.open(source_file, "rb") as file: - for _, element in etree.iterparse( - file, - tag="{*}record", - encoding="utf-8", - recover=True, - ): - record_string = etree.tostring(element, encoding="utf-8") - record = cls.parse_bs4_in_isolated_thread(record_string) - yield record - element.clear() + yield from super().read_source_file(source_file) + for record in cls._yield_api_records_for_deleting(): + yield record.encode() @classmethod def _yield_api_records_for_deleting(cls) -> Iterator[Tag]: diff --git a/transmogrifier/sources/xmltransformer.py b/transmogrifier/sources/xmltransformer.py index 65b8269..2005b6d 100644 --- a/transmogrifier/sources/xmltransformer.py +++ b/transmogrifier/sources/xmltransformer.py @@ -17,9 +17,9 @@ class XMLTransformer(Transformer): """XML transformer class.""" @classmethod - def parse_source_file(cls, source_file: str) -> Iterator[Tag]: + def read_source_file(cls, source_file: str) -> Iterator[bytes]: """ - Parse XML file and return source records as bs4 Tags via an iterator. + Read XML file and yield each element as bytes. Args: source_file: A file containing source records to be transformed. @@ -31,9 +31,7 @@ def parse_source_file(cls, source_file: str) -> Iterator[Tag]: encoding="utf-8", recover=True, ): - record_string = etree.tostring(element, encoding="utf-8") - record = cls.parse_bs4_in_isolated_thread(record_string) - yield record + yield etree.tostring(element, encoding="utf-8", with_tail=False) element.clear() @classmethod diff --git a/uv.lock b/uv.lock index c2d2ca3..29ecbfe 100644 --- a/uv.lock +++ b/uv.lock @@ -671,18 +671,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, ] -[[package]] -name = "jsonlines" -version = "4.0.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "attrs" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/35/87/bcda8e46c88d0e34cad2f09ee2d0c7f5957bccdb9791b0b934ec84d84be4/jsonlines-4.0.0.tar.gz", hash = "sha256:0c6d2c09117550c089995247f605ae4cf77dd1533041d366351f6f298822ea74", size = 11359, upload-time = "2023-09-01T12:34:44.187Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/62/d9ba6323b9202dd2fe166beab8a86d29465c41a0288cbe229fac60c1ab8d/jsonlines-4.0.0-py3-none-any.whl", hash = "sha256:185b334ff2ca5a91362993f42e83588a360cf95ce4b71a73548502bda52a7c55", size = 8701, upload-time = "2023-09-01T12:34:42.563Z" }, -] - [[package]] name = "librt" version = "0.13.0" @@ -1700,7 +1688,6 @@ dependencies = [ { name = "beautifulsoup4" }, { name = "boto3" }, { name = "click" }, - { name = "jsonlines" }, { name = "lxml" }, { name = "pandas" }, { name = "pyarrow" }, @@ -1733,7 +1720,6 @@ requires-dist = [ { name = "beautifulsoup4", specifier = "==4.12.3" }, { name = "boto3" }, { name = "click" }, - { name = "jsonlines" }, { name = "lxml" }, { name = "pandas" }, { name = "pyarrow" },