diff --git a/tests/conftest.py b/tests/conftest.py index 191a543..4838a6a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,7 @@ # ruff: noqa: D205, D209 +import os import pytest @@ -10,6 +11,7 @@ generate_sample_records_with_simulated_partitions, ) from timdex_dataset_api import TIMDEXDataset +from timdex_dataset_api.dataset import TIMDEXDatasetConfig @pytest.fixture(autouse=True) @@ -90,3 +92,58 @@ def _records_iter(num_records): ) return _records_iter + + +@pytest.fixture +def dataset_with_runs_location(tmp_path) -> str: + """Fixture to simulate a dataset with multiple full and daily ETL runs.""" + location = str(tmp_path / "dataset_with_runs") + os.mkdir(location) + + timdex_dataset = TIMDEXDataset( + location, config=TIMDEXDatasetConfig(max_rows_per_group=75, max_rows_per_file=75) + ) + timdex_dataset.load() + + run_params = [] + + # simulate ETL runs for 'alma' + run_params.extend( + [ + (40, "alma", "2024-12-01", "full", "index", "run-1"), + (20, "alma", "2024-12-15", "daily", "index", "run-2"), + (100, "alma", "2025-01-01", "full", "index", "run-3"), + (50, "alma", "2025-01-02", "daily", "index", "run-4"), + (25, "alma", "2025-01-03", "daily", "index", "run-5"), + (10, "alma", "2025-01-04", "daily", "delete", "run-6"), + (9, "alma", "2025-01-05", "daily", "index", "run-7"), + ] + ) + + # simulate ETL runs for 'dspace' + run_params.extend( + [ + (30, "dspace", "2024-12-02", "full", "index", "run-8"), + (10, "dspace", "2024-12-16", "daily", "index", "run-9"), + (90, "dspace", "2025-02-01", "full", "index", "run-10"), + (40, "dspace", "2025-02-02", "daily", "index", "run-11"), + (15, "dspace", "2025-02-03", "daily", "index", "run-12"), + (5, "dspace", "2025-02-04", "daily", "delete", "run-13"), + (4, "dspace", "2025-02-05", "daily", "index", "run-14"), + ] + ) + + # write to dataset + for params in run_params: + num_records, source, run_date, run_type, action, run_id = params + records = generate_sample_records( + num_records, + source=source, + run_date=run_date, + run_type=run_type, + action=action, + run_id=run_id, + ) + timdex_dataset.write(records) + + return location diff --git a/tests/test_runs.py b/tests/test_runs.py new file mode 100644 index 0000000..0867906 --- /dev/null +++ b/tests/test_runs.py @@ -0,0 +1,85 @@ +# ruff: noqa: SLF001, D205, D209, PLR2004 + +import datetime +from unittest.mock import patch + +import pytest + +from timdex_dataset_api import TIMDEXDataset +from timdex_dataset_api.run import TIMDEXRunManager + + +@pytest.fixture +def timdex_run_manager(dataset_with_runs_location): + timdex_dataset = TIMDEXDataset(dataset_with_runs_location) + return TIMDEXRunManager(timdex_dataset=timdex_dataset) + + +def test_timdex_run_manager_init(dataset_with_runs_location): + timdex_dataset = TIMDEXDataset(dataset_with_runs_location) + timdex_run_manager = TIMDEXRunManager(timdex_dataset=timdex_dataset) + assert timdex_run_manager._runs_metadata_cache is None + + +def test_timdex_run_manager_parse_single_parquet_file_success(timdex_run_manager): + """Parse run metadata from first parquet file in fixture dataset. We know the details + of this ETL run in advance given the deterministic fixture that generated it.""" + parquet_filepath = timdex_run_manager.timdex_dataset.dataset.files[0] + run_metadata = timdex_run_manager._parse_run_metadata_from_parquet_file( + parquet_filepath + ) + assert run_metadata["source"] == "alma" + assert run_metadata["run_date"] == datetime.date(2024, 12, 1) + assert run_metadata["run_type"] == "full" + assert run_metadata["run_id"] == "run-1" + assert run_metadata["num_rows"] == 40 + assert run_metadata["filename"] == parquet_filepath + + +def test_timdex_run_manager_parse_multiple_parquet_files(timdex_run_manager): + parquet_metadata_df = timdex_run_manager._get_parquet_files_run_metadata() + + # assert 16 rows for this per-file dataframe, despite only 14 distinct ETL "runs" + assert len(parquet_metadata_df) == 16 + + # assert each source has metadata for 8 parquet files + assert parquet_metadata_df.source.value_counts().to_dict() == {"alma": 8, "dspace": 8} + + +def test_timdex_run_manager_get_runs_df(timdex_run_manager): + runs_df = timdex_run_manager.get_runs_metadata() + + # assert two "large" runs have multiple parquet files + assert len(runs_df[runs_df.parquet_files_count > 1]) == 2 + + # assert 7 distinct runs per source, despite more parquet files + assert runs_df.source.value_counts().to_dict() == {"alma": 7, "dspace": 7} + + +def test_timdex_run_manager_get_source_current_run_parquet_files_success( + timdex_run_manager, +): + ordered_parquet_files = timdex_run_manager.get_current_source_parquet_files("alma") + + # assert 6 parquet files, despite being 8 total for alma + # this represents the last full run and all daily since + assert len(ordered_parquet_files) + + # assert sorted reverse chronologically + assert "year=2025/month=01/day=05" in ordered_parquet_files[0] + assert "year=2025/month=01/day=01" in ordered_parquet_files[-1] + + +def test_timdex_run_manager_caches_runs_dataframe(timdex_run_manager): + runs_df = timdex_run_manager.get_runs_metadata() + assert timdex_run_manager._runs_metadata_cache is not None + + with patch.object( + timdex_run_manager, "_get_parquet_files_run_metadata" + ) as mocked_intermediate_method: + mocked_intermediate_method.side_effect = Exception( + "I am not reached, cache is used." + ) + runs_df_2 = timdex_run_manager.get_runs_metadata() + + assert runs_df.equals(runs_df_2) diff --git a/timdex_dataset_api/run.py b/timdex_dataset_api/run.py new file mode 100644 index 0000000..0e11536 --- /dev/null +++ b/timdex_dataset_api/run.py @@ -0,0 +1,186 @@ +"""timdex_dataset_api/run.py""" + +import concurrent.futures +import logging +import time +from typing import TYPE_CHECKING + +import pandas as pd +import pyarrow.parquet as pq + +if TYPE_CHECKING: + from timdex_dataset_api.dataset import TIMDEXDataset + +logger = logging.getLogger(__name__) + + +class TIMDEXRunManager: + """Manages and provides access to ETL run metadata from the TIMDEX parquet dataset.""" + + def __init__(self, timdex_dataset: "TIMDEXDataset"): + self.timdex_dataset: TIMDEXDataset = timdex_dataset + if self.timdex_dataset.dataset is None: + self.timdex_dataset.load() + + self._runs_metadata_cache: pd.DataFrame | None = None + + def clear_cache(self) -> None: + self._runs_metadata_cache = None + + def get_runs_metadata(self, *, refresh: bool = False) -> pd.DataFrame: + """Get metadata for all runs in dataset, grouped by run_id. + + The dataframe returned includes the following columns: + - source + - run_date + - run_type + - run_id + - num_rows: total number of records for that run_id + - parquet_files: list of parquet file(s) that are associated with that run + + Args: + refresh: If True, force refresh of cached metadata + """ + start_time = time.perf_counter() + + if self._runs_metadata_cache is not None and not refresh: + return self._runs_metadata_cache + + ungrouped_runs_df = self._get_parquet_files_run_metadata() + if ungrouped_runs_df.empty: + return ungrouped_runs_df + + # group by run_id + grouped_runs_df = ( + ungrouped_runs_df.groupby("run_id") + .agg( + { + "source": "first", + "run_date": "first", + "run_type": "first", + "num_rows": "sum", + "filename": list, + } + ) + .reset_index() + ) + + # add additional metadata + grouped_runs_df = grouped_runs_df.rename(columns={"filename": "parquet_files"}) + grouped_runs_df["parquet_files_count"] = grouped_runs_df["parquet_files"].apply( + lambda x: len(x) + ) + + # sort by run date and source + grouped_runs_df = grouped_runs_df.sort_values( + ["run_date", "source"], ascending=False + ) + + # cache the result + self._runs_metadata_cache = grouped_runs_df + + logger.info( + f"Dataset runs metadata retrieved, elapsed: " + f"{round(time.perf_counter() - start_time, 2)}s, runs: {len(grouped_runs_df)}" + ) + return grouped_runs_df + + def get_current_source_parquet_files(self, source: str) -> list[str]: + """Get reverse chronological list of current parquet files for a source. + + Args: + source: The source identifier to filter runs + """ + runs_df = self.get_runs_metadata() + source_runs_df = runs_df[runs_df.source == source].copy() + + # get last "full" run + full_runs_df = source_runs_df[source_runs_df.run_type == "full"] + if len(full_runs_df) == 0: + raise RuntimeError( + f"Could not find the most recent 'full' run for source: '{source}'" + ) + last_full_run = full_runs_df.iloc[0] + + # get all "daily" runs since + daily_runs_df = source_runs_df[ + (source_runs_df.run_type == "daily") + & (source_runs_df.run_date >= last_full_run.run_date) + ] + + ordered_parquet_files = [] + for _, daily_run in daily_runs_df.iterrows(): + ordered_parquet_files.extend(daily_run.parquet_files) + ordered_parquet_files.extend(last_full_run.parquet_files) + + return ordered_parquet_files + + def _get_parquet_files_run_metadata(self, max_workers: int = 250) -> pd.DataFrame: + """Retrieve run metadata from parquet file(s) in dataset. + + A single ETL run may still be spread across multiple Parquet files making this + data ungrouped by run. + + Args: + max_workers: Maximum number of parallel workers for processing + - a high number is generally safe given the lightweight nature of the + thread's work, just reading a few parquet file header bytes + """ + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [] + for parquet_filepath in self.timdex_dataset.dataset.files: # type: ignore[attr-defined] + future = executor.submit( + self._parse_run_metadata_from_parquet_file, + parquet_filepath, + ) + futures.append(future) + + done, not_done = concurrent.futures.wait( + futures, return_when=concurrent.futures.ALL_COMPLETED + ) + + results = [] + for future in done: + try: + if result := future.result(): + results.append(result) + except Exception: + logger.exception("Error reading run metadata from parquet file.") + + return pd.DataFrame(results) if results else pd.DataFrame() + + def _parse_run_metadata_from_parquet_file(self, parquet_filepath: str) -> dict: + """Parse source, run_date, run_type, and run_id from a single Parquet file. + + The TIMDEX parquet dataset has a characteristic that we can use for extracting + run information from a single row in a parquet file: all rows in the parquet file + share the column values source, run_date, run_type, and run_id. + + Taking this a step further, we can extract these values without even touching a + single proper row from the parquet file, but from reading the parquet file + column statistics. In this way, we can extract run information from a parquet + file by only reading the lightweight parquet file metadata. + + Args: + parquet_filepath: Path to the parquet file + """ + parquet_file = pq.ParquetFile( + parquet_filepath, + filesystem=self.timdex_dataset.filesystem, # type: ignore[union-attr] + ) + file_meta = parquet_file.metadata.to_dict() + num_rows = file_meta["num_rows"] + columns_meta = file_meta["row_groups"][0]["columns"] # type: ignore[typeddict-item] + source = columns_meta[3]["statistics"]["max"] + run_date = columns_meta[4]["statistics"]["max"] + run_type = columns_meta[5]["statistics"]["max"] + run_id = columns_meta[7]["statistics"]["max"] + + return { + "source": source, + "run_date": run_date, + "run_type": run_type, + "run_id": run_id, + "num_rows": num_rows, + "filename": parquet_filepath, + }