openmetadata-ingestion 1.6.0.0rc1__py3-none-any.whl → 1.6.0.0rc3__py3-none-any.whl

metadata/data_quality/runner/base_test_suite_source.py

@@ -15,25 +15,27 @@ Base source for the data quality used to instantiate a data quality runner with
 from copy import deepcopy
 from typing import Optional, cast
 
-from sqlalchemy import MetaData
-
+from metadata.data_quality.builders.validator_builder import ValidatorBuilder
 from metadata.data_quality.interface.test_suite_interface import TestSuiteInterface
-from metadata.data_quality.interface.test_suite_interface_factory import (
-    test_suite_interface_factory,
-)
 from metadata.data_quality.runner.core import DataTestsRunner
 from metadata.generated.schema.entity.data.table import Table
-from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
-    DatalakeConnection,
-)
 from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+from metadata.generated.schema.entity.services.serviceType import ServiceType
+from metadata.generated.schema.metadataIngestion.testSuitePipeline import (
+    TestSuitePipeline,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-
-NON_SQA_DATABASE_CONNECTIONS = (DatalakeConnection,)
+from metadata.sampler.models import SampleConfig
+from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.profiler_utils import get_context_entities
+from metadata.utils.service_spec.service_spec import (
+    import_sampler_class,
+    import_test_suite_class,
+)
 
 
 class BaseTestSuiteRunner:
@@ -45,11 +47,16 @@ class BaseTestSuiteRunner:
         ometa_client: OpenMetadata,
         entity: Table,
     ):
+        self.validator_builder_class = ValidatorBuilder
         self._interface = None
         self.entity = entity
         self.service_conn_config = self._copy_service_config(config, self.entity.database)  # type: ignore
+        self._interface_type: str = self.service_conn_config.type.value.lower()
+
+        self.source_config = TestSuitePipeline.model_validate(
+            config.source.sourceConfig.config
+        )
         self.ometa_client = ometa_client
-        self.sqa_metadata = self._set_sqa_metadata()
 
     @property
     def interface(self) -> Optional[TestSuiteInterface]:
@@ -87,28 +94,47 @@ class BaseTestSuiteRunner:
 
         return config_copy
 
-    def _set_sqa_metadata(self):
-        """Set sqlalchemy metadata"""
-        if not isinstance(self.service_conn_config, NON_SQA_DATABASE_CONNECTIONS):
-            return MetaData()
-        return None
-
     def create_data_quality_interface(self) -> TestSuiteInterface:
         """Create data quality interface
 
         Returns:
             TestSuiteInterface: a data quality interface
         """
-        data_quality_interface: TestSuiteInterface = (
-            test_suite_interface_factory.create(
-                self.service_conn_config,
-                self.ometa_client,
-                self.entity,
-                sqa_metadata=self.sqa_metadata,
-            )
+        schema_entity, database_entity, _ = get_context_entities(
+            entity=self.entity, metadata=self.ometa_client
+        )
+        test_suite_class = import_test_suite_class(
+            ServiceType.Database,
+            source_type=self._interface_type,
+            source_config_type=self.service_conn_config.type.value,
+        )
+        sampler_class = import_sampler_class(
+            ServiceType.Database,
+            source_type=self._interface_type,
+            source_config_type=self.service_conn_config.type.value,
+        )
+        # This is shared between the sampler and DQ interfaces
+        sampler_interface: SamplerInterface = sampler_class.create(
+            service_connection_config=self.service_conn_config,
+            ometa_client=self.ometa_client,
+            entity=self.entity,
+            schema_entity=schema_entity,
+            database_entity=database_entity,
+            default_sample_config=SampleConfig(
+                profile_sample=self.source_config.profileSample,
+                profile_sample_type=self.source_config.profileSampleType,
+                sampling_method_type=self.source_config.samplingMethodType,
+            ),
+        )
+
+        self.interface: TestSuiteInterface = test_suite_class.create(
+            service_connection_config=self.service_conn_config,
+            ometa_client=self.ometa_client,
+            sampler=sampler_interface,
+            table_entity=self.entity,
+            validator_builder=self.validator_builder_class,
         )
-        self.interface = data_quality_interface
-        return data_quality_interface
+        return self.interface
 
     def get_data_quality_runner(self) -> DataTestsRunner:
         """Get a data quality runner

metadata/data_quality/validations/base_test_handler.py

@@ -19,10 +19,9 @@ import reprlib
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Callable, List, Optional, Type, TypeVar, Union
 
+from pydantic import BaseModel
+
 from metadata.data_quality.validations import utils
-from metadata.data_quality.validations.runtime_param_setter.param_setter import (
-    RuntimeParameterSetter,
-)
 from metadata.generated.schema.tests.basic import (
     TestCaseResult,
     TestCaseStatus,
@@ -37,6 +36,7 @@ if TYPE_CHECKING:
 
 T = TypeVar("T", bound=Callable)
 R = TypeVar("R")
+S = TypeVar("S", bound=BaseModel)
 
 
 class BaseTestValidator(ABC):
@@ -45,8 +45,6 @@ class BaseTestValidator(ABC):
     This can be useful to resolve complex test parameters based on the parameters gibven by the user.
     """
 
-    runtime_parameter_setter: Optional[Type[RuntimeParameterSetter]] = None
-
     def __init__(
         self,
         runner: Union[QueryRunner, List["DataFrame"]],
@@ -120,10 +118,10 @@ class BaseTestValidator(ABC):
             failed_rows = (
                 failed_rows if failed_rows is not None else (row_count - passed_rows)
             )
-            test_case_result.passedRows = passed_rows
-            test_case_result.failedRows = failed_rows
-            test_case_result.passedRowsPercentage = (passed_rows / row_count) * 100
-            test_case_result.failedRowsPercentage = (failed_rows / row_count) * 100  # type: ignore
+            test_case_result.passedRows = int(passed_rows)
+            test_case_result.failedRows = int(failed_rows)
+            test_case_result.passedRowsPercentage = float(passed_rows / row_count) * 100
+            test_case_result.failedRowsPercentage = float(failed_rows / row_count) * 100  # type: ignore
 
         return test_case_result
 
@@ -168,3 +166,10 @@ class BaseTestValidator(ABC):
     def get_predicted_value(self) -> Optional[str]:
         """Get predicted value"""
         return None
+
+    def get_runtime_parameters(self, setter_class: Type[S]) -> S:
+        """Get runtime parameters"""
+        for param in self.test_case.parameterValues or []:
+            if param.name == setter_class.__name__:
+                return setter_class.model_validate_json(param.value)
+        raise ValueError(f"Runtime parameter {setter_class.__name__} not found")

metadata/data_quality/validations/column/sqlalchemy/columnValueLengthsToBeBetween.py

@@ -40,7 +40,7 @@ class ColumnValueLengthsToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValueMaxToBeBetween.py

@@ -38,7 +38,7 @@ class ColumnValueMaxToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValueMeanToBeBetween.py

@@ -39,7 +39,7 @@ class ColumnValueMeanToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValueMedianToBeBetween.py

@@ -39,7 +39,7 @@ class ColumnValueMedianToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValueMinToBeBetween.py

@@ -39,7 +39,7 @@ class ColumnValueMinToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValueStdDevToBeBetween.py

@@ -39,7 +39,7 @@ class ColumnValueStdDevToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesMissingCount.py

@@ -42,7 +42,7 @@ class ColumnValuesMissingCountValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column, **kwargs) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesSumToBeBetween.py

@@ -39,7 +39,7 @@ class ColumnValuesSumToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeAtExpectedLocation.py

@@ -37,7 +37,7 @@ class ColumnValuesToBeAtExpectedLocationValidator(
     def _fetch_data(self, columns: List[str]) -> Iterator:
         """Fetch data from the runner object"""
         self.runner = cast(QueryRunner, self.runner)
-        inspection = inspect(self.runner.table)
+        inspection = inspect(self.runner.dataset)
         table_columns: List[Column] = inspection.c if inspection is not None else []
         cols = [col for col in table_columns if col.name in columns]
         for col in cols:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeBetween.py

@@ -39,7 +39,7 @@ class ColumnValuesToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeInSet.py

@@ -39,7 +39,7 @@ class ColumnValuesToBeInSetValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column, **kwargs) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeNotInSet.py

@@ -39,7 +39,7 @@ class ColumnValuesToBeNotInSetValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column, **kwargs) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeNotNull.py

@@ -42,7 +42,7 @@ class ColumnValuesToBeNotNullValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeUnique.py

@@ -17,7 +17,6 @@ from typing import Optional
 
 from sqlalchemy import Column, inspect
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.orm.util import AliasedClass
 
 from metadata.data_quality.validations.column.base.columnValuesToBeUnique import (
     BaseColumnValuesToBeUniqueValidator,
@@ -41,7 +40,7 @@ class ColumnValuesToBeUniqueValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:
@@ -53,12 +52,7 @@ class ColumnValuesToBeUniqueValidator(
         """
         count = Metrics.COUNT.value(column).fn()
         unique_count = Metrics.UNIQUE_COUNT.value(column).query(
-            sample=self.runner._sample  # pylint: disable=protected-access
-            if isinstance(
-                self.runner._sample,  # pylint: disable=protected-access
-                AliasedClass,
-            )
-            else self.runner.table,
+            sample=self.runner.dataset,
             session=self.runner._session,  # pylint: disable=protected-access
         )  # type: ignore
 

metadata/data_quality/validations/column/sqlalchemy/columnValuesToMatchRegex.py

@@ -43,7 +43,7 @@ class ColumnValuesToMatchRegexValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(

metadata/data_quality/validations/column/sqlalchemy/columnValuesToNotMatchRegex.py

@@ -43,7 +43,7 @@ class ColumnValuesToNotMatchRegexValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column, **kwargs) -> Optional[int]:

metadata/data_quality/validations/runtime_param_setter/param_setter.py

@@ -15,14 +15,13 @@ from pydantic import BaseModel
 
 from metadata.generated.schema.entity.data.table import Table
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.profiler.processor.sampler.sqlalchemy.sampler import SQASampler
 
 
 class RuntimeParameterSetter(ABC):
     """
     Set runtime parameters for a test case. This can be useful for resolving parameters based on the
     state of OpenMetadata entities.
-    You can see an example implemnentation in `TableDiffParamsSetter`.
+    You can see an example implementation in `TableDiffParamsSetter`.
     """
 
     def __init__(
@@ -30,7 +29,7 @@ class RuntimeParameterSetter(ABC):
         ometa_client: OpenMetadata,
         service_connection_config,
         table_entity: Table,
-        sampler: SQASampler,
+        sampler,
     ):
         self.ometa_client = ometa_client
         self.service_connection_config = service_connection_config

metadata/data_quality/validations/runtime_param_setter/param_setter_factory.py

@@ -13,8 +13,8 @@ Module that defines the RuntimeParameterFactory class.
 This class is responsible for creating instances of the RuntimeParameterSetter 
 based on the test case.
 """
-
-from typing import Optional
+import sys
+from typing import Dict, Set, Type
 
 from metadata.data_quality.validations.runtime_param_setter.param_setter import (
     RuntimeParameterSetter,
@@ -22,7 +22,35 @@ from metadata.data_quality.validations.runtime_param_setter.param_setter import
 from metadata.data_quality.validations.runtime_param_setter.table_diff_params_setter import (
     TableDiffParamsSetter,
 )
+from metadata.data_quality.validations.table.sqlalchemy.tableDiff import (
+    TableDiffValidator,
+)
+from metadata.generated.schema.entity.data.table import Table
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.sampler.sampler_interface import SamplerInterface
+
+
+def removesuffix(s: str, suffix: str) -> str:
+    """A custom implementation of removesuffix for python versions < 3.9
+
+    Args:
+        s (str): The string to remove the suffix from
+        suffix (str): The suffix to remove
+
+    Returns:
+        str: The string with the suffix removed
+    """
+    if sys.version_info >= (3, 9):
+        return s.removesuffix(suffix)
+    if s.endswith(suffix):
+        return s[: -len(suffix)]
+    return s
+
+
+def validator_name(test_case_class: Type) -> str:
+    return removesuffix(
+        test_case_class.__name__[0].lower() + test_case_class.__name__[1:], "Validator"
+    )
 
 
 class RuntimeParameterSetterFactory:
@@ -30,25 +58,25 @@ class RuntimeParameterSetterFactory:
 
     def __init__(self) -> None:
         """Set"""
-        self._setter_map = {
-            TableDiffParamsSetter: {"tableDiff"},
+        self._setter_map: Dict[str, Set[Type[RuntimeParameterSetter]]] = {
+            validator_name(TableDiffValidator): {TableDiffParamsSetter},
         }
 
-    def get_runtime_param_setter(
+    def get_runtime_param_setters(
         self,
         name: str,
         ometa: OpenMetadata,
         service_connection_config,
-        table_entity,
-        sampler,
-    ) -> Optional[RuntimeParameterSetter]:
+        table_entity: Table,
+        sampler: SamplerInterface,
+    ) -> Set[RuntimeParameterSetter]:
         """Get the runtime parameter setter"""
-        for setter_cls, validator_names in self._setter_map.items():
-            if name in validator_names:
-                return setter_cls(
-                    ometa,
-                    service_connection_config,
-                    table_entity,
-                    sampler,
-                )
-        return None
+        return {
+            setter(
+                ometa,
+                service_connection_config,
+                table_entity,
+                sampler,
+            )
+            for setter in self._setter_map.get(name, set())
+        }

metadata/data_quality/validations/runtime_param_setter/table_diff_params_setter.py

@@ -10,7 +10,7 @@
 #  limitations under the License.
 """Module that defines the TableDiffParamsSetter class."""
 from ast import literal_eval
-from typing import List, Optional
+from typing import List, Optional, Set
 from urllib.parse import urlparse
 
 from metadata.data_quality.validations import utils
@@ -75,7 +75,9 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
             DatabaseService, table2.service.id, nullable=False
         )
         key_columns = self.get_key_columns(test_case)
-        extra_columns = self.get_extra_columns(key_columns, test_case)
+        extra_columns = self.get_extra_columns(
+            key_columns, test_case, self.table_entity.columns, table2.columns
+        )
         return TableDiffRuntimeParameters(
             table_profile_config=self.table_entity.tableProfilerConfig,
             table1=TableParameter(
@@ -111,19 +113,18 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
                     case_sensitive=case_sensitive_columns,
                 ),
             ),
-            keyColumns=key_columns,
-            extraColumns=extra_columns,
+            keyColumns=list(key_columns),
+            extraColumns=list(extra_columns),
             whereClause=self.build_where_clause(test_case),
         )
 
-    # pylint: disable=protected-access
     def build_where_clause(self, test_case) -> Optional[str]:
         param_where_clause = self.get_parameter(test_case, "where", None)
         partition_where_clause = (
             None
             if not (
-                self.sampler._partition_details
-                and self.sampler._partition_details.enablePartitioning
+                self.sampler.partition_details
+                and self.sampler.partition_details.enablePartitioning
             )
             else self.sampler.get_partitioned_query().whereclause.compile(
                 compile_kwargs={"literal_binds": True}
@@ -135,21 +136,25 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
         return " AND ".join(where_clauses)
 
     def get_extra_columns(
-        self, key_columns: List[str], test_case
-    ) -> Optional[List[str]]:
+        self,
+        key_columns: Set[str],
+        test_case,
+        left_columns: List[Column],
+        right_columns: List[Column],
+    ) -> Optional[Set[str]]:
         extra_columns_param = self.get_parameter(test_case, "useColumns", None)
         if extra_columns_param is not None:
             extra_columns: List[str] = literal_eval(extra_columns_param)
             self.validate_columns(extra_columns)
-            return extra_columns
+            return set(extra_columns)
         if extra_columns_param is None:
             extra_columns_param = []
-            for column in self.table_entity.columns:
+            for column in left_columns + right_columns:
                 if column.name.root not in key_columns:
                     extra_columns_param.insert(0, column.name.root)
-        return extra_columns_param
+        return set(extra_columns_param)
 
-    def get_key_columns(self, test_case) -> List[str]:
+    def get_key_columns(self, test_case) -> Set[str]:
         key_columns_param = self.get_parameter(test_case, "keyColumns", "[]")
         key_columns: List[str] = literal_eval(key_columns_param)
         if key_columns:
@@ -168,13 +173,13 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
                 "Could not find primary key or unique constraint columns.\n",
                 "Specify 'keyColumns' to explicitly set the columns to use as keys.",
             )
-        return key_columns
+        return set(key_columns)
 
     @staticmethod
     def filter_relevant_columns(
         columns: List[Column],
-        key_columns: List[str],
-        extra_columns: List[str],
+        key_columns: Set[str],
+        extra_columns: Set[str],
         case_sensitive: bool,
     ) -> List[Column]:
         validated_columns = (

metadata/data_quality/validations/table/sqlalchemy/tableCustomSQLQuery.py

@@ -37,7 +37,13 @@ class TableCustomSQLQueryValidator(BaseTableCustomSQLQueryValidator, SQAValidato
                 text(sql_expression)
             )
             if strategy == Strategy.COUNT:
-                return cursor.scalar()
+                result = cursor.scalar()
+                if not isinstance(result, int):
+                    raise ValueError(
+                        f"When using COUNT strategy, the result must be an integer. Received: {type(result)}\n"
+                        "Example: SELECT COUNT(*) FROM table_name WHERE my_value IS NOT NULL"
+                    )
+                return result
             return cursor.fetchall()
         except Exception as exc:
             self.runner._session.rollback()  # pylint: disable=protected-access

metadata/data_quality/validations/table/sqlalchemy/tableDiff.py

@@ -48,6 +48,7 @@ from metadata.generated.schema.tests.basic import (
     TestCaseStatus,
     TestResultValue,
 )
+from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.orm.converter.base import build_orm_col
 from metadata.profiler.orm.functions.md5 import MD5
 from metadata.profiler.orm.functions.substr import Substr
@@ -178,7 +179,7 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
     runtime_params: TableDiffRuntimeParameters
 
     def run_validation(self) -> TestCaseResult:
-        self.runtime_params = self.get_runtime_params()
+        self.runtime_params = self.get_runtime_parameters(TableDiffRuntimeParameters)
         try:
             self._validate_dialects()
             return self._run()
@@ -273,15 +274,16 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
         ).with_schema()
         result = []
         for column in table1.key_columns + table1.extra_columns:
-            col1_type = self._get_column_python_type(
-                table1._schema[column]  # pylint: disable=protected-access
-            )
-            # Skip columns that are not in the second table. We cover this case in get_changed_added_columns.
-            if table2._schema.get(column) is None:  # pylint: disable=protected-access
+            col1 = table1._schema.get(column)  # pylint: disable=protected-access
+            if col1 is None:
+                # Skip columns that are not in the first table. We cover this case in get_changed_added_columns.
                 continue
-            col2_type = self._get_column_python_type(
-                table2._schema[column]  # pylint: disable=protected-access
-            )
+            col2 = table2._schema.get(column)  # pylint: disable=protected-access
+            if col2 is None:
+                # Skip columns that are not in the second table. We cover this case in get_changed_added_columns.
+                continue
+            col1_type = self._get_column_python_type(col1)
+            col2_type = self._get_column_python_type(col2)
             if is_numeric(col1_type) and is_numeric(col2_type):
                 continue
             if col1_type != col2_type:
@@ -429,7 +431,7 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
             self.runtime_params.table_profile_config.profileSampleType
             == ProfileSampleType.ROWS
         ):
-            row_count = self.get_row_count()
+            row_count = self.get_total_row_count()
             if row_count is None:
                 raise ValueError("Row count is required for ROWS profile sample type")
             return int(
@@ -438,13 +440,6 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
             )
         raise ValueError("Invalid profile sample type")
 
-    def get_runtime_params(self) -> TableDiffRuntimeParameters:
-        raw = self.get_test_case_param_value(
-            self.test_case.parameterValues, "runtimeParams", str
-        )
-        runtime_params = TableDiffRuntimeParameters.model_validate_json(raw)
-        return runtime_params
-
     def get_row_diff_test_case_result(
         self,
         threshold: int,
@@ -640,5 +635,13 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
         )
 
     def get_row_count(self) -> Optional[int]:
-        self.runner._sample = None  # pylint: disable=protected-access
         return self._compute_row_count(self.runner, None)
+
+    def get_total_row_count(self) -> Optional[int]:
+        row_count = Metrics.ROW_COUNT()
+        try:
+            row = self.runner.select_first_from_table(row_count.fn())
+            return dict(row).get(Metrics.ROW_COUNT.name)
+        except Exception as e:
+            logger.error(f"Error getting row count: {e}")
+            return None