diff --git a/src/databricks/labs/lakebridge/reconcile/compare.py b/src/databricks/labs/lakebridge/reconcile/compare.py index c1b62896e8..7ae7c02497 100644 --- a/src/databricks/labs/lakebridge/reconcile/compare.py +++ b/src/databricks/labs/lakebridge/reconcile/compare.py @@ -1,4 +1,21 @@ +"""Source/target compare and reconciliation. + +Three flows share ``_aliased_join`` / ``_join_prepare_persist`` / ``_filter_to_value_mismatches`` where noted: + +1. **Hash row reconcile** — ``reconcile_data``: full outer join on keys, prefixed ``src``/``tgt`` + columns, compare ``hash_value_recon``, missing-side and value-mismatch helpers. +2. **Aggregate reconcile** — ``prepare_persisted_aggregate_join`` then + ``reconcile_agg_data_per_rule``: full or cross join, ``ColumnMapping`` pairs, + ``_mismatch_rows_for_aggregate_mappings``. +3. **Capture (column-level)** — ``capture_mismatch_data_and_columns`` / ``_get_mismatch_df``: + inner join on keys (aliases ``base``/``compare``), per-column ``_base``/``_compare``/``_match`` + projections; keeps all key-matched rows with booleans (does not filter to mismatches only). + +See `lakebridge#745` (Data Compare consolidation). +""" + import logging +from collections.abc import Callable from functools import reduce from pyspark.sql import DataFrame from pyspark.sql.functions import col, expr, lit @@ -22,6 +39,9 @@ _HASH_COLUMN_NAME = "hash_value_recon" _SAMPLE_ROWS = 50 +_CAPTURE_SOURCE_ALIAS = "base" +_CAPTURE_TARGET_ALIAS = "compare" + def _raise_column_mismatch_exception(msg: str, source_missing: list[str], target_missing: list[str]) -> Exception: error_msg = ( @@ -47,12 +67,219 @@ def _build_column_selector(table_name, column_name): return f'{table_name}.{DialectUtils.ansi_normalize_identifier(column_name)} as {alias}' +def _aliased_join( + source: DataFrame, + target: DataFrame, + *, + source_alias: str, + target_alias: str, + how: str, + on=None, +) -> DataFrame: + src = source.alias(source_alias) + tgt = target.alias(target_alias) + if how == "cross": + return src.join(other=tgt, how="cross") + if on is None: + raise ValueError("join condition 'on' is required when how is not 'cross'") + return src.join(other=tgt, on=on, how=how) + + +def _inner_join_for_capture_mismatch( + source: DataFrame, + target: DataFrame, + key_columns: list[str], +) -> DataFrame: + """Inner join on shared key column names (capture / column-level mismatch path, issue #745).""" + return _aliased_join( + source, + target, + source_alias=_CAPTURE_SOURCE_ALIAS, + target_alias=_CAPTURE_TARGET_ALIAS, + how="inner", + on=key_columns, + ) + + +def _persist_reconcile_dataframe(df: DataFrame, persistence: AbstractReconIntermediatePersist) -> DataFrame: + return persistence.write_and_read_df_with_volumes(df) + + +def _join_prepare_persist( + source: DataFrame, + target: DataFrame, + persistence: AbstractReconIntermediatePersist, + *, + source_alias: str, + target_alias: str, + how: str, + on=None, + prepare: Callable[[DataFrame], DataFrame], +) -> DataFrame: + joined = _aliased_join( + source, + target, + source_alias=source_alias, + target_alias=target_alias, + how=how, + on=on, + ) + return _persist_reconcile_dataframe(prepare(joined), persistence) + + +def _select_prefixed_columns_for_hash_reconcile( + joined: DataFrame, + *, + source: DataFrame, + target: DataFrame, + source_alias: str, + target_alias: str, +) -> DataFrame: + return joined.selectExpr( + *[f'{_build_column_selector(source_alias, col_name)}' for col_name in source.columns], + *[f'{_build_column_selector(target_alias, col_name)}' for col_name in target.columns], + ) + + +def _select_aggregate_joined_columns( + joined: DataFrame, + *, + source: DataFrame, + target: DataFrame, +) -> DataFrame: + joined_cols = source.columns + target.columns + normalized_joined_cols = [DialectUtils.ansi_normalize_identifier(c) for c in joined_cols] + return joined.select(*normalized_joined_cols) + + +def prepare_persisted_aggregate_join( + source: DataFrame, + target: DataFrame, + key_columns: list[str] | None, + persistence: AbstractReconIntermediatePersist, +) -> DataFrame: + """Full/cross join source and target for aggregate reconciliation, normalize columns, persist. + + Uses the same ``_join_prepare_persist`` pipeline as hash reconciliation (issue #745). + Replaces the former ``join_aggregate_data`` entry point. + """ + source_alias = "src" + target_alias = "tgt" + if key_columns: + how = "full" + join_condition = _generate_agg_join_condition(source_alias, target_alias, key_columns) + else: + how = "cross" + join_condition = None + return _join_prepare_persist( + source, + target, + persistence, + source_alias=source_alias, + target_alias=target_alias, + how=how, + on=join_condition, + prepare=lambda joined: _select_aggregate_joined_columns(joined, source=source, target=target), + ) + + def _build_mismatch_column(table, column): return col(DialectUtils.ansi_normalize_identifier(column)).alias( DialectUtils.unnormalize_identifier(column.replace(f'{table}_', '').lower()) ) +def _mismatch_projection_for_prefixed_columns(df: DataFrame, side_alias: str): + return [ + _build_mismatch_column(side_alias, col_name) for col_name in df.columns if col_name.startswith(f"{side_alias}_") + ] + + +def _joined_rows_missing_on_side( + df: DataFrame, + *, + absent_side_alias: str, + present_side_alias: str, + compare_basename: str = _HASH_COLUMN_NAME, +) -> DataFrame: + """Rows where the join matched only one side: ``absent`` side has null ``{alias}_{compare_basename}``.""" + return ( + df.filter(col(f"{absent_side_alias}_{compare_basename}").isNull()) + .select(*_mismatch_projection_for_prefixed_columns(df, present_side_alias)) + .drop(compare_basename) + ) + + +def _filter_to_value_mismatches( + df: DataFrame, + *, + values_equal, + match_flag_col: str, + row_predicate=None, +) -> DataFrame: + """Keep rows where ``values_equal`` is false (after optional ``row_predicate``). + + Shared by hash reconcile (single pair of columns, both non-null) and aggregate + reconcile (reduced AND of source/target column equalities). + """ + out = df + if row_predicate is not None: + out = out.filter(row_predicate) + return out.withColumn(match_flag_col, values_equal).filter(col(match_flag_col) == lit(False)) + + +def _value_mismatch_where_both_present( + df: DataFrame, + left_col: str, + right_col: str, + *, + match_col_name: str, +) -> DataFrame: + """Hash-style compare: both columns non-null, keep rows where they differ.""" + presence = col(left_col).isNotNull() & col(right_col).isNotNull() + return _filter_to_value_mismatches( + df, + row_predicate=presence, + values_equal=col(left_col) == col(right_col), + match_flag_col=match_col_name, + ) + + +def _mismatch_rows_for_prefixed_compare_column( + df: DataFrame, + *, + source_alias: str, + target_alias: str, + compare_basename: str, + match_flag_col: str, +) -> DataFrame: + """Value mismatches when both sides have ``{alias}_{compare_basename}``; project source-side prefixed columns.""" + src_c = f"{source_alias}_{compare_basename}" + tgt_c = f"{target_alias}_{compare_basename}" + return ( + _value_mismatch_where_both_present(df, src_c, tgt_c, match_col_name=match_flag_col) + .select(*_mismatch_projection_for_prefixed_columns(df, source_alias)) + .drop(compare_basename) + ) + + +def _data_reconcile_output( + *, + mismatch_df: DataFrame | None, + missing_in_src: DataFrame, + missing_in_tgt: DataFrame, +) -> DataReconcileOutput: + mismatch_count = mismatch_df.count() if mismatch_df is not None else 0 + return DataReconcileOutput( + mismatch_count=mismatch_count, + missing_in_src_count=missing_in_src.count(), + missing_in_tgt_count=missing_in_tgt.count(), + missing_in_src=missing_in_src.limit(_SAMPLE_ROWS), + missing_in_tgt=missing_in_tgt.limit(_SAMPLE_ROWS), + mismatch=MismatchOutput(mismatch_df=mismatch_df), + ) + + def reconcile_data( source: DataFrame, target: DataFrame, @@ -64,83 +291,42 @@ def reconcile_data( target_alias = "tgt" if report_type not in {"data", "all"}: key_columns = [_HASH_COLUMN_NAME] - df = ( - source.alias(source_alias) - .join( - other=target.alias(target_alias), - on=_generate_join_condition(source_alias, target_alias, key_columns), - how="full", - ) - .selectExpr( - *[f'{_build_column_selector(source_alias, col_name)}' for col_name in source.columns], - *[f'{_build_column_selector(target_alias, col_name)}' for col_name in target.columns], - ) + df = _join_prepare_persist( + source, + target, + persistence, + source_alias=source_alias, + target_alias=target_alias, + how="full", + on=_generate_join_condition(source_alias, target_alias, key_columns), + prepare=lambda joined: _select_prefixed_columns_for_hash_reconcile( + joined, + source=source, + target=target, + source_alias=source_alias, + target_alias=target_alias, + ), ) - - 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 - missing_in_src = ( - df.filter(col(f"{source_alias}_{_HASH_COLUMN_NAME}").isNull()) - .select( - *[ - _build_mismatch_column(target_alias, col_name) - for col_name in df.columns - if col_name.startswith(f'{target_alias}_') - ] - ) - .drop(f"{_HASH_COLUMN_NAME}") - ) - - missing_in_tgt = ( - df.filter(col(f"{target_alias}_{_HASH_COLUMN_NAME}").isNull()) - .select( - *[ - _build_mismatch_column(source_alias, col_name) - for col_name in df.columns - if col_name.startswith(f'{source_alias}_') - ] - ) - .drop(f"{_HASH_COLUMN_NAME}") - ) - mismatch_count = 0 - if mismatch: - mismatch_count = mismatch.count() - - missing_in_src_count = missing_in_src.count() - missing_in_tgt_count = missing_in_tgt.count() - - return DataReconcileOutput( - mismatch_count=mismatch_count, - missing_in_src_count=missing_in_src_count, - missing_in_tgt_count=missing_in_tgt_count, - missing_in_src=missing_in_src.limit(_SAMPLE_ROWS), - missing_in_tgt=missing_in_tgt.limit(_SAMPLE_ROWS), - mismatch=MismatchOutput(mismatch_df=mismatch), + missing_in_src = _joined_rows_missing_on_side(df, absent_side_alias=source_alias, present_side_alias=target_alias) + missing_in_tgt = _joined_rows_missing_on_side(df, absent_side_alias=target_alias, present_side_alias=source_alias) + return _data_reconcile_output( + mismatch_df=mismatch, + missing_in_src=missing_in_src, + missing_in_tgt=missing_in_tgt, ) def _get_mismatch_data(df: DataFrame, src_alias: str, tgt_alias: str) -> DataFrame: - return ( - df.filter( - (col(f"{src_alias}_{_HASH_COLUMN_NAME}").isNotNull()) - & (col(f"{tgt_alias}_{_HASH_COLUMN_NAME}").isNotNull()) - ) - .withColumn( - "hash_match", - col(f"{src_alias}_{_HASH_COLUMN_NAME}") == col(f"{tgt_alias}_{_HASH_COLUMN_NAME}"), - ) - .filter(col("hash_match") == lit(False)) - .select( - *[ - _build_mismatch_column(src_alias, col_name) - for col_name in df.columns - if col_name.startswith(f'{src_alias}_') - ] - ) - .drop(f"{_HASH_COLUMN_NAME}") + return _mismatch_rows_for_prefixed_compare_column( + df, + source_alias=src_alias, + target_alias=tgt_alias, + compare_basename=_HASH_COLUMN_NAME, + match_flag_col="hash_match", ) @@ -153,6 +339,7 @@ def _build_capture_df(df: DataFrame) -> DataFrame: def capture_mismatch_data_and_columns(source: DataFrame, target: DataFrame, key_columns: list[str]) -> MismatchOutput: + """Inner-join capture with per-column ``_match`` flags (not full-outer hash reconcile). Shares ``_aliased_join``.""" source_df = _build_capture_df(source) target_df = _build_capture_df(target) unnormalized_key_columns = [DialectUtils.unnormalize_identifier(column) for column in key_columns] @@ -194,35 +381,43 @@ def _unnormalize_mismatch_df_col(column, suffix): return unnormalized -def _get_mismatch_df(source: DataFrame, target: DataFrame, key_columns: list[str], column_list: list[str]): +def _capture_mismatch_base_compare_projections(column_list: list[str]): + source_alias, compare_alias = _CAPTURE_SOURCE_ALIAS, _CAPTURE_TARGET_ALIAS source_aliased = [ - col('base.' + DialectUtils.ansi_normalize_identifier(column)).alias( - _unnormalize_mismatch_df_col(column, '_base') + col(f"{source_alias}." + DialectUtils.ansi_normalize_identifier(column)).alias( + _unnormalize_mismatch_df_col(column, "_base") ) for column in column_list ] target_aliased = [ - col('compare.' + DialectUtils.ansi_normalize_identifier(column)).alias( - _unnormalize_mismatch_df_col(column, '_compare') + col(f"{compare_alias}." + DialectUtils.ansi_normalize_identifier(column)).alias( + _unnormalize_mismatch_df_col(column, "_compare") ) for column in column_list ] + return source_aliased, target_aliased + - match_expr = [ - expr(f"{_normalize_mismatch_df_col(column,'_base')}=={_normalize_mismatch_df_col(column,'_compare')}").alias( - _unnormalize_mismatch_df_col(column, '_match') +def _capture_mismatch_per_column_match_exprs(column_list: list[str]): + return [ + expr(f"{_normalize_mismatch_df_col(column, '_base')}=={_normalize_mismatch_df_col(column, '_compare')}").alias( + _unnormalize_mismatch_df_col(column, "_match") ) for column in column_list ] + + +def _get_mismatch_df(source: DataFrame, target: DataFrame, key_columns: list[str], column_list: list[str]): + source_aliased, target_aliased = _capture_mismatch_base_compare_projections(column_list) + match_expr = _capture_mismatch_per_column_match_exprs(column_list) key_cols = [col(DialectUtils.ansi_normalize_identifier(column)) for column in key_columns] select_expr = key_cols + source_aliased + target_aliased + match_expr logger.info(f"KEY COLUMNS: {key_columns}") logger.info(f"SELECT COLUMNS: {select_expr}") - mismatch_df = ( - source.alias('base').join(other=target.alias('compare'), on=key_columns, how="inner").select(*select_expr) - ) + joined = _inner_join_for_capture_mismatch(source, target, key_columns) + mismatch_df = joined.select(*select_expr) compare_columns = [ DialectUtils.ansi_normalize_identifier(column) for column in mismatch_df.columns if column not in key_columns @@ -312,43 +507,23 @@ def _generate_match_columns(select_cols: list[ColumnMapping]): return items -def _get_mismatch_agg_data( +def _mismatch_rows_for_aggregate_mappings( df: DataFrame, select_cols: list[ColumnMapping], group_cols: list[ColumnMapping] | None, ) -> DataFrame: - # TODO: Integrate with _get_mismatch_data function - """ - For each rule select columns, generate a match column to compare the aggregated data between Source and Target - - e.g., select_cols = [(source_min_col1, target_min_col1), (source_count_col3, target_count_col3)] - - source_min_col1 | target_min_col1 | match_min_col1 | agg_data_match | - -----------------|--------------------|----------------|-------------------| - 11 | 12 |source_min_col1 == target_min_col1 | False | - - :param df: Joined DataFrame with aggregated data from Source and Target - :param select_cols: Rule specific select columns - :param group_cols: Rule specific group by columns, if any - :return: DataFrame with match__ and agg_data_match columns - to identify the aggregate data mismatch between Source and Target - """ + """Rows where aggregated source/target measures disagree (after optional group-by presence filter).""" df_with_match_cols = df - if group_cols: - # Filter Conditions are in the format of: source_group_by_col1 is not null and target_group_by_col1 is not null filter_conditions = _agg_conditions(group_cols) df_with_match_cols = df_with_match_cols.filter(filter_conditions) - - # Generate match columns for the select columns. e.g., match__ for match_column_name, match_column in _generate_match_columns(select_cols): df_with_match_cols = df_with_match_cols.withColumn(match_column_name, match_column) - - # e.g., source_min_col1 == target_min_col1 and source_count_col3 == target_count_col3 select_conditions = _agg_conditions(select_cols, "select") - - return df_with_match_cols.withColumn("agg_data_match", select_conditions).filter( - col("agg_data_match") == lit(False) + return _filter_to_value_mismatches( + df_with_match_cols, + values_equal=select_conditions, + match_flag_col="agg_data_match", ) @@ -385,8 +560,7 @@ def reconcile_agg_data_per_rule( joined_df_with_rule_cols = joined_df.select(*df_rule_columns) - # Data mismatch between Source and Target aggregated data - mismatch = _get_mismatch_agg_data(joined_df_with_rule_cols, rule_select_columns, rule_group_columns) + mismatch = _mismatch_rows_for_aggregate_mappings(joined_df_with_rule_cols, rule_select_columns, rule_group_columns) # Data missing in Source DataFrame rule_target_columns = set(target_columns).intersection([mapping.target_name for mapping in rule_select_columns]) @@ -403,52 +577,12 @@ def reconcile_agg_data_per_rule( ) # TODO write `missing_in_tgt` to delta - mismatch_count = 0 - if mismatch: - mismatch_count = mismatch.count() - - rule_reconcile_output = DataReconcileOutput( - mismatch_count=mismatch_count, - missing_in_src_count=missing_in_src.count(), - missing_in_tgt_count=missing_in_tgt.count(), - missing_in_src=missing_in_src.limit(_SAMPLE_ROWS), - missing_in_tgt=missing_in_tgt.limit(_SAMPLE_ROWS), - mismatch=MismatchOutput(mismatch_df=mismatch), + return _data_reconcile_output( + mismatch_df=mismatch, + missing_in_src=missing_in_src, + missing_in_tgt=missing_in_tgt, ) - return rule_reconcile_output - - -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" - target_alias = "tgt" - - # Generates group by columns in the format of: - # [(source_group_by_col1, target_group_by_col1), (source_group_by_col2, target_group_by_col2) ... ] - if key_columns: - # If there are Group By columns, do Full join on the grouped columns - df = source.alias(source_alias).join( - other=target.alias(target_alias), - on=_generate_agg_join_condition(source_alias, target_alias, key_columns), - how="full", - ) - else: - # If there is no Group By condition, do Cross join as there is only one record - df = source.alias(source_alias).join( - other=target.alias(target_alias), - how="cross", - ) - - 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) - persisted = persistence.write_and_read_df_with_volumes(joined_df) - return persisted +# Backward-compatible alias for existing imports/callers +join_aggregate_data = prepare_persisted_aggregate_join diff --git a/tests/integration/reconcile/conftest.py b/tests/integration/reconcile/conftest.py index 848b887e7c..c588e8c715 100644 --- a/tests/integration/reconcile/conftest.py +++ b/tests/integration/reconcile/conftest.py @@ -29,7 +29,7 @@ ) from databricks.labs.lakebridge.contexts.application import ApplicationContext from databricks.labs.lakebridge.reconcile.recon_capture import AbstractReconIntermediatePersist -from databricks.labs.lakebridge.reconcile.recon_config import Table, Transformation +from databricks.labs.lakebridge.reconcile.recon_config import JdbcReaderOptions, Table, Transformation logger = logging.getLogger(__name__) @@ -71,6 +71,15 @@ TERADATA_SCHEMA = "lf_test_user" TERADATA_TABLE = "diamonds" +# Spark JDBC parallel read: numeric partition column + stride bounds (not row filters) +DIAMONDS_JDBC_READER_OPTIONS = JdbcReaderOptions( + num_partitions=4, + partition_column="carat", + lower_bound="0.0", + upper_bound="1.0", + fetchsize=100, +) + @pytest.fixture def recon_catalog(make_catalog) -> str: