databricks-labs-lakebridge 0.10.6__py3-none-any.whl → 0.10.7__py3-none-any.whl

databricks/labs/lakebridge/reconcile/execute.py

@@ -1,81 +1,26 @@
 import logging
-import sys
 import os
-from datetime import datetime
-from uuid import uuid4
+import sys
 
-from pyspark.errors import PySparkException
-from pyspark.sql import DataFrame, SparkSession
-from sqlglot import Dialect
+from databricks.connect import DatabricksSession
+from databricks.labs.blueprint.installation import Installation
+from databricks.sdk import WorkspaceClient
 
 from databricks.labs.lakebridge.config import (
-    DatabaseConfig,
     TableRecon,
     ReconcileConfig,
-    ReconcileMetadataConfig,
 )
-from databricks.labs.lakebridge.transpiler.sqlglot.dialect_utils import get_dialect
-from databricks.labs.lakebridge.reconcile.compare import (
-    capture_mismatch_data_and_columns,
-    reconcile_data,
-    join_aggregate_data,
-    reconcile_agg_data_per_rule,
-)
-from databricks.labs.lakebridge.reconcile.connectors.data_source import DataSource
-from databricks.labs.lakebridge.reconcile.connectors.source_adapter import create_adapter
 from databricks.labs.lakebridge.reconcile.exception import (
-    DataSourceRuntimeException,
-    InvalidInputException,
     ReconciliationException,
 )
-from databricks.labs.lakebridge.reconcile.query_builder.aggregate_query import AggregateQueryBuilder
-from databricks.labs.lakebridge.reconcile.query_builder.count_query import CountQueryBuilder
-from databricks.labs.lakebridge.reconcile.query_builder.hash_query import HashQueryBuilder
-from databricks.labs.lakebridge.reconcile.query_builder.sampling_query import (
-    SamplingQueryBuilder,
-)
-from databricks.labs.lakebridge.reconcile.query_builder.threshold_query import (
-    ThresholdQueryBuilder,
-)
-from databricks.labs.lakebridge.reconcile.recon_capture import (
-    ReconCapture,
-    generate_final_reconcile_output,
-    ReconIntermediatePersist,
-    generate_final_reconcile_aggregate_output,
-)
+from databricks.labs.lakebridge.reconcile.trigger_recon_aggregate_service import TriggerReconAggregateService
 from databricks.labs.lakebridge.reconcile.recon_config import (
-    Schema,
-    Table,
-    AggregateQueryRules,
-    SamplingOptions,
     RECONCILE_OPERATION_NAME,
     AGG_RECONCILE_OPERATION_NAME,
 )
-from databricks.labs.lakebridge.reconcile.recon_output_config import (
-    DataReconcileOutput,
-    ReconcileOutput,
-    ReconcileProcessDuration,
-    SchemaReconcileOutput,
-    ThresholdOutput,
-    ReconcileRecordCount,
-    AggregateQueryOutput,
-)
-from databricks.labs.lakebridge.reconcile.sampler import SamplerFactory
-from databricks.labs.lakebridge.reconcile.schema_compare import SchemaCompare
-from databricks.labs.lakebridge.transpiler.execute import verify_workspace_client
-from databricks.sdk import WorkspaceClient
-from databricks.labs.blueprint.installation import Installation
-from databricks.connect import DatabricksSession
+from databricks.labs.lakebridge.reconcile.trigger_recon_service import TriggerReconService
 
 logger = logging.getLogger(__name__)
-_SAMPLE_ROWS = 50
-
-
-def validate_input(input_value: str, list_of_value: set, message: str):
-    if input_value not in list_of_value:
-        error_message = f"{message} --> {input_value} is not one of {list_of_value}"
-        logger.error(error_message)
-        raise InvalidInputException(error_message)
 
 
 def main(*argv) -> None:
@@ -118,7 +63,7 @@ def _trigger_recon(
     reconcile_config: ReconcileConfig,
 ):
     try:
-        recon_output = recon(
+        recon_output = TriggerReconService.trigger_recon(
             ws=w,
             spark=DatabricksSession.builder.getOrCreate(),
             table_recon=table_recon,
@@ -154,7 +99,8 @@ def _trigger_reconcile_aggregates(
       after logging the error details.
     """
     try:
-        recon_output = reconcile_aggregates(
+        reconcile_config.report_type = "aggregate"
+        recon_output = TriggerReconAggregateService.trigger_recon_aggregates(
             ws=ws,
             spark=DatabricksSession.builder.getOrCreate(),
             table_recon=table_recon,
@@ -167,753 +113,6 @@ def _trigger_reconcile_aggregates(
         raise e
 
 
-def recon(
-    ws: WorkspaceClient,
-    spark: SparkSession,
-    table_recon: TableRecon,
-    reconcile_config: ReconcileConfig,
-    local_test_run: bool = False,
-) -> ReconcileOutput:
-    """[EXPERIMENTAL] Reconcile the data between the source and target tables."""
-    # verify the workspace client and add proper product and version details
-    # TODO For now we are utilising the
-    #  verify_workspace_client from transpile/execute.py file. Later verify_workspace_client function has to be
-    #  refactored
-
-    ws_client: WorkspaceClient = verify_workspace_client(ws)
-
-    # validate the report type
-    report_type = reconcile_config.report_type.lower()
-    logger.info(f"report_type: {report_type}, data_source: {reconcile_config.data_source} ")
-    validate_input(report_type, {"schema", "data", "row", "all"}, "Invalid report type")
-
-    source, target = initialise_data_source(
-        engine=get_dialect(reconcile_config.data_source),
-        spark=spark,
-        ws=ws_client,
-        secret_scope=reconcile_config.secret_scope,
-    )
-
-    recon_id = str(uuid4())
-    # initialise the Reconciliation
-    reconciler = Reconciliation(
-        source,
-        target,
-        reconcile_config.database_config,
-        report_type,
-        SchemaCompare(spark=spark),
-        get_dialect(reconcile_config.data_source),
-        spark,
-        metadata_config=reconcile_config.metadata_config,
-    )
-
-    # initialise the recon capture class
-    recon_capture = ReconCapture(
-        database_config=reconcile_config.database_config,
-        recon_id=recon_id,
-        report_type=report_type,
-        source_dialect=get_dialect(reconcile_config.data_source),
-        ws=ws_client,
-        spark=spark,
-        metadata_config=reconcile_config.metadata_config,
-        local_test_run=local_test_run,
-    )
-
-    for table_conf in table_recon.tables:
-        recon_process_duration = ReconcileProcessDuration(start_ts=str(datetime.now()), end_ts=None)
-        schema_reconcile_output = SchemaReconcileOutput(is_valid=True)
-        data_reconcile_output = DataReconcileOutput()
-        try:
-            src_schema, tgt_schema = _get_schema(
-                source=source, target=target, table_conf=table_conf, database_config=reconcile_config.database_config
-            )
-        except DataSourceRuntimeException as e:
-            schema_reconcile_output = SchemaReconcileOutput(is_valid=False, exception=str(e))
-        else:
-            if report_type in {"schema", "all"}:
-                schema_reconcile_output = _run_reconcile_schema(
-                    reconciler=reconciler, table_conf=table_conf, src_schema=src_schema, tgt_schema=tgt_schema
-                )
-                logger.warning("Schema comparison is completed.")
-
-            if report_type in {"data", "row", "all"}:
-                data_reconcile_output = _run_reconcile_data(
-                    reconciler=reconciler, table_conf=table_conf, src_schema=src_schema, tgt_schema=tgt_schema
-                )
-                logger.warning(f"Reconciliation for '{report_type}' report completed.")
-
-        recon_process_duration.end_ts = str(datetime.now())
-        # Persist the data to the delta tables
-        recon_capture.start(
-            data_reconcile_output=data_reconcile_output,
-            schema_reconcile_output=schema_reconcile_output,
-            table_conf=table_conf,
-            recon_process_duration=recon_process_duration,
-            record_count=reconciler.get_record_count(table_conf, report_type),
-        )
-        if report_type != "schema":
-            ReconIntermediatePersist(
-                spark=spark, path=generate_volume_path(table_conf, reconcile_config.metadata_config)
-            ).clean_unmatched_df_from_volume()
-
-    return _verify_successful_reconciliation(
-        generate_final_reconcile_output(
-            recon_id=recon_id,
-            spark=spark,
-            metadata_config=reconcile_config.metadata_config,
-            local_test_run=local_test_run,
-        )
-    )
-
-
-def _verify_successful_reconciliation(
-    reconcile_output: ReconcileOutput, operation_name: str = "reconcile"
-) -> ReconcileOutput:
-    for table_output in reconcile_output.results:
-        if table_output.exception_message or (
-            table_output.status.column is False
-            or table_output.status.row is False
-            or table_output.status.schema is False
-            or table_output.status.aggregate is False
-        ):
-            raise ReconciliationException(
-                f" Reconciliation failed for one or more tables. Please check the recon metrics for more details."
-                f" **{operation_name}** failed.",
-                reconcile_output=reconcile_output,
-            )
-
-    logger.info("Reconciliation completed successfully.")
-    return reconcile_output
-
-
-def generate_volume_path(table_conf: Table, metadata_config: ReconcileMetadataConfig):
-    catalog = metadata_config.catalog
-    schema = metadata_config.schema
-    return f"/Volumes/{catalog}/{schema}/{metadata_config.volume}/{table_conf.source_name}_{table_conf.target_name}/"
-
-
-def initialise_data_source(
-    ws: WorkspaceClient,
-    spark: SparkSession,
-    engine: Dialect,
-    secret_scope: str,
-):
-    source = create_adapter(engine=engine, spark=spark, ws=ws, secret_scope=secret_scope)
-    target = create_adapter(engine=get_dialect("databricks"), spark=spark, ws=ws, secret_scope=secret_scope)
-
-    return source, target
-
-
-def _get_missing_data(
-    reader: DataSource,
-    sampler: SamplingQueryBuilder,
-    missing_df: DataFrame,
-    catalog: str,
-    schema: str,
-    table_name: str,
-) -> DataFrame:
-    sample_query = sampler.build_query(missing_df)
-    return reader.read_data(
-        catalog=catalog,
-        schema=schema,
-        table=table_name,
-        query=sample_query,
-        options=None,
-    )
-
-
-def reconcile_aggregates(
-    ws: WorkspaceClient,
-    spark: SparkSession,
-    table_recon: TableRecon,
-    reconcile_config: ReconcileConfig,
-    local_test_run: bool = False,
-):
-    """[EXPERIMENTAL] Reconcile the aggregated data between the source and target tables.
-    for e.g., COUNT, SUM, AVG of columns between source and target with or without any specific key/group by columns
-    Supported Aggregate functions: MIN, MAX, COUNT, SUM, AVG, MEAN, MODE, PERCENTILE, STDDEV, VARIANCE, MEDIAN
-    """
-    # verify the workspace client and add proper product and version details
-    # TODO For now we are utilising the
-    #  verify_workspace_client from transpile/execute.py file. Later verify_workspace_client function has to be
-    #  refactored
-
-    ws_client: WorkspaceClient = verify_workspace_client(ws)
-
-    report_type = ""
-    if report_type:
-        logger.info(f"report_type: {report_type}")
-    logger.info(f"data_source: {reconcile_config.data_source}")
-
-    # Read the reconcile_config and initialise the source and target data sources. Target is always Databricks
-    source, target = initialise_data_source(
-        engine=get_dialect(reconcile_config.data_source),
-        spark=spark,
-        ws=ws_client,
-        secret_scope=reconcile_config.secret_scope,
-    )
-
-    # Generate Unique recon_id for every run
-    recon_id = str(uuid4())
-
-    # initialise the Reconciliation
-    reconciler = Reconciliation(
-        source,
-        target,
-        reconcile_config.database_config,
-        report_type,
-        SchemaCompare(spark=spark),
-        get_dialect(reconcile_config.data_source),
-        spark,
-        metadata_config=reconcile_config.metadata_config,
-    )
-
-    # initialise the recon capture class
-    recon_capture = ReconCapture(
-        database_config=reconcile_config.database_config,
-        recon_id=recon_id,
-        report_type=report_type,
-        source_dialect=get_dialect(reconcile_config.data_source),
-        ws=ws_client,
-        spark=spark,
-        metadata_config=reconcile_config.metadata_config,
-        local_test_run=local_test_run,
-    )
-
-    # Get the Aggregated Reconciliation Output for each table
-    for table_conf in table_recon.tables:
-        recon_process_duration = ReconcileProcessDuration(start_ts=str(datetime.now()), end_ts=None)
-        try:
-            src_schema, tgt_schema = _get_schema(
-                source=source,
-                target=target,
-                table_conf=table_conf,
-                database_config=reconcile_config.database_config,
-            )
-        except DataSourceRuntimeException as e:
-            raise ReconciliationException(message=str(e)) from e
-
-        assert table_conf.aggregates, "Aggregates must be defined for Aggregates Reconciliation"
-
-        table_reconcile_agg_output_list: list[AggregateQueryOutput] = _run_reconcile_aggregates(
-            reconciler=reconciler,
-            table_conf=table_conf,
-            src_schema=src_schema,
-            tgt_schema=tgt_schema,
-        )
-
-        recon_process_duration.end_ts = str(datetime.now())
-
-        # Persist the data to the delta tables
-        recon_capture.store_aggregates_metrics(
-            reconcile_agg_output_list=table_reconcile_agg_output_list,
-            table_conf=table_conf,
-            recon_process_duration=recon_process_duration,
-        )
-
-        (
-            ReconIntermediatePersist(
-                spark=spark,
-                path=generate_volume_path(table_conf, reconcile_config.metadata_config),
-            ).clean_unmatched_df_from_volume()
-        )
-
-    return _verify_successful_reconciliation(
-        generate_final_reconcile_aggregate_output(
-            recon_id=recon_id,
-            spark=spark,
-            metadata_config=reconcile_config.metadata_config,
-            local_test_run=local_test_run,
-        ),
-        operation_name=AGG_RECONCILE_OPERATION_NAME,
-    )
-
-
-class Reconciliation:
-
-    def __init__(
-        self,
-        source: DataSource,
-        target: DataSource,
-        database_config: DatabaseConfig,
-        report_type: str,
-        schema_comparator: SchemaCompare,
-        source_engine: Dialect,
-        spark: SparkSession,
-        metadata_config: ReconcileMetadataConfig,
-    ):
-        self._source = source
-        self._target = target
-        self._report_type = report_type
-        self._database_config = database_config
-        self._schema_comparator = schema_comparator
-        self._target_engine = get_dialect("databricks")
-        self._source_engine = source_engine
-        self._spark = spark
-        self._metadata_config = metadata_config
-
-    def reconcile_data(
-        self,
-        table_conf: Table,
-        src_schema: list[Schema],
-        tgt_schema: list[Schema],
-    ) -> DataReconcileOutput:
-        data_reconcile_output = self._get_reconcile_output(table_conf, src_schema, tgt_schema)
-        reconcile_output = data_reconcile_output
-        if self._report_type in {"data", "all"}:
-            reconcile_output = self._get_sample_data(table_conf, data_reconcile_output, src_schema, tgt_schema)
-            if table_conf.get_threshold_columns("source"):
-                reconcile_output.threshold_output = self._reconcile_threshold_data(table_conf, src_schema, tgt_schema)
-
-        if self._report_type == "row" and table_conf.get_threshold_columns("source"):
-            logger.warning("Threshold comparison is ignored for 'row' report type")
-
-        return reconcile_output
-
-    def reconcile_schema(
-        self,
-        src_schema: list[Schema],
-        tgt_schema: list[Schema],
-        table_conf: Table,
-    ):
-        return self._schema_comparator.compare(src_schema, tgt_schema, self._source_engine, table_conf)
-
-    def reconcile_aggregates(
-        self,
-        table_conf: Table,
-        src_schema: list[Schema],
-        tgt_schema: list[Schema],
-    ) -> list[AggregateQueryOutput]:
-        return self._get_reconcile_aggregate_output(table_conf, src_schema, tgt_schema)
-
-    def _get_reconcile_output(
-        self,
-        table_conf,
-        src_schema,
-        tgt_schema,
-    ):
-        src_hash_query = HashQueryBuilder(table_conf, src_schema, "source", self._source_engine).build_query(
-            report_type=self._report_type
-        )
-        tgt_hash_query = HashQueryBuilder(table_conf, tgt_schema, "target", self._source_engine).build_query(
-            report_type=self._report_type
-        )
-        src_data = self._source.read_data(
-            catalog=self._database_config.source_catalog,
-            schema=self._database_config.source_schema,
-            table=table_conf.source_name,
-            query=src_hash_query,
-            options=table_conf.jdbc_reader_options,
-        )
-        tgt_data = self._target.read_data(
-            catalog=self._database_config.target_catalog,
-            schema=self._database_config.target_schema,
-            table=table_conf.target_name,
-            query=tgt_hash_query,
-            options=table_conf.jdbc_reader_options,
-        )
-
-        volume_path = generate_volume_path(table_conf, self._metadata_config)
-        return reconcile_data(
-            source=src_data,
-            target=tgt_data,
-            key_columns=table_conf.join_columns,
-            report_type=self._report_type,
-            spark=self._spark,
-            path=volume_path,
-        )
-
-    def _get_reconcile_aggregate_output(
-        self,
-        table_conf,
-        src_schema,
-        tgt_schema,
-    ):
-        """
-        Creates a single Query, for the aggregates having the same group by columns. (Ex: 1)
-        If there are no group by columns, all the aggregates are clubbed together in a single query. (Ex: 2)
-        Examples:
-            1.  {
-                      "type": "MIN",
-                      "agg_cols": ["COL1"],
-                      "group_by_cols": ["COL4"]
-                    },
-                    {
-                      "type": "MAX",
-                      "agg_cols": ["COL2"],
-                      "group_by_cols": ["COL9"]
-                    },
-                    {
-                      "type": "COUNT",
-                      "agg_cols": ["COL2"],
-                      "group_by_cols": ["COL9"]
-                    },
-                    {
-                      "type": "AVG",
-                      "agg_cols": ["COL3"],
-                      "group_by_cols": ["COL4"]
-                    },
-              Query 1: SELECT MIN(COL1), AVG(COL3) FROM :table GROUP BY COL4
-              Rules: ID  | Aggregate Type | Column | Group By Column
-                         #1,   MIN,                      COL1,     COL4
-                         #2,   AVG,                     COL3,      COL4
-              -------------------------------------------------------
-              Query 2: SELECT MAX(COL2), COUNT(COL2) FROM :table GROUP BY COL9
-              Rules: ID  | Aggregate Type | Column | Group By Column
-                         #1,   MAX,                      COL2,     COL9
-                         #2,   COUNT,                COL2,      COL9
-          2.  {
-              "type": "MAX",
-              "agg_cols": ["COL1"]
-            },
-            {
-              "type": "SUM",
-              "agg_cols": ["COL2"]
-            },
-            {
-              "type": "MAX",
-              "agg_cols": ["COL3"]
-            }
-          Query: SELECT MAX(COL1), SUM(COL2), MAX(COL3) FROM :table
-          Rules: ID  | Aggregate Type | Column | Group By Column
-                     #1, MAX, COL1,
-                     #2, SUM, COL2,
-                     #3, MAX, COL3,
-        """
-
-        src_query_builder = AggregateQueryBuilder(
-            table_conf,
-            src_schema,
-            "source",
-            self._source_engine,
-        )
-
-        # build Aggregate queries for source,
-        src_agg_queries: list[AggregateQueryRules] = src_query_builder.build_queries()
-
-        # There could be one or more queries per table based on the group by columns
-
-        # build Aggregate queries for target(Databricks),
-        tgt_agg_queries: list[AggregateQueryRules] = AggregateQueryBuilder(
-            table_conf,
-            tgt_schema,
-            "target",
-            self._target_engine,
-        ).build_queries()
-
-        volume_path = generate_volume_path(table_conf, self._metadata_config)
-
-        table_agg_output: list[AggregateQueryOutput] = []
-
-        # Iterate over the grouped aggregates and reconcile the data
-        # Zip all the keys, read the source, target data for each Aggregate query
-        # and reconcile on the aggregate data
-        # For e.g., (source_query_GRP1, target_query_GRP1), (source_query_GRP2, target_query_GRP2)
-        for src_query_with_rules, tgt_query_with_rules in zip(src_agg_queries, tgt_agg_queries):
-            # For each Aggregate query, read the Source and Target Data and add a hash column
-
-            rules_reconcile_output: list[AggregateQueryOutput] = []
-            src_data = None
-            tgt_data = None
-            joined_df = None
-            data_source_exception = None
-            try:
-                src_data = self._source.read_data(
-                    catalog=self._database_config.source_catalog,
-                    schema=self._database_config.source_schema,
-                    table=table_conf.source_name,
-                    query=src_query_with_rules.query,
-                    options=table_conf.jdbc_reader_options,
-                )
-                tgt_data = self._target.read_data(
-                    catalog=self._database_config.target_catalog,
-                    schema=self._database_config.target_schema,
-                    table=table_conf.target_name,
-                    query=tgt_query_with_rules.query,
-                    options=table_conf.jdbc_reader_options,
-                )
-                # Join the Source and Target Aggregated data
-                joined_df = join_aggregate_data(
-                    source=src_data,
-                    target=tgt_data,
-                    key_columns=src_query_with_rules.group_by_columns,
-                    spark=self._spark,
-                    path=f"{volume_path}{src_query_with_rules.group_by_columns_as_str}",
-                )
-            except DataSourceRuntimeException as e:
-                data_source_exception = e
-
-            # For each Aggregated Query, reconcile the data based on the rule
-            for rule in src_query_with_rules.rules:
-                if data_source_exception:
-                    rule_reconcile_output = DataReconcileOutput(exception=str(data_source_exception))
-                else:
-                    rule_reconcile_output = reconcile_agg_data_per_rule(
-                        joined_df, src_data.columns, tgt_data.columns, rule
-                    )
-                rules_reconcile_output.append(AggregateQueryOutput(rule=rule, reconcile_output=rule_reconcile_output))
-
-            # For each table, there could be many Aggregated queries.
-            # Collect the list of Rule Reconcile output per each Aggregate query and append it to the list
-            table_agg_output.extend(rules_reconcile_output)
-        return table_agg_output
-
-    def _get_sample_data(
-        self,
-        table_conf,
-        reconcile_output,
-        src_schema,
-        tgt_schema,
-    ):
-        mismatch = None
-        missing_in_src = None
-        missing_in_tgt = None
-
-        if (
-            reconcile_output.mismatch_count > 0
-            or reconcile_output.missing_in_src_count > 0
-            or reconcile_output.missing_in_tgt_count > 0
-        ):
-            src_sampler = SamplingQueryBuilder(table_conf, src_schema, "source", self._source_engine)
-            tgt_sampler = SamplingQueryBuilder(table_conf, tgt_schema, "target", self._target_engine)
-            if reconcile_output.mismatch_count > 0:
-                mismatch = self._get_mismatch_data(
-                    src_sampler,
-                    tgt_sampler,
-                    reconcile_output.mismatch_count,
-                    reconcile_output.mismatch.mismatch_df,
-                    table_conf.join_columns,
-                    table_conf.source_name,
-                    table_conf.target_name,
-                    table_conf.sampling_options,
-                )
-
-            if reconcile_output.missing_in_src_count > 0:
-                missing_in_src = _get_missing_data(
-                    self._target,
-                    tgt_sampler,
-                    reconcile_output.missing_in_src,
-                    self._database_config.target_catalog,
-                    self._database_config.target_schema,
-                    table_conf.target_name,
-                )
-
-            if reconcile_output.missing_in_tgt_count > 0:
-                missing_in_tgt = _get_missing_data(
-                    self._source,
-                    src_sampler,
-                    reconcile_output.missing_in_tgt,
-                    self._database_config.source_catalog,
-                    self._database_config.source_schema,
-                    table_conf.source_name,
-                )
-
-        return DataReconcileOutput(
-            mismatch=mismatch,
-            mismatch_count=reconcile_output.mismatch_count,
-            missing_in_src_count=reconcile_output.missing_in_src_count,
-            missing_in_tgt_count=reconcile_output.missing_in_tgt_count,
-            missing_in_src=missing_in_src,
-            missing_in_tgt=missing_in_tgt,
-        )
-
-    def _get_mismatch_data(
-        self,
-        src_sampler,
-        tgt_sampler,
-        mismatch_count,
-        mismatch,
-        key_columns,
-        src_table: str,
-        tgt_table: str,
-        sampling_options: SamplingOptions,
-    ):
-
-        tgt_sampling_query = tgt_sampler.build_query_with_alias()
-
-        sampling_model_target = self._target.read_data(
-            catalog=self._database_config.target_catalog,
-            schema=self._database_config.target_schema,
-            table=tgt_table,
-            query=tgt_sampling_query,
-            options=None,
-        )
-
-        # Uses pre-calculated `mismatch_count` from `reconcile_output.mismatch_count` to avoid from recomputing `mismatch` for RandomSampler.
-        mismatch_sampler = SamplerFactory.get_sampler(sampling_options)
-        df = mismatch_sampler.sample(mismatch, mismatch_count, key_columns, sampling_model_target).cache()
-
-        src_mismatch_sample_query = src_sampler.build_query(df)
-        tgt_mismatch_sample_query = tgt_sampler.build_query(df)
-
-        src_data = self._source.read_data(
-            catalog=self._database_config.source_catalog,
-            schema=self._database_config.source_schema,
-            table=src_table,
-            query=src_mismatch_sample_query,
-            options=None,
-        )
-        tgt_data = self._target.read_data(
-            catalog=self._database_config.target_catalog,
-            schema=self._database_config.target_schema,
-            table=tgt_table,
-            query=tgt_mismatch_sample_query,
-            options=None,
-        )
-
-        return capture_mismatch_data_and_columns(source=src_data, target=tgt_data, key_columns=key_columns)
-
-    def _reconcile_threshold_data(
-        self,
-        table_conf: Table,
-        src_schema: list[Schema],
-        tgt_schema: list[Schema],
-    ):
-
-        src_data, tgt_data = self._get_threshold_data(table_conf, src_schema, tgt_schema)
-
-        source_view = f"source_{table_conf.source_name}_df_threshold_vw"
-        target_view = f"target_{table_conf.target_name}_df_threshold_vw"
-
-        src_data.createOrReplaceTempView(source_view)
-        tgt_data.createOrReplaceTempView(target_view)
-
-        return self._compute_threshold_comparison(table_conf, src_schema)
-
-    def _get_threshold_data(
-        self,
-        table_conf: Table,
-        src_schema: list[Schema],
-        tgt_schema: list[Schema],
-    ) -> tuple[DataFrame, DataFrame]:
-        src_threshold_query = ThresholdQueryBuilder(
-            table_conf, src_schema, "source", self._source_engine
-        ).build_threshold_query()
-        tgt_threshold_query = ThresholdQueryBuilder(
-            table_conf, tgt_schema, "target", self._target_engine
-        ).build_threshold_query()
-
-        src_data = self._source.read_data(
-            catalog=self._database_config.source_catalog,
-            schema=self._database_config.source_schema,
-            table=table_conf.source_name,
-            query=src_threshold_query,
-            options=table_conf.jdbc_reader_options,
-        )
-        tgt_data = self._target.read_data(
-            catalog=self._database_config.target_catalog,
-            schema=self._database_config.target_schema,
-            table=table_conf.target_name,
-            query=tgt_threshold_query,
-            options=table_conf.jdbc_reader_options,
-        )
-
-        return src_data, tgt_data
-
-    def _compute_threshold_comparison(self, table_conf: Table, src_schema: list[Schema]) -> ThresholdOutput:
-        threshold_comparison_query = ThresholdQueryBuilder(
-            table_conf, src_schema, "target", self._target_engine
-        ).build_comparison_query()
-
-        threshold_result = self._target.read_data(
-            catalog=self._database_config.target_catalog,
-            schema=self._database_config.target_schema,
-            table=table_conf.target_name,
-            query=threshold_comparison_query,
-            options=table_conf.jdbc_reader_options,
-        )
-        threshold_columns = table_conf.get_threshold_columns("source")
-        failed_where_cond = " OR ".join([name + "_match = 'Failed'" for name in threshold_columns])
-        mismatched_df = threshold_result.filter(failed_where_cond)
-        mismatched_count = mismatched_df.count()
-        threshold_df = None
-        if mismatched_count > 0:
-            threshold_df = mismatched_df.limit(_SAMPLE_ROWS)
-
-        return ThresholdOutput(threshold_df=threshold_df, threshold_mismatch_count=mismatched_count)
-
-    def get_record_count(self, table_conf: Table, report_type: str) -> ReconcileRecordCount:
-        if report_type != "schema":
-            source_count_query = CountQueryBuilder(table_conf, "source", self._source_engine).build_query()
-            target_count_query = CountQueryBuilder(table_conf, "target", self._target_engine).build_query()
-            source_count_row = self._source.read_data(
-                catalog=self._database_config.source_catalog,
-                schema=self._database_config.source_schema,
-                table=table_conf.source_name,
-                query=source_count_query,
-                options=None,
-            ).first()
-            target_count_row = self._target.read_data(
-                catalog=self._database_config.target_catalog,
-                schema=self._database_config.target_schema,
-                table=table_conf.target_name,
-                query=target_count_query,
-                options=None,
-            ).first()
-
-            source_count = int(source_count_row[0]) if source_count_row is not None else 0
-            target_count = int(target_count_row[0]) if target_count_row is not None else 0
-
-            return ReconcileRecordCount(source=int(source_count), target=int(target_count))
-        return ReconcileRecordCount()
-
-
-def _get_schema(
-    source: DataSource,
-    target: DataSource,
-    table_conf: Table,
-    database_config: DatabaseConfig,
-) -> tuple[list[Schema], list[Schema]]:
-    src_schema = source.get_schema(
-        catalog=database_config.source_catalog,
-        schema=database_config.source_schema,
-        table=table_conf.source_name,
-    )
-    tgt_schema = target.get_schema(
-        catalog=database_config.target_catalog,
-        schema=database_config.target_schema,
-        table=table_conf.target_name,
-    )
-
-    return src_schema, tgt_schema
-
-
-def _run_reconcile_data(
-    reconciler: Reconciliation,
-    table_conf: Table,
-    src_schema: list[Schema],
-    tgt_schema: list[Schema],
-) -> DataReconcileOutput:
-    try:
-        return reconciler.reconcile_data(table_conf=table_conf, src_schema=src_schema, tgt_schema=tgt_schema)
-    except DataSourceRuntimeException as e:
-        return DataReconcileOutput(exception=str(e))
-
-
-def _run_reconcile_schema(
-    reconciler: Reconciliation,
-    table_conf: Table,
-    src_schema: list[Schema],
-    tgt_schema: list[Schema],
-):
-    try:
-        return reconciler.reconcile_schema(table_conf=table_conf, src_schema=src_schema, tgt_schema=tgt_schema)
-    except PySparkException as e:
-        return SchemaReconcileOutput(is_valid=False, exception=str(e))
-
-
-def _run_reconcile_aggregates(
-    reconciler: Reconciliation,
-    table_conf: Table,
-    src_schema: list[Schema],
-    tgt_schema: list[Schema],
-) -> list[AggregateQueryOutput]:
-    try:
-        return reconciler.reconcile_aggregates(table_conf, src_schema, tgt_schema)
-    except DataSourceRuntimeException as e:
-        return [AggregateQueryOutput(reconcile_output=DataReconcileOutput(exception=str(e)), rule=None)]
-
-
 if __name__ == "__main__":
     if "DATABRICKS_RUNTIME_VERSION" not in os.environ:
         raise SystemExit("Only intended to run in Databricks Runtime")