TIMX 494 - TIMDEXRunManager for producing ETL run metadata

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

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)

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,
+        }