TIMX 738 - Parallel Workers and Alma Performance Improvements

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 `<record>` 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.
 ```

pyproject.toml

@@ -11,7 +11,6 @@ dependencies = [
   "beautifulsoup4==4.12.3",
   "boto3",
   "click",
-  "jsonlines",
   "lxml",
   "pandas",
   "pyarrow",

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):

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"<?xml")
+
+
+def test_transformer_next_iter_stores_original_json_line(tmp_path):
+    source_file = tmp_path / "gismit-2024-01-01-full-extracted-records-to-index.jsonl"
+    line = (
+        Path("tests/fixtures/aardvark/aardvark_record_all_fields.jsonl")
+        .read_bytes()
+        .strip()
+    )
+    source_file.write_bytes(line + b"\n")
+    transformer = Transformer.load("gismit", str(source_file))
+    assert next(transformer).source_record == line

tests/sources/xml/test_marc.py

@@ -1,11 +1,14 @@
 # ruff: noqa: E501, SLF001
 import logging
+from unittest import mock
 
 import pytest
 from bs4 import BeautifulSoup  # type: ignore[import-untyped]
 
 import transmogrifier.models as timdex
+from transmogrifier.config import load_external_config
 from transmogrifier.exceptions import SkippedRecordEvent
+from transmogrifier.sources.xml import marc
 from transmogrifier.sources.xml.marc import Marc
 
 
@@ -2071,6 +2074,50 @@ def test_loc_crosswalk_code_to_name_returns_name(caplog, loc_country_crosswalk):
     )
 
 
+@pytest.mark.parametrize(
+    "crosswalk_path", ["config/loc-countries.xml", "config/loc-languages.xml"]
+)
+def test_loc_crosswalk_code_to_name_cached_matches_uncached(crosswalk_path):
+    # duplicate and obsolete codes, plus a sample of current codes and an invalid code
+    crosswalk = load_external_config(crosswalk_path, "xml")
+    code_elements = crosswalk.find_all("code")
+    all_codes = [str(code.string) for code in code_elements]
+    codes = [
+        *{code for code in all_codes if all_codes.count(code) > 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"

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(

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