persist in agg recon after source and target join

Author: m-abulazm

src/databricks/labs/lakebridge/reconcile/compare.py

@@ -5,7 +5,9 @@
 
 from databricks.labs.lakebridge.reconcile.connectors.dialect_utils import DialectUtils
 from databricks.labs.lakebridge.reconcile.exception import ColumnMismatchException
-from databricks.labs.lakebridge.reconcile.recon_capture import AbstractReconIntermediatePersist
+from databricks.labs.lakebridge.reconcile.recon_capture import (
+    AbstractReconIntermediatePersist,
+)
 from databricks.labs.lakebridge.reconcile.recon_output_config import (
     DataReconcileOutput,
     MismatchOutput,
@@ -56,7 +58,7 @@ def reconcile_data(
     target: DataFrame,
     key_columns: list[str],
     report_type: str,
-    inter_persist: AbstractReconIntermediatePersist,
+    persistence: AbstractReconIntermediatePersist,
 ) -> DataReconcileOutput:
     source_alias = "src"
     target_alias = "tgt"
@@ -75,7 +77,7 @@ def reconcile_data(
         )
     )
 
-    df = inter_persist.write_and_read_df_with_volumes(df)
+    df = persistence.write_and_read_df_with_volumes(df)
     # Checkpoint after joining source and target to backpressure
 
     mismatch = _get_mismatch_data(df, source_alias, target_alias) if report_type in {"all", "data"} else None
@@ -414,7 +416,12 @@ def reconcile_agg_data_per_rule(
     return rule_reconcile_output
 
 
-def join_aggregate_data(source: DataFrame, target: DataFrame, key_columns: list[str] | None) -> DataFrame:
+def join_aggregate_data(
+    source: DataFrame,
+    target: DataFrame,
+    key_columns: list[str] | None,
+    persistence: AbstractReconIntermediatePersist,
+) -> DataFrame:
     # TODO:  Integrate with reconcile_data function
 
     source_alias = "src"
@@ -439,5 +446,5 @@ def join_aggregate_data(source: DataFrame, target: DataFrame, key_columns: list[
 
     joined_cols = source.columns + target.columns
     normalized_joined_cols = [DialectUtils.ansi_normalize_identifier(col) for col in joined_cols]
-    joined_df = df.select(*normalized_joined_cols)
+    joined_df = persistence.write_and_read_df_with_volumes(df.select(*normalized_joined_cols))
     return joined_df

src/databricks/labs/lakebridge/reconcile/reconciliation.py

@@ -150,7 +150,7 @@ def _get_reconcile_output(
             target=tgt_data,
             key_columns=table_conf.join_columns,
             report_type=self._report_type,
-            inter_persist=self.intermediate_persist,
+            persistence=self.intermediate_persist,
         )
 
     def _get_reconcile_aggregate_output(
@@ -264,6 +264,7 @@ def _get_reconcile_aggregate_output(
                     source=src_data,
                     target=tgt_data,
                     key_columns=src_query_with_rules.group_by_columns,
+                    persistence=self.intermediate_persist,
                 )
             except DataSourceRuntimeException as e:
                 data_source_exception = e

tests/integration/reconcile/test_aggregates_reconcile.py

@@ -3,7 +3,6 @@
 from dataclasses import dataclass
 from pathlib import Path
 
-from unittest.mock import patch
 
 import pytest
 from pyspark.testing import assertDataFrameEqual

tests/integration/reconcile/test_compare.py

@@ -45,7 +45,7 @@ def test_compare_data_for_report_all(
         target=target,
         key_columns=["s_suppkey", "s_nationkey"],
         report_type="all",
-        inter_persist=FakeReconIntermediatePersist(),
+        persistence=FakeReconIntermediatePersist(),
     )
     expected = DataReconcileOutput(
         mismatch_count=1,
@@ -97,7 +97,7 @@ def test_compare_data_for_report_hash(mock_spark, tmp_path: Path):
         target=target,
         key_columns=["s_suppkey", "s_nationkey"],
         report_type="hash",
-        inter_persist=FakeReconIntermediatePersist(),
+        persistence=FakeReconIntermediatePersist(),
     )
     expected = DataReconcileOutput(
         missing_in_src=missing_in_src,
@@ -280,7 +280,7 @@ def test_compare_data_special_column_names(mock_spark, tmp_path: Path):
         target=target,
         key_columns=["`s``supp#`", "`s_nation#`"],
         report_type="all",
-        inter_persist=FakeReconIntermediatePersist(),
+        persistence=FakeReconIntermediatePersist(),
     )
     expected = DataReconcileOutput(
         mismatch_count=1,