snowpark-checkpoints-validators 0.2.0rc1__py3-none-any.whl → 0.3.0__py3-none-any.whl

snowflake/snowpark_checkpoints/spark_migration.py

@@ -0,0 +1,255 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from typing import Callable, Optional, TypeVar
+
+import pandas as pd
+
+from pyspark.sql import DataFrame as SparkDataFrame
+
+from snowflake.snowpark import DataFrame as SnowparkDataFrame
+from snowflake.snowpark.types import PandasDataFrame
+from snowflake.snowpark_checkpoints.errors import SparkMigrationError
+from snowflake.snowpark_checkpoints.job_context import SnowparkJobContext
+from snowflake.snowpark_checkpoints.snowpark_sampler import (
+    SamplingAdapter,
+    SamplingStrategy,
+)
+from snowflake.snowpark_checkpoints.utils.constants import FAIL_STATUS, PASS_STATUS
+from snowflake.snowpark_checkpoints.utils.logging_utils import log
+from snowflake.snowpark_checkpoints.utils.telemetry import STATUS_KEY, report_telemetry
+from snowflake.snowpark_checkpoints.utils.utils_checks import (
+    _replace_special_characters,
+    _update_validation_result,
+)
+
+
+fn = TypeVar("F", bound=Callable)
+LOGGER = logging.getLogger(__name__)
+
+
+@log
+def check_with_spark(
+    job_context: Optional[SnowparkJobContext],
+    spark_function: fn,
+    checkpoint_name: str,
+    sample_number: Optional[int] = 100,
+    sampling_strategy: Optional[SamplingStrategy] = SamplingStrategy.RANDOM_SAMPLE,
+    output_path: Optional[str] = None,
+) -> Callable[[fn], fn]:
+    """Validate function output with Spark instance.
+
+    Will take the input snowpark dataframe of this function, sample data, convert
+    it to a Spark dataframe and then execute `spark_function`. Subsequently
+    the output of that function will be compared to the output of this function
+    for the same sample of data.
+
+    Args:
+        job_context (SnowparkJobContext): The job context containing configuration and details for the validation.
+        spark_function (fn): The equivalent PySpark function to compare against the Snowpark implementation.
+        checkpoint_name (str): A name for the checkpoint. Defaults to None.
+        sample_number (Optional[int], optional): The number of rows for validation. Defaults to 100.
+        sampling_strategy (Optional[SamplingStrategy], optional): The strategy used for sampling data.
+            Defaults to SamplingStrategy.RANDOM_SAMPLE.
+        output_path (Optional[str], optional): The path to store the validation results. Defaults to None.
+
+    Returns:
+        Callable[[fn], fn]: A decorator that wraps the original Snowpark function with validation logic.
+
+    """
+
+    def check_with_spark_decorator(snowpark_fn):
+        @log(log_args=False)
+        def wrapper(*args, **kwargs):
+            LOGGER.info(
+                "Starting output validation between Snowpark function '%s' and Spark function '%s'",
+                snowpark_fn.__name__,
+                spark_function.__name__,
+            )
+            _checkpoint_name = checkpoint_name
+            if checkpoint_name is None:
+                LOGGER.warning(
+                    "No checkpoint name provided. Using '%s' as the checkpoint name",
+                    snowpark_fn.__name__,
+                )
+                _checkpoint_name = snowpark_fn.__name__
+            _checkpoint_name = _replace_special_characters(_checkpoint_name)
+
+            sampler = SamplingAdapter(
+                job_context,
+                sample_number=sample_number,
+                sampling_strategy=sampling_strategy,
+            )
+            sampler.process_args(args)
+            snowpark_sample_args = sampler.get_sampled_snowpark_args()
+            pyspark_sample_args = sampler.get_sampled_spark_args()
+
+            # Run the sampled data in snowpark
+            LOGGER.info("Running the Snowpark function with sampled args")
+            snowpark_test_results = snowpark_fn(*snowpark_sample_args, **kwargs)
+            LOGGER.info("Running the Spark function with sampled args")
+            spark_test_results = spark_function(*pyspark_sample_args, **kwargs)
+
+            LOGGER.info("Comparing the results of the Snowpark and Spark functions")
+            result, exception = _assert_return(
+                snowpark_test_results,
+                spark_test_results,
+                job_context,
+                _checkpoint_name,
+                output_path,
+            )
+            if not result:
+                LOGGER.error(
+                    "Validation failed. The results of the Snowpark function '%s' and Spark function '%s' do not match",
+                    snowpark_fn.__name__,
+                    spark_function.__name__,
+                )
+                raise exception from None
+            LOGGER.info(
+                "Validation passed. The results of the Snowpark function '%s' and Spark function '%s' match",
+                snowpark_fn.__name__,
+                spark_function.__name__,
+            )
+
+            # Run the original function in snowpark
+            return snowpark_fn(*args, **kwargs)
+
+        return wrapper
+
+    return check_with_spark_decorator
+
+
+@report_telemetry(
+    params_list=["snowpark_results", "spark_results"],
+    return_indexes=[(STATUS_KEY, 0)],
+    multiple_return=True,
+)
+def _assert_return(
+    snowpark_results, spark_results, job_context, checkpoint_name, output_path=None
+) -> tuple[bool, Optional[Exception]]:
+    """Assert and validate the results from Snowpark and Spark transformations.
+
+    Args:
+        snowpark_results (Any): Results from the Snowpark transformation.
+        spark_results (Any): Results from the Spark transformation to compare against.
+        job_context (Any): Additional context about the job. Defaults to None.
+        checkpoint_name (Any): Name of the checkpoint for logging. Defaults to None.
+        output_path (Optional[str], optional): The path to store the validation results. Defaults to None.
+
+    Raises:
+        AssertionError: If the Snowpark and Spark results do not match.
+        TypeError: If the results cannot be compared.
+
+    """
+    if isinstance(snowpark_results, SnowparkDataFrame) and isinstance(
+        spark_results, SparkDataFrame
+    ):
+        LOGGER.debug("Comparing two DataFrame results for equality")
+        cmp = compare_spark_snowpark_dfs(spark_results, snowpark_results)
+
+        if not cmp.empty:
+            exception_result = SparkMigrationError(
+                "DataFrame difference:\n", job_context, checkpoint_name, cmp
+            )
+            return False, exception_result
+        job_context._mark_pass(checkpoint_name)
+        _update_validation_result(checkpoint_name, PASS_STATUS, output_path)
+        return True, None
+    else:
+        LOGGER.debug("Comparing two scalar results for equality")
+        if snowpark_results != spark_results:
+            exception_result = SparkMigrationError(
+                "Return value difference:\n",
+                job_context,
+                checkpoint_name,
+                f"{snowpark_results} != {spark_results}",
+            )
+            _update_validation_result(checkpoint_name, FAIL_STATUS, output_path)
+            return False, exception_result
+        job_context._mark_pass(checkpoint_name)
+        _update_validation_result(checkpoint_name, PASS_STATUS, output_path)
+        return True, None
+
+
+def compare_spark_snowpark_dfs(
+    spark_df: SparkDataFrame, snowpark_df: SnowparkDataFrame
+) -> PandasDataFrame:
+    """Compare two dataframes for equality.
+
+    Args:
+        spark_df (SparkDataFrame): The Spark dataframe to compare.
+        snowpark_df (SnowparkDataFrame): The Snowpark dataframe to compare.
+
+    Returns:Pandas DataFrame containing the differences between the two dataframes.
+
+    """
+    snowpark_df = snowpark_df.to_pandas()
+    snowpark_df.columns = snowpark_df.columns.str.upper()
+    spark_df = spark_df.toPandas()
+    spark_df.columns = spark_df.columns.str.upper()
+    spark_cols = set(spark_df.columns)
+    snowpark_cols = set(snowpark_df.columns)
+    cmp = pd.DataFrame([])
+    left = spark_cols - snowpark_cols
+    right = snowpark_cols - spark_cols
+    if left != set():
+        cmp = _compare_dfs(spark_df, snowpark_df, "spark", "snowpark")
+    if right != set():
+        right_cmp = _compare_dfs(snowpark_df, spark_df, "snowpark", "spark")
+        cmp = right_cmp if cmp.empty else pd.concat([cmp, right_cmp], ignore_index=True)
+    if left == set() and right == set():
+        if spark_df.shape == snowpark_df.shape:
+            cmp = spark_df.compare(snowpark_df, result_names=("spark", "snowpark"))
+        else:
+            cmp = spark_df.merge(snowpark_df, indicator=True, how="outer").loc[
+                lambda x: x["_merge"] != "both"
+            ]
+            cmp = cmp.replace(
+                {"left_only": "spark_only", "right_only": "snowpark_only"}
+            )
+
+    return cmp
+
+
+def _compare_dfs(
+    df_a: pd.DataFrame, df_b: pd.DataFrame, left_label: str, right_label: str
+) -> PandasDataFrame:
+    """Compare two dataframes for equality.
+
+    Args:
+        df_a (PandasDataFrame): The first dataframe to compare.
+        df_b (PandasDataFrame): The second dataframe to compare.
+        left_label (str): The label for the first dataframe.
+        right_label (str): The label for the second dataframe.
+
+    :return: Pandas DataFrame containing the differences between the two dataframes.
+
+    """
+    df_a["side"] = "a"
+    df_b["side"] = "b"
+    a_only = [col for col in df_a.columns if col not in df_b.columns] + ["side"]
+    b_only = [col for col in df_b.columns if col not in df_a.columns] + ["side"]
+    cmp = (
+        df_a[a_only]
+        .merge(df_b[b_only], indicator=True, how="left")
+        .loc[lambda x: x["_merge"] != "both"]
+    )
+    cmp = cmp.replace(
+        {"left_only": f"{left_label}_only", "right_only": f"{right_label}_only"}
+    )
+    cmp = cmp.drop(columns="side")
+    return cmp

snowflake/snowpark_checkpoints/utils/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

snowflake/snowpark_checkpoints/utils/constants.py

@@ -0,0 +1,134 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Skip type
+from enum import IntEnum
+from typing import Final
+
+
+class CheckpointMode(IntEnum):
+
+    """Enum class representing the validation mode."""
+
+    SCHEMA = 1
+    """Validate against a schema file"""
+    DATAFRAME = 2
+    """Validate against a dataframe"""
+
+
+# Execution mode
+SCHEMA_EXECUTION_MODE: Final[str] = "Schema"
+DATAFRAME_EXECUTION_MODE: Final[str] = "Dataframe"
+
+
+# File position on stack
+STACK_POSITION_CHECKPOINT: Final[int] = 6
+
+# Validation status
+PASS_STATUS: Final[str] = "PASS"
+FAIL_STATUS: Final[str] = "FAIL"
+
+# Validation result keys
+DEFAULT_KEY: Final[str] = "default"
+
+
+# Skip type
+SKIP_ALL: Final[str] = "skip_all"
+
+# Supported types
+BOOLEAN_TYPE: Final[str] = "boolean"
+BINARY_TYPE: Final[str] = "binary"
+BYTE_TYPE: Final[str] = "byte"
+CHAR_TYPE: Final[str] = "char"
+DATE_TYPE: Final[str] = "date"
+DAYTIMEINTERVAL_TYPE: Final[str] = "daytimeinterval"
+DECIMAL_TYPE: Final[str] = "decimal"
+DOUBLE_TYPE: Final[str] = "double"
+FLOAT_TYPE: Final[str] = "float"
+INTEGER_TYPE: Final[str] = "integer"
+LONG_TYPE: Final[str] = "long"
+SHORT_TYPE: Final[str] = "short"
+STRING_TYPE: Final[str] = "string"
+TIMESTAMP_TYPE: Final[str] = "timestamp"
+TIMESTAMP_NTZ_TYPE: Final[str] = "timestamp_ntz"
+VARCHAR_TYPE: Final[str] = "varchar"
+
+# Pandas data types
+PANDAS_BOOLEAN_DTYPE: Final[str] = "bool"
+PANDAS_DATETIME_DTYPE: Final[str] = "datetime64[ns]"
+PANDAS_FLOAT_DTYPE: Final[str] = "float64"
+PANDAS_INTEGER_DTYPE: Final[str] = "int64"
+PANDAS_OBJECT_DTYPE: Final[str] = "object"
+PANDAS_TIMEDELTA_DTYPE: Final[str] = "timedelta64[ns]"
+
+# Schemas keys
+COLUMNS_KEY: Final[str] = "columns"
+COUNT_KEY: Final[str] = "rows_count"
+DECIMAL_PRECISION_KEY: Final[str] = "decimal_precision"
+FALSE_COUNT_KEY: Final[str] = "false_count"
+FORMAT_KEY: Final[str] = "format"
+NAME_KEY: Final[str] = "name"
+MARGIN_ERROR_KEY: Final[str] = "margin_error"
+MAX_KEY: Final[str] = "max"
+MEAN_KEY: Final[str] = "mean"
+MIN_KEY: Final[str] = "min"
+NULL_COUNT_KEY: Final[str] = "rows_null_count"
+NULLABLE_KEY: Final[str] = "nullable"
+ROWS_NOT_NULL_COUNT_KEY: Final[str] = "rows_not_null_count"
+TRUE_COUNT_KEY: Final[str] = "true_count"
+TYPE_KEY: Final[str] = "type"
+ROWS_COUNT_KEY: Final[str] = "rows_count"
+FORMAT_KEY: Final[str] = "format"
+
+DATAFRAME_CUSTOM_DATA_KEY: Final[str] = "custom_data"
+DATAFRAME_PANDERA_SCHEMA_KEY: Final[str] = "pandera_schema"
+
+# Default values
+DEFAULT_DATE_FORMAT: Final[str] = "%Y-%m-%d"
+
+# SQL Column names
+TABLE_NAME_COL: Final[str] = "TABLE_NAME"
+CREATED_COL: Final[str] = "CREATED"
+
+# SQL Table names
+INFORMATION_SCHEMA_TABLE_NAME: Final[str] = "INFORMATION_SCHEMA"
+TABLES_TABLE_NAME: Final[str] = "TABLES"
+
+# SQL Query
+EXCEPT_HASH_AGG_QUERY: Final[
+    str
+] = "SELECT HASH_AGG(*) FROM IDENTIFIER(:1) EXCEPT SELECT HASH_AGG(*) FROM IDENTIFIER(:2)"
+
+# Table checkpoints name
+CHECKPOINT_TABLE_NAME_FORMAT: Final[str] = "{}_snowpark"
+
+# Write mode
+OVERWRITE_MODE: Final[str] = "overwrite"
+
+# Validation modes
+VALIDATION_MODE_KEY: Final[str] = "validation_mode"
+PIPELINES_KEY: Final[str] = "pipelines"
+
+# File name
+CHECKPOINT_JSON_OUTPUT_FILE_FORMAT_NAME: Final[str] = "{}.json"
+CHECKPOINTS_JSON_FILE_NAME: Final[str] = "checkpoints.json"
+SNOWPARK_CHECKPOINTS_OUTPUT_DIRECTORY_NAME: Final[str] = "snowpark-checkpoints-output"
+CHECKPOINT_PARQUET_OUTPUT_FILE_FORMAT_NAME: Final[str] = "{}.parquet"
+VALIDATION_RESULTS_JSON_FILE_NAME: Final[str] = "checkpoint_validation_results.json"
+
+# Environment variables
+SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH_ENV_VAR: Final[
+    str
+] = "SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH"

snowflake/snowpark_checkpoints/utils/extra_config.py

@@ -0,0 +1,132 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import os
+
+from typing import Optional
+
+from snowflake.snowpark_checkpoints.io_utils.io_file_manager import get_io_file_manager
+from snowflake.snowpark_checkpoints.utils.constants import (
+    SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH_ENV_VAR,
+)
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+# noinspection DuplicatedCode
+def _get_checkpoint_contract_file_path() -> str:
+    return os.environ.get(
+        SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH_ENV_VAR, get_io_file_manager().getcwd()
+    )
+
+
+def _set_conf_io_strategy() -> None:
+    try:
+        from snowflake.snowpark_checkpoints.io_utils.io_default_strategy import (
+            IODefaultStrategy,
+        )
+        from snowflake.snowpark_checkpoints_configuration.io_utils.io_file_manager import (
+            EnvStrategy as ConfEnvStrategy,
+        )
+        from snowflake.snowpark_checkpoints_configuration.io_utils.io_file_manager import (
+            get_io_file_manager as get_conf_io_file_manager,
+        )
+
+        is_default_strategy = isinstance(
+            get_io_file_manager().strategy, IODefaultStrategy
+        )
+
+        if is_default_strategy:
+            return
+
+        class CustomConfEnvStrategy(ConfEnvStrategy):
+            def file_exists(self, path: str) -> bool:
+                return get_io_file_manager().file_exists(path)
+
+            def read(
+                self, file_path: str, mode: str = "r", encoding: Optional[str] = None
+            ) -> Optional[str]:
+                return get_io_file_manager().read(file_path, mode, encoding)
+
+            def getcwd(self) -> str:
+                return get_io_file_manager().getcwd()
+
+        get_conf_io_file_manager().set_strategy(CustomConfEnvStrategy())
+
+    except ImportError:
+        LOGGER.debug(
+            "snowpark-checkpoints-configuration is not installed. Cannot get a checkpoint metadata instance."
+        )
+
+
+# noinspection DuplicatedCode
+def _get_metadata():
+    try:
+        from snowflake.snowpark_checkpoints_configuration.checkpoint_metadata import (
+            CheckpointMetadata,
+        )
+
+        path = _get_checkpoint_contract_file_path()
+        _set_conf_io_strategy()
+        LOGGER.debug("Loading checkpoint metadata from '%s'", path)
+        metadata = CheckpointMetadata(path)
+        return True, metadata
+
+    except ImportError:
+        LOGGER.debug(
+            "snowpark-checkpoints-configuration is not installed. Cannot get a checkpoint metadata instance."
+        )
+        return False, None
+
+
+def is_checkpoint_enabled(checkpoint_name: Optional[str] = None) -> bool:
+    """Check if a checkpoint is enabled.
+
+    Args:
+        checkpoint_name (Optional[str], optional): The name of the checkpoint.
+
+    Returns:
+        bool: True if the checkpoint is enabled, False otherwise.
+
+    """
+    enabled, metadata = _get_metadata()
+    if enabled and checkpoint_name is not None:
+        config = metadata.get_checkpoint(checkpoint_name)
+        return config.enabled
+    return True
+
+
+def get_checkpoint_file(checkpoint_name: str) -> Optional[str]:
+    """Retrieve the configuration for a specified checkpoint.
+
+    This function fetches the checkpoint configuration if metadata is enabled.
+    It extracts the file name from the checkpoint metadata or
+    from the call stack if not explicitly provided in the metadata.
+
+    Args:
+        checkpoint_name (str): The name of the checkpoint to retrieve the configuration for.
+
+    Returns:
+        Optional[dict]: A dictionary containing the file name,
+                        or None if metadata is not enabled.
+
+    """
+    enabled, metadata = _get_metadata()
+    if enabled:
+        config = metadata.get_checkpoint(checkpoint_name)
+        return config.file
+    return None

snowflake/snowpark_checkpoints/utils/logging_utils.py

@@ -0,0 +1,67 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from functools import wraps
+from typing import Callable, Optional, TypeVar
+
+from typing_extensions import ParamSpec
+
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+def log(
+    _func: Optional[Callable[P, R]] = None,
+    *,
+    logger: Optional[logging.Logger] = None,
+    log_args: bool = True,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Log the function call and any exceptions that occur.
+
+    Args:
+        _func: The function to log.
+        logger: The logger to use for logging. If not provided, a logger will be created using the
+                function's module name.
+        log_args: Whether to log the arguments passed to the function.
+
+    Returns:
+        A decorator that logs the function call and any exceptions that occur.
+
+    """
+
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            _logger = logging.getLogger(func.__module__) if logger is None else logger
+            if log_args:
+                args_repr = [repr(a) for a in args]
+                kwargs_repr = [f"{k}={v!r}" for k, v in kwargs.items()]
+                formatted_args = ", ".join([*args_repr, *kwargs_repr])
+                _logger.debug("%s called with args %s", func.__name__, formatted_args)
+            try:
+                return func(*args, **kwargs)
+            except Exception:
+                _logger.exception("An error occurred in %s", func.__name__)
+                raise
+
+        return wrapper
+
+    # Handle the case where the decorator is used without parentheses
+    if _func is None:
+        return decorator
+    return decorator(_func)