Add Decimal support to min_max generator (#1013)

[Jgprog117]

Summary

Resolves #1013

• Extended dq_generate_min_max method to support Python's Decimal type in addition to int and float for min/max validation checks.
• Extended check functions to support Decimal values for limits
• Extended serialization and deserialization of checks to support Decimal during round trip

Problem

The current implementation of dq_generate_min_max only recognizes int and float as numeric types:

def _is_num(value):
    return isinstance(value, (int, float))

Changes

• Import Decimal from decimal module
• Update _is_num() to include Decimal in isinstance check
• Add comprehensive unit tests for Decimal, int, and float types

This enables proper data quality checks for decimal-precise financial and scientific data where floating-point precision issues would cause false positives.

Linked issues

Resolves #1013

Tests

• manually tested
• added unit tests
• added integration tests
• added end-to-end tests
• added performance tests

[github-actions]

All commits in PR should be signed ('git commit -S ...'). See https://docs.github.com/en/authentication/managing-commit-signature-verification/signing-commits

[Copilot]

Pull request overview

Extends dq_generate_min_max to recognize Python Decimal values as numeric inputs and adds unit tests covering Decimal min/max scenarios.

Changes:

• Added Decimal to numeric-type detection in dq_generate_min_max.
• Added unit tests validating min/max rule generation for Decimal, int, and float.
• Added a test intended to ensure mixed numeric families (e.g., int + Decimal) don’t generate a rule.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
src/databricks/labs/dqx/profiler/generator.py	Treats Decimal as a supported numeric type for min/max rule generation.
tests/unit/test_generator_numeric.py	Adds unit tests for Decimal min/max behavior and mixed-type handling.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[ghanse]

Looking good. Requesting some extra tests.

[ghanse]

I believe we should update the valid argument types for our check functions if we intend to create rules with Decimal arguments.

One thing to consider is that rules generated from profiling DecimalType columns cannot be saved/loaded with Decimal arguments using the current methods.

We should check that the generated checks can be safely saved/loaded without value modification. Profiling to generating to saving rules is a very common workflow.

[mwojtyczka]

We need to also add support for Decimal in the checks.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[mwojtyczka]

LGTM

Updated tests and missing handling of decimal for tolerance

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 18 out of 18 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

src/databricks/labs/dqx/checks_serializer.py:453

• The PR description/title focus on adding Decimal support to dq_generate_min_max, but this change set also introduces a substantial checks serialization/storage API refactor (new SerializerFactory/DataFrameConverter/Checks* classes and updated call sites). Please update the PR description to reflect the broader scope, or consider splitting the refactor into a separate PR to make review and release-impact assessment clearer.

class ChecksNormalizer:
    """
    Handles normalization and denormalization of check dictionaries.
    E.g. responsible for converting Decimal values to/from serializable format.
    """

    @staticmethod
    def normalize(checks: list[dict]) -> list[dict]:
        """
        Recursively normalize checks dictionary to make it JSON/YAML serializable.

        Args:
            checks: List of check dictionaries that may contain non-serializable values.

        Returns:
            List of normalized check dictionaries.
        """

        def normalize_value(val: Any) -> Any:
            """Recursively normalize a value."""
            if isinstance(val, dict):
                return {k: normalize_value(v) for k, v in val.items()}
            # normalize_bound_args handles None, primitives, lists, tuples, Decimal, etc.
            return normalize_bound_args(val)

        return [normalize_value(check) for check in checks]

    @staticmethod
    def denormalize_value(val: Any) -> Any:
        """Recursively convert special markers (e.g. Decimal) back to original objects."""
        if isinstance(val, dict):
            # Check if this is a Decimal marker
            if "__decimal__" in val and len(val) == 1:
                return Decimal(val["__decimal__"])
            # Otherwise, recursively process the dict
            return {k: ChecksNormalizer.denormalize_value(v) for k, v in val.items()}
        if isinstance(val, (list, tuple)):
            return type(val)(ChecksNormalizer.denormalize_value(v) for v in val)
        return val

    @staticmethod
    def denormalize(checks: list[dict]) -> list[dict]:
        """
        Recursively convert special markers back to objects after deserialization.
        Converts {"__decimal__": "0.01"} back to Decimal("0.01").

        Args:
            checks: List of check dictionaries that may contain special markers.

        Returns:
            List of check dictionaries with special markers converted to objects.
        """
        return [ChecksNormalizer.denormalize_value(check) for check in checks]


class FileFormatSerializer(ABC):
    """
    Abstract base class for file format serializers.
    """

    @abstractmethod
    def serialize(self, data: list[dict]) -> str:
        """Serialize data to string format."""

    @abstractmethod
    def deserialize(self, file_like: TextIO) -> list[dict]:
        """Deserialize data from file-like object."""


class JsonSerializer(FileFormatSerializer):
    """JSON format serializer implementation."""

    def serialize(self, data: list[dict]) -> str:
        """Serialize data to JSON string."""
        return json.dumps(data)

    def deserialize(self, file_like: TextIO) -> list[dict]:
        """Deserialize data from JSON file."""
        return json.load(file_like) or []


class YamlSerializer(FileFormatSerializer):
    """YAML format serializer implementation."""

    def serialize(self, data: list[dict]) -> str:
        """Serialize data to YAML string."""
        return yaml.safe_dump(data)

    def deserialize(self, file_like: TextIO) -> list[dict]:
        """Deserialize data from YAML file."""
        return yaml.safe_load(file_like) or []


class SerializerFactory:
    """
    Factory for creating appropriate serializers based on file extension.
    """

    _serializers: dict[str, type[FileFormatSerializer]] = {
        ".json": JsonSerializer,
        ".yaml": YamlSerializer,
        ".yml": YamlSerializer,
    }

    @classmethod
    def get_supported_extensions(cls) -> tuple[str, ...]:
        """
        Get tuple of supported file extensions.

        Returns:
            Tuple of supported file extensions (e.g., (".json", ".yaml", ".yml")).
        """
        return tuple(cls._serializers.keys())

    @classmethod
    def create_serializer(cls, extension: str | None = None) -> FileFormatSerializer:
        """
        Create a serializer based on file extension.

        Args:
            extension: File extension (e.g., ".json", ".yaml", ".yml").
                       If None or empty, defaults to YAML.

        Returns:
            Appropriate serializer instance. Defaults to YAML if extension not recognized or not provided.
        """
        if not extension:
            return YamlSerializer()
        ext = extension.lower()
        serializer_class = cls._serializers.get(ext, YamlSerializer)
        return serializer_class()

    @classmethod
    def register_format(cls, extension: str, serializer_class: type[FileFormatSerializer]) -> None:
        """
        Register a new file format serializer.

        Args:
            extension: File extension
            serializer_class: Serializer class implementing FileFormatSerializer interface.
        """
        cls._serializers[extension.lower()] = serializer_class


class ChecksSerializer:
    """
    Handles serialization of DQRule objects to dictionaries and file formats.
    """

    @staticmethod
    def serialize(checks: list[DQRule]) -> list[dict]:
        """
        Converts a list of quality checks defined as *DQRule* objects to a list of quality checks
        defined as Python dictionaries.

        Args:
            checks: List of DQRule instances to convert.

        Returns:
            List of dictionaries representing the DQRule instances.

        Raises:
            InvalidCheckError: If any item in the list is not a DQRule instance.
        """
        dq_rules = []
        for check in checks:
            if not isinstance(check, DQRule):
                raise InvalidCheckError(f"Expected DQRule instance, got {type(check).__name__}")
            dq_rules.append(check.to_dict())
        return dq_rules

    @staticmethod
    def serialize_to_bytes(checks: list[dict], extension: str) -> bytes:
        """
        Serializes a list of checks to bytes in json or yaml (default) format.

        Args:
            checks: List of checks to serialize.
            extension: File extension (e.g., ".json", ".yaml", ".yml").
        Returns:
            Serialized checks as bytes.
        """
        serializer = SerializerFactory.create_serializer(extension)
        normalized_checks = ChecksNormalizer.normalize(checks)
        serialized_str = serializer.serialize(normalized_checks)
        return serialized_str.encode("utf-8")


class ChecksDeserializer:
    """
    Handles deserialization of dictionaries to DQRule objects and from file formats.
    """

    def __init__(self, custom_checks: dict[str, Callable] | None = None):
        """
        Initialize the deserializer.

        Args:
            custom_checks: Dictionary with custom check functions.
        """
        self.custom_checks = custom_checks

    def deserialize(self, checks: list[dict]) -> list[DQRule]:
        """
        Converts a list of quality checks defined as Python dictionaries to a list of `DQRule` objects.

        Args:
            checks: list of dictionaries describing checks. Each check is a dictionary
                consisting of following fields:
                - *check* - Column expression to evaluate. This expression should return string value if it's evaluated to true
                    or *null* if it's evaluated to *false*
                - *name* - name that will be given to a resulting column. Autogenerated if not provided
                - *criticality* (optional) - possible values are *error* (data going only into "bad" dataframe),
                and *warn* (data is going into both dataframes)
                - *filter* (optional) - Expression for filtering data quality checks
                - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check.

        Returns:
            list of data quality check rules

        Raises:
            InvalidCheckError: If any dictionary is invalid or unsupported.
        """
        status = ChecksValidator.validate_checks(checks, self.custom_checks)
        if status.has_errors:
            raise InvalidCheckError(str(status))

        dq_rule_checks: list[DQRule] = []
        for check_def in checks:
            logger.debug(f"Processing check definition: {check_def}")

            check = check_def.get("check", {})
            name = check_def.get("name", None)
            func_name = check.get("function")
            func = resolve_check_function(func_name, self.custom_checks, fail_on_missing=True)
            assert func  # should already be validated

            func_args = check.get("arguments", {})
            for_each_column = check.get("for_each_column")
            column = func_args.get("column")  # should be defined for single-column checks only
            columns = func_args.get("columns")  # should be defined for multi-column checks only
            assert not (column and columns)  # should already be validated
            criticality = check_def.get("criticality", "error")
            filter_str = check_def.get("filter")
            user_metadata = check_def.get("user_metadata")

            # Exclude `column` and `columns` from check_func_kwargs
            # as these are always included in the check function call
            check_func_kwargs = {k: v for k, v in func_args.items() if k not in {"column", "columns"}}

            # treat non-registered function as row-level checks
            if for_each_column:
                dq_rule_checks += DQForEachColRule(
                    columns=for_each_column,
                    name=name,
                    check_func=func,
                    criticality=criticality,
                    filter=filter_str,
                    check_func_kwargs=check_func_kwargs,
                    user_metadata=user_metadata,
                ).get_rules()
            else:
                rule_type = CHECK_FUNC_REGISTRY.get(func_name)
                if rule_type == "dataset":
                    dq_rule_checks.append(
                        DQDatasetRule(
                            column=column,
                            columns=columns,
                            check_func=func,
                            check_func_kwargs=check_func_kwargs,
                            name=name,
                            criticality=criticality,
                            filter=filter_str,
                            user_metadata=user_metadata,
                        )
                    )
                else:  # default to row-level rule
                    dq_rule_checks.append(
                        DQRowRule(
                            column=column,
                            columns=columns,
                            check_func=func,
                            check_func_kwargs=check_func_kwargs,
                            name=name,
                            criticality=criticality,
                            filter=filter_str,
                            user_metadata=user_metadata,
                        )
                    )

        return dq_rule_checks

    @staticmethod
    def deserialize_from_file(extension: str, file_like: TextIO) -> list[dict]:
        """
        Deserialize checks from a file-like object based on file extension.
        Automatically denormalizes special markers back to objects.

        Args:
            extension: File extension (e.g., ".json", ".yaml", ".yml").
            file_like: File-like object to read from.

        Returns:
            List of check dictionaries with special markers converted to objects.
        """
        serializer = SerializerFactory.create_serializer(extension)
        checks = serializer.deserialize(file_like)
        return ChecksNormalizer.denormalize(checks)


class DataFrameConverter:
    """
    Handles conversion between DataFrames and check dictionaries.
    """

    @staticmethod
    def from_dataframe(df: DataFrame, run_config_name: str = "default") -> list[dict]:
        """
        Converts a list of quality checks defined in a DataFrame to a list of quality checks
        defined as Python dictionaries.

        Args:
            df: DataFrame with data quality check rules. Each row should define a check. Rows should
            have the following columns:
                - *name* - Name that will be given to a resulting column. Autogenerated if not provided.
                - *criticality* (optional) - Possible values are *error* (data going only into "bad" dataframe) and *warn* (data is going into both dataframes).
                - *check* - DQX check function used in the check; A *StructType* column defining the data quality check.
                - *filter* - Expression for filtering data quality checks.
                - *run_config_name* (optional) - Run configuration name for storing checks across runs.
                - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check.
            run_config_name: Run configuration name for filtering quality rules, e.g. input table or job name (use "default" if not provided).

        Returns:
                List of data quality check specifications as a Python dictionary
        """
        check_rows = df.where(f"run_config_name = '{run_config_name}'").collect()
        collect_limit = 500
        if len(check_rows) > collect_limit:
            warnings.warn(
                f"Collecting large number of rows from DataFrame: {len(check_rows)}",
                category=UserWarning,
                stacklevel=2,
            )

        checks = []
        for row in check_rows:
            check_dict = {
                "name": row.name,
                "criticality": row.criticality,
                "check": {
                    "function": row.check["function"],
                    "arguments": (
                        {k: safe_json_load(v) for k, v in row.check["arguments"].items()}
                        if row.check["arguments"] is not None
                        else {}
                    ),
                },
            }
            if "for_each_column" in row.check and row.check["for_each_column"]:
                check_dict["check"]["for_each_column"] = row.check["for_each_column"]
            if row.filter is not None:
                check_dict["filter"] = row.filter
            if row.user_metadata is not None:
                check_dict["user_metadata"] = row.user_metadata
            # Denormalize special markers back to objects
            checks.append(ChecksNormalizer.denormalize_value(check_dict))
        return checks

    @staticmethod
    def to_dataframe(
        spark: SparkSession,
        checks: list[dict],
        run_config_name: str = "default",
    ) -> DataFrame:
        """
        Converts a list of quality checks defined as Python dictionaries to a DataFrame.

        Args:
            spark: Spark session.
            checks: list of check specifications as Python dictionaries. Each check consists of the following fields:
                - *check* - Column expression to evaluate. This expression should return string value if it's evaluated to
                   true (it will be used as an error/warning message) or *null* if it's evaluated to *false*
                - *name* - Name that will be given to a resulting column. Autogenerated if not provided
                - *criticality* (optional) - Possible values are *error* (data going only into "bad" dataframe) and *warn*
                   (data is going into both dataframes)
                - *filter* (optional) - Expression for filtering data quality checks
                - *user_metadata* (optional) - User-defined key-value pairs added to metadata generated by the check.
            run_config_name: Run configuration name for storing quality checks across runs, e.g. input table or job name (use "default" if not provided)

        Returns:
            DataFrame with data quality check rules

        Raises:
            InvalidCheckError: If any check is invalid or unsupported.
        """
        deserializer = ChecksDeserializer()
        dq_rule_checks: list[DQRule] = deserializer.deserialize(checks)

        dq_rule_rows = []
        for dq_rule_check in dq_rule_checks:
            arguments = dict(dq_rule_check.check_func_kwargs)

            if dq_rule_check.column is not None:
                arguments["column"] = dq_rule_check.column

            if dq_rule_check.columns is not None:
                arguments["columns"] = dq_rule_check.columns

            json_arguments = {k: json.dumps(normalize_bound_args(v)) for k, v in arguments.items()}
            dq_rule_rows.append(
                [
                    dq_rule_check.name,
                    dq_rule_check.criticality,
                    {"function": dq_rule_check.check_func.__name__, "arguments": json_arguments},
                    dq_rule_check.filter,
                    run_config_name,
                    dq_rule_check.user_metadata,
                ]
            )
        return spark.createDataFrame(dq_rule_rows, CHECKS_TABLE_SCHEMA)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.