openmetadata-ingestion 1.5.13.0__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

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/models.py

@@ -4,13 +4,17 @@ from typing import List, Optional
 
 from pydantic import BaseModel
 
-from metadata.generated.schema.entity.data.table import Column
+from metadata.generated.schema.entity.data.table import Column, TableProfilerConfig
+from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseServiceType,
+)
 
 
 class TableParameter(BaseModel):
     serviceUrl: str
     path: str
     columns: List[Column]
+    database_service_type: DatabaseServiceType
 
 
 class TableDiffRuntimeParameters(BaseModel):
@@ -19,3 +23,4 @@ class TableDiffRuntimeParameters(BaseModel):
     keyColumns: List[str]
     extraColumns: List[str]
     whereClause: Optional[str]
+    table_profile_config: Optional[TableProfilerConfig]

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.sqlalchemy.sampler import SQASampler
+
+
+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: SQASampler,
+    ) -> 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,9 +75,13 @@ 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(
+                database_service_type=service1.serviceType,
                 path=self.get_data_diff_table_path(
                     self.table_entity.fullyQualifiedName.root
                 ),
@@ -94,6 +98,7 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
                 ),
             ),
             table2=TableParameter(
+                database_service_type=service2.serviceType,
                 path=self.get_data_diff_table_path(table2_fqn),
                 serviceUrl=self.get_data_diff_url(
                     service2,
@@ -108,18 +113,19 @@ 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
-            or not self.sampler._partition_details.enablePartitioning
+            if not (
+                self.sampler.partition_details
+                and self.sampler.partition_details.enablePartitioning
+            )
             else self.sampler.get_partitioned_query().whereclause.compile(
                 compile_kwargs={"literal_binds": True}
             )
@@ -130,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:
@@ -163,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/tableDiff.py

@@ -10,8 +10,11 @@
 #  limitations under the License.
 # pylint: disable=missing-module-docstring
 import logging
+import random
+import string
 import traceback
 from decimal import Decimal
+from functools import reduce
 from itertools import islice
 from typing import Dict, Iterable, List, Optional, Tuple, cast
 from urllib.parse import urlparse
@@ -22,23 +25,35 @@ from data_diff.diff_tables import DiffResultWrapper
 from data_diff.errors import DataDiffMismatchingKeyTypesError
 from data_diff.utils import ArithAlphanumeric, CaseInsensitiveDict
 from sqlalchemy import Column as SAColumn
+from sqlalchemy import create_engine, literal, select
 
 from metadata.data_quality.validations import utils
 from metadata.data_quality.validations.base_test_handler import BaseTestValidator
 from metadata.data_quality.validations.mixins.sqa_validator_mixin import (
     SQAValidatorMixin,
 )
-from metadata.data_quality.validations.models import TableDiffRuntimeParameters
-from metadata.generated.schema.entity.data.table import Column
+from metadata.data_quality.validations.models import (
+    TableDiffRuntimeParameters,
+    TableParameter,
+)
+from metadata.generated.schema.entity.data.table import Column, ProfileSampleType
 from metadata.generated.schema.entity.services.connections.database.sapHanaConnection import (
     SapHanaScheme,
 )
+from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseServiceType,
+)
 from metadata.generated.schema.tests.basic import (
     TestCaseResult,
     TestCaseStatus,
     TestResultValue,
 )
-from metadata.profiler.orm.registry import Dialects
+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
+from metadata.profiler.orm.registry import Dialects, PythonDialects
+from metadata.utils.collections import CaseInsensitiveList
 from metadata.utils.logger import test_suite_logger
 
 logger = test_suite_logger()
@@ -57,6 +72,73 @@ SUPPORTED_DIALECTS = [
 ]
 
 
+def build_sample_where_clause(
+    table: TableParameter, key_columns: List[str], salt: str, hex_nounce: str
+) -> str:
+    sql_alchemy_columns = [
+        build_orm_col(i, c, table.database_service_type)
+        for i, c in enumerate(table.columns)
+        if c.name.root in key_columns
+    ]
+    reduced_concat = reduce(
+        lambda c1, c2: c1.concat(c2), sql_alchemy_columns + [literal(salt)]
+    )
+    sqa_dialect = create_engine(
+        f"{PythonDialects[table.database_service_type.name].value}://"
+    ).dialect
+    return str(
+        select()
+        .filter(
+            Substr(
+                MD5(reduced_concat),
+                1,
+                8,
+            )
+            < hex_nounce
+        )
+        .whereclause.compile(
+            dialect=sqa_dialect,
+            compile_kwargs={"literal_binds": True},
+        )
+    )
+
+
+def compile_and_clauses(elements) -> str:
+    """Compile a list of elements into a string with 'and' clauses.
+
+    Args:
+        elements: A string or a list of strings or lists
+
+    Returns:
+        A string with 'and' clauses
+
+    Raises:
+        ValueError: If the input is not a string or a list
+
+    Examples:
+        >>> compile_and_clauses("a")
+        'a'
+        >>> compile_and_clauses(["a", "b"])
+        'a and b'
+        >>> compile_and_clauses([["a", "b"], "c"])
+        '(a and b) and c'
+    """
+    if isinstance(elements, str):
+        return elements
+    if isinstance(elements, list):
+        if len(elements) == 1:
+            return compile_and_clauses(elements[0])
+        return " and ".join(
+            (
+                f"({compile_and_clauses(e)})"
+                if isinstance(e, list)
+                else compile_and_clauses(e)
+            )
+            for e in elements
+        )
+    raise ValueError("Input must be a string or a list")
+
+
 class UnsupportedDialectError(Exception):
     def __init__(self, param: str, dialect: str):
         super().__init__(f"Unsupported dialect in param {param}: {dialect}")
@@ -97,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()
@@ -192,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:
@@ -241,17 +324,20 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
 
     def get_table_diff(self) -> DiffResultWrapper:
         """Calls data_diff.diff_tables with the parameters from the test case."""
+        left_where, right_where = self.sample_where_clause()
         table1 = data_diff.connect_to_table(
             self.runtime_params.table1.serviceUrl,
             self.runtime_params.table1.path,
             self.runtime_params.keyColumns,  # type: ignore
             case_sensitive=self.get_case_sensitive(),
+            where=left_where,
         )
         table2 = data_diff.connect_to_table(
             self.runtime_params.table2.serviceUrl,
             self.runtime_params.table2.path,
             self.runtime_params.keyColumns,  # type: ignore
             case_sensitive=self.get_case_sensitive(),
+            where=right_where,
         )
         data_diff_kwargs = {
             "key_columns": self.runtime_params.keyColumns,
@@ -270,12 +356,89 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
         """Returns the where clause from the test case parameters or None if it is a blank string."""
         return self.runtime_params.whereClause or None
 
-    def get_runtime_params(self) -> TableDiffRuntimeParameters:
-        raw = self.get_test_case_param_value(
-            self.test_case.parameterValues, "runtimeParams", str
+    def sample_where_clause(self) -> Tuple[Optional[str], Optional[str]]:
+        """We use a where clause to sample the data for the diff. This is useful because with data diff
+        we do not have access to the underlying 'SELECT' statement. This method generates a where clause
+        that selects a random sample of the data based on the profile sample configuration.
+        The method uses the md5 hash of the key columns and a random salt to select a random sample of the data.
+        This ensures that the same data is selected from the two tables of the comparison.
+
+        Example:
+            -- Table 1 --  | -- Table 2 --
+            id | name      | id | name
+            1  | Alice     | 1  | Alice
+            2  | Bob       | 2  | Bob
+            3  | Charlie   | 3  | Charlie
+            4  | David     | 4  | Edward
+            5  | Edward    | 6  | Frank
+
+        If we want a sample of 20% of the data, the where clause will intend to select one of the rows
+        on Table 1 and the hash will ensure that the same row is selected on Table 2. We want to avoid selecting rows
+        with different ids because the comparison will not be sensible.
+        """
+        if (
+            # no sample configuration
+            self.runtime_params.table_profile_config is None
+            or self.runtime_params.table_profile_config.profileSample is None
+            # sample is 100% or in other words no sample is required
+            or (
+                self.runtime_params.table_profile_config.profileSampleType
+                == ProfileSampleType.PERCENTAGE
+                and self.runtime_params.table_profile_config.profileSample == 100
+            )
+        ):
+            return None, None
+        if DatabaseServiceType.Mssql in [
+            self.runtime_params.table1.database_service_type,
+            self.runtime_params.table2.database_service_type,
+        ]:
+            logger.warning("Sampling not supported in MSSQL. Skipping sampling.")
+            return None, None
+        nounce = self.calculate_nounce()
+        # SQL MD5 returns a 32 character hex string even with leading zeros so we need to
+        # pad the nounce to 8 characters in preserve lexical order.
+        # example: SELECT md5('j(R1wzR*y[^GxWJ5B>L{-HLETRD');
+        hex_nounce = hex(nounce)[2:].rjust(8, "0")
+        # TODO: using strings for this is sub-optimal. But using bytes buffers requires a by-database
+        # implementaiton. We can use this as default and add database specific implementations as the
+        # need arises.
+        salt = "".join(
+            random.choices(string.ascii_letters + string.digits, k=5)
+        )  # 1 / ~62^5 should be enough entropy. Use letters and digits to avoid messing with SQL syntax
+        key_columns = (
+            CaseInsensitiveList(self.runtime_params.keyColumns)
+            if not self.get_case_sensitive()
+            else self.runtime_params.keyColumns
+        )
+        return tuple(
+            build_sample_where_clause(table, key_columns, salt, hex_nounce)
+            for table in [self.runtime_params.table1, self.runtime_params.table2]
         )
-        runtime_params = TableDiffRuntimeParameters.model_validate_json(raw)
-        return runtime_params
+
+    def calculate_nounce(self, max_nounce=2**32 - 1) -> int:
+        """Calculate the nounce based on the profile sample configuration. The nounce is
+        the sample fraction projected to a number on a scale of 0 to max_nounce"""
+        if (
+            self.runtime_params.table_profile_config.profileSampleType
+            == ProfileSampleType.PERCENTAGE
+        ):
+            return int(
+                max_nounce
+                * self.runtime_params.table_profile_config.profileSample
+                / 100
+            )
+        if (
+            self.runtime_params.table_profile_config.profileSampleType
+            == ProfileSampleType.ROWS
+        ):
+            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(
+                max_nounce
+                * (self.runtime_params.table_profile_config.profileSample / row_count)
+            )
+        raise ValueError("Invalid profile sample type")
 
     def get_row_diff_test_case_result(
         self,
@@ -470,3 +633,15 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
         return utils.get_bool_test_case_param(
             self.test_case.parameterValues, "caseSensitiveColumns"
         )
+
+    def get_row_count(self) -> Optional[int]:
+        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