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,
     download_input_files,
     init_run,
     run_ab_transforms,
@@ -25,6 +27,8 @@
 
 logger = logging.getLogger(__name__)
 
+CONFIG = Config()
+
 
 @click.group(context_settings={"help_option_names": ["-h", "--help"]})
 @click.option(
@@ -181,19 +185,31 @@ def run_diff(
         input_files=input_files_list,
         use_local_s3=download_files,
     )
+
     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
@@ -21,4 +22,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,114 @@
+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.
+
+    This dataset is partitioned by source and 'has_diff' boolean.
+    """
+    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

abdiff/webapp/utils.py

@@ -6,6 +6,14 @@
 import pandas as pd
 from flask import g
 
+SPARSE_MATRIX_SKIP_FIELDS = [
+    "record_a",
+    "record_b",
+    "ab_diff",
+    "modified_timdex_fields",
+    "has_diff",
+]
+
 
 def get_run_directory(run_timestamp: str) -> str:
     return str(Path(g.job_directory) / "runs" / run_timestamp)
@@ -16,7 +24,7 @@ def get_record_a_b_versions(
 ) -> tuple[dict, dict]:
     """Retrieve A and B versions of a single record from diffs dataset."""
     with duckdb.connect() as conn:
-        parquet_glob_pattern = f"{run_directory}/diffs/**/*.parquet"
+        parquet_glob_pattern = f"{run_directory}/records/**/*.parquet"
         conn.execute(
             """
             select record_a, record_b
@@ -62,11 +70,13 @@ def duckdb_query_run_metrics(
     with duckdb.connect() as conn:
 
         # prepare view of record diff matrix
-        parquet_glob_pattern = f"{run_directory}/metrics/**/*.parquet"
+        parquet_glob_pattern = f"{run_directory}/records/**/*.parquet"
         conn.execute(
             f"""
             create view record_diff_matrix as (
-                select * from read_parquet(
+                select
+                * exclude ({",".join([f'"{col}"' for col in SPARSE_MATRIX_SKIP_FIELDS])})
+                from read_parquet(
                     '{parquet_glob_pattern}',
                     hive_partitioning=true
                 )
@@ -146,9 +156,10 @@ def get_record_field_diff_summary(run_directory: str, timdex_record_id: str) ->
         )
     record_row = results_df.iloc[0].to_dict()
 
-    skip_fields = ["timdex_record_id", "source", "has_diff"]
     timdex_fields = {
-        field: value for field, value in record_row.items() if field not in skip_fields
+        field: value
+        for field, value in record_row.items()
+        if field not in SPARSE_MATRIX_SKIP_FIELDS
     }
     fields_with_diffs = [field for field, value in timdex_fields.items() if value == 1]
 

tests/conftest.py

@@ -16,7 +16,7 @@
 from click.testing import CliRunner
 from freezegun import freeze_time
 
-from abdiff.core import calc_ab_diffs, init_job, init_run
+from abdiff.core import calc_ab_diffs, create_final_records, init_job, init_run
 from abdiff.core.calc_ab_metrics import (
     _prepare_duckdb_context,
     create_record_diff_matrix_dataset,
@@ -526,3 +526,14 @@ def duckdb_context_with_diff_matrix(
         function_duckdb_connection, diff_matrix_dataset_filepath
     )
     return function_duckdb_connection, fields, sources
+
+
+@pytest.fixture
+def final_records_dataset_path(
+    run_directory, diffs_dataset_directory, diff_matrix_dataset_filepath
+):
+    return create_final_records(
+        run_directory=run_directory,
+        diffs_dataset_path=diffs_dataset_directory,
+        metrics_dataset_path=diff_matrix_dataset_filepath,
+    )

tests/fixtures/jobs/example-job-2/job.json

@@ -1,6 +1,6 @@
 {
   "job_directory": "tests/fixtures/jobs/example-job-2",
   "job_message": "small job for testing purposes",
-  "image_tag_a": "transmogrifier-example-job-2-395e612:latest",
-  "image_tag_b": "transmogrifier-example-job-2-cf1024c:latest"
+  "image_tag_a": "transmogrifier-abdiff-395e612:latest",
+  "image_tag_b": "transmogrifier-abdiff-cf1024c:latest"
 }