Create final records dataset from diffs and metrics

Why these changes are being introduced:

With the full pipeline mostly formed it was observed that we could
simplify the final artifacts produced by a run by combining the 'diffs'
and 'metrics' datasets into a final 'records' dataset.

This dovetails with the new optional env var 'PRESERVE_ARTIFACTS' where
every other artifact except this one would be removed after successful creation.

How this addresses that need:
* Adds new core function 'create_final_records' that writes a final
dataset containing all records and diff information needed for statistical
and individual record analysis.

Side effects of this change:
* Fewer final disparate datasets for a run

Relevant ticket(s):
* https://mitlibraries.atlassian.net/browse/TIMX-370

Author: ghukill

abdiff/cli.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import shutil
 from datetime import timedelta
 from itertools import chain
 from time import perf_counter
@@ -14,6 +15,7 @@
     calc_ab_diffs,
     calc_ab_metrics,
     collate_ab_transforms,
+    create_final_records,
     init_run,
     run_ab_transforms,
 )
@@ -24,6 +26,8 @@
 
 logger = logging.getLogger(__name__)
 
+CONFIG = Config()
+
 
 @click.group(context_settings={"help_option_names": ["-h", "--help"]})
 @click.option(
@@ -166,19 +170,31 @@ def run_diff(job_directory: str, input_files: str, message: str) -> None:
         image_tag_b=job_data["image_tag_b"],
         input_files=input_files_list,
     )
+
     collated_dataset_path = collate_ab_transforms(
         run_directory=run_directory,
         ab_transformed_file_lists=ab_transformed_file_lists,
     )
+
     diffs_dataset_path = calc_ab_diffs(
         run_directory=run_directory,
         collated_dataset_path=collated_dataset_path,
     )
-    calc_ab_metrics(
+
+    if not CONFIG.preserve_artifacts:
+        shutil.rmtree(collated_dataset_path)
+
+    metrics_dataset_path = calc_ab_metrics(
         run_directory=run_directory,
         diffs_dataset_path=diffs_dataset_path,
     )
 
+    create_final_records(run_directory, diffs_dataset_path, metrics_dataset_path)
+
+    if not CONFIG.preserve_artifacts:
+        shutil.rmtree(diffs_dataset_path)
+        shutil.rmtree(metrics_dataset_path)
+
 
 @main.command()
 @click.option(

abdiff/core/__init__.py

@@ -7,6 +7,7 @@
 from abdiff.core.calc_ab_diffs import calc_ab_diffs
 from abdiff.core.calc_ab_metrics import calc_ab_metrics
 from abdiff.core.collate_ab_transforms import collate_ab_transforms
+from abdiff.core.create_final_records import create_final_records
 from abdiff.core.init_job import init_job
 from abdiff.core.init_run import init_run
 from abdiff.core.run_ab_transforms import run_ab_transforms
@@ -19,4 +20,5 @@
     "collate_ab_transforms",
     "calc_ab_diffs",
     "calc_ab_metrics",
+    "create_final_records",
 ]

abdiff/core/calc_ab_diffs.py

@@ -89,7 +89,6 @@ def get_diffed_batches_iter(
     ) as executor:
         pending_futures = []
         for batch_count, batch in enumerate(batches_iter):
-            logger.info(f"Submitting batch {batch_count} for processing")
             future = executor.submit(process_batch, batch)
             pending_futures.append((batch_count, future))
 

abdiff/core/calc_ab_metrics.py

@@ -20,9 +20,10 @@
 def calc_ab_metrics(
     run_directory: str,
     diffs_dataset_path: str,
-) -> dict:
+) -> str:
 
-    os.makedirs(Path(run_directory) / "metrics", exist_ok=True)
+    metrics_dataset = Path(run_directory) / "metrics"
+    os.makedirs(metrics_dataset, exist_ok=True)
 
     # build field diffs dataframe
     field_matrix_dataset_filepath = create_record_diff_matrix_dataset(
@@ -37,7 +38,7 @@ def calc_ab_metrics(
         run_directory=run_directory, new_data={"metrics": metrics_data}
     )
 
-    return metrics_data
+    return str(metrics_dataset)
 
 
 def create_record_diff_matrix_dataset(

abdiff/core/create_final_records.py

@@ -0,0 +1,112 @@
+import logging
+from collections.abc import Generator
+from pathlib import Path
+
+import duckdb
+import pyarrow as pa
+import pyarrow.dataset as ds
+
+from abdiff.config import Config
+from abdiff.core.utils import load_dataset, write_to_dataset
+
+logger = logging.getLogger(__name__)
+
+CONFIG = Config()
+
+READ_BATCH_SIZE = 1_000
+
+
+def create_final_records(
+    run_directory: str, diffs_dataset_path: str, metrics_dataset_path: str
+) -> str:
+    """Produce a single, final dataset that contains all records and diff information.
+
+    This dataset is produced by joining the "diffs" dataset (which contains the full
+    A and B records, and the JSON diff) with the "metrics" dataset (which is a sparse
+    matrix of TIMDEX fields and boolean 1 or 0 if that record has a diff for that field).
+    This dataset should be sufficient for supporting any webapp data needs.
+    """
+    logger.info("Creating final records dataset from 'diffs' and 'metrics' datasets.")
+
+    diffs_dataset = load_dataset(diffs_dataset_path)
+    metrics_dataset = load_dataset(metrics_dataset_path)
+
+    # get list of unique columns from metrics dataset, and create final dataset schema
+    metrics_timdex_field_columns = [
+        name
+        for name in metrics_dataset.schema.names
+        if name not in diffs_dataset.schema.names
+    ]
+    metrics_columns = (
+        pa.field(name, pa.int64())
+        for name in metrics_dataset.schema.names
+        if name in metrics_timdex_field_columns
+    )
+    final_records_dataset_schema = pa.schema(
+        (
+            pa.field("timdex_record_id", pa.string()),
+            pa.field("source", pa.string()),
+            pa.field("record_a", pa.binary()),
+            pa.field("record_b", pa.binary()),
+            pa.field("ab_diff", pa.string()),
+            pa.field("modified_timdex_fields", pa.list_(pa.string())),
+            pa.field("has_diff", pa.string()),
+            *metrics_columns,  # type: ignore[arg-type]
+        )
+    )
+
+    records_dataset_path = str(Path(run_directory) / "records")
+    write_to_dataset(
+        get_final_records_iter(
+            diffs_dataset, metrics_dataset, metrics_timdex_field_columns
+        ),
+        base_dir=records_dataset_path,
+        schema=final_records_dataset_schema,
+        partition_columns=["source", "has_diff"],
+    )
+
+    return records_dataset_path
+
+
+def get_final_records_iter(
+    diffs_dataset: ds.Dataset,
+    metrics_dataset: ds.Dataset,
+    metrics_timdex_field_columns: list[str],
+) -> Generator[pa.RecordBatch, None, None]:
+
+    with duckdb.connect(":memory:") as conn:
+
+        # register datasets in DuckDB for use
+        conn.register("diffs", diffs_dataset.to_table())
+        conn.register("metrics", metrics_dataset.to_table())
+
+        # prepare select columns
+        select_columns = ",".join(
+            [
+                "d.timdex_record_id",
+                "d.source",
+                "d.record_a",
+                "d.record_b",
+                "d.ab_diff",
+                "d.modified_timdex_fields",
+                "d.has_diff",
+                *[f"m.{name}" for name in metrics_timdex_field_columns],
+            ]
+        )
+
+        results = conn.execute(
+            f"""
+            select {select_columns}
+            from diffs d
+            inner join metrics m on m.timdex_record_id = d.timdex_record_id
+            """
+        ).fetch_record_batch(READ_BATCH_SIZE)
+
+        count = 0
+        while True:
+            try:
+                count += 1
+                logger.info(f"Yielding final records dataset batch: {count}")
+                yield results.read_next_batch()
+            except StopIteration:
+                break

tests/test_calc_ab_metrics.py

@@ -163,8 +163,10 @@ def test_full_metrics_data_has_expected_structure(diff_matrix_dataset_filepath):
 
 
 def test_core_function_updates_run_data(run_directory, diffs_dataset_directory):
-    metrics = calc_ab_metrics(run_directory, diffs_dataset_directory)
-    run_data = read_run_json(run_directory)
+    metrics_dataset_filepath = calc_ab_metrics(run_directory, diffs_dataset_directory)
+
+    metrics_dataset = load_dataset(metrics_dataset_filepath)
+    assert isinstance(metrics_dataset, ds.Dataset)
 
-    assert isinstance(metrics, dict)
-    assert run_data["metrics"] == metrics
+    run_data = read_run_json(run_directory)
+    assert isinstance(run_data["metrics"], dict)

tests/test_cli.py

@@ -136,16 +136,22 @@ def test_view_job_webapp_host_and_port_configurable(
 @patch("abdiff.cli.collate_ab_transforms")
 @patch("abdiff.cli.calc_ab_diffs")
 @patch("abdiff.cli.calc_ab_metrics")
+@patch("abdiff.cli.create_final_records")
 def test_run_diff_success(
     mock_init_run,
     mock_run_ab_transforms,
     mock_collate_ab_transforms,
     mock_calc_ab_diffs,
     mock_calc_ab_metrics,
+    mock_create_final_records,
     caplog,
     runner,
     example_job_directory,
+    monkeypatch,
 ):
+    # skip any attempts to remove any mocked artifacts
+    monkeypatch.setenv("PRESERVE_ARTIFACTS", "true")
+
     # mock initialization of run
     mock_init_run.return_value = str(
         Path(example_job_directory) / "runs" / "2024-10-15_12-00-00"
@@ -170,7 +176,10 @@ def test_run_diff_success(
     mock_calc_ab_diffs.return_value = "path/to/run/diffs"
 
     # mock metrics generation
-    mock_calc_ab_metrics.return_value = {"msg": "these are the from the diffs metrics"}
+    mock_calc_ab_metrics.return_value = "path/to/run/metrics"
+
+    # mock final records dataset
+    mock_create_final_records.return_value = "path/to/run/records"
 
     caplog.set_level("DEBUG")
     result = runner.invoke(
@@ -191,3 +200,4 @@ def test_run_diff_success(
     mock_collate_ab_transforms.assert_called()
     mock_calc_ab_diffs.assert_called()
     mock_calc_ab_metrics.assert_called()
+    mock_create_final_records.assert_called()