snowpark-checkpoints-validators 0.1.0rc3__py3-none-any.whl → 0.1.1__py3-none-any.whl

snowflake/snowpark_checkpoints/snowpark_sampler.py

@@ -0,0 +1,99 @@
+# 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.
+
+from typing import Optional
+
+import pandas
+
+from snowflake.snowpark import DataFrame as SnowparkDataFrame
+from snowflake.snowpark_checkpoints.job_context import SnowparkJobContext
+
+
+class SamplingStrategy:
+    RANDOM_SAMPLE = 1
+    LIMIT = 2
+
+
+class SamplingError(Exception):
+    pass
+
+
+class SamplingAdapter:
+    def __init__(
+        self,
+        job_context: Optional[SnowparkJobContext],
+        sample_frac: Optional[float] = None,
+        sample_number: Optional[int] = None,
+        sampling_strategy: SamplingStrategy = SamplingStrategy.RANDOM_SAMPLE,
+    ):
+        self.pandas_sample_args = []
+        self.job_context = job_context
+        if sample_frac and not (0 <= sample_frac <= 1):
+            raise ValueError(
+                f"'sample_size' value {sample_frac} is out of range (0 <= sample_size <= 1)"
+            )
+
+        self.sample_frac = sample_frac
+        self.sample_number = sample_number
+        self.sampling_strategy = sampling_strategy
+
+    def process_args(self, input_args):
+        # create the intermediate pandas
+        # data frame for the test data
+        for arg in input_args:
+            if isinstance(arg, SnowparkDataFrame):
+                if arg.count() == 0:
+                    raise SamplingError(
+                        "Input DataFrame is empty. Cannot sample from an empty DataFrame."
+                    )
+
+                if self.sampling_strategy == SamplingStrategy.RANDOM_SAMPLE:
+                    if self.sample_frac:
+                        df_sample = arg.sample(frac=self.sample_frac).to_pandas()
+                    else:
+                        df_sample = arg.sample(n=self.sample_number).to_pandas()
+                else:
+                    df_sample = arg.limit(self.sample_number).to_pandas()
+
+                self.pandas_sample_args.append(df_sample)
+            else:
+                self.pandas_sample_args.append(arg)
+
+    def get_sampled_pandas_args(self):
+        return self.pandas_sample_args
+
+    def get_sampled_snowpark_args(self):
+        if self.job_context is None:
+            raise SamplingError("Need a job context to compare with Spark")
+        snowpark_sample_args = []
+        for arg in self.pandas_sample_args:
+            if isinstance(arg, pandas.DataFrame):
+                snowpark_df = self.job_context.snowpark_session.create_dataframe(arg)
+                snowpark_sample_args.append(snowpark_df)
+            else:
+                snowpark_sample_args.append(arg)
+        return snowpark_sample_args
+
+    def get_sampled_spark_args(self):
+        if self.job_context is None:
+            raise SamplingError("Need a job context to compare with Spark")
+        pyspark_sample_args = []
+        for arg in self.pandas_sample_args:
+            if isinstance(arg, pandas.DataFrame):
+                pyspark_df = self.job_context.spark_session.createDataFrame(arg)
+                pyspark_sample_args.append(pyspark_df)
+            else:
+                pyspark_sample_args.append(arg)
+        return pyspark_sample_args

snowflake/snowpark_checkpoints/spark_migration.py

@@ -0,0 +1,222 @@
+# 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.
+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.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)
+
+
+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):
+        _checkpoint_name = checkpoint_name
+        if checkpoint_name is None:
+            _checkpoint_name = snowpark_fn.__name__
+        _checkpoint_name = _replace_special_characters(_checkpoint_name)
+
+        def wrapper(*args, **kwargs):
+            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
+            snowpark_test_results = snowpark_fn(*snowpark_sample_args, **kwargs)
+            spark_test_results = spark_function(*pyspark_sample_args, **kwargs)
+            result, exception = _assert_return(
+                snowpark_test_results,
+                spark_test_results,
+                job_context,
+                _checkpoint_name,
+                output_path,
+            )
+            if not result:
+                raise exception from None
+            # 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
+    ):
+        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:
+
+        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/checkpoint_logger.py

@@ -0,0 +1,52 @@
+# 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 threading
+
+
+class CheckpointLogger:
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            with cls._lock:
+                if not cls._instance:
+                    cls._instance = super().__new__(cls, *args, **kwargs)
+                    cls._instance._initialize()
+        return cls._instance
+
+    def _initialize(self):
+        # Create formatter
+        formatter = logging.Formatter(
+            "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+        )
+
+        self.logger = logging.getLogger("CheckpointLogger")
+        self.logger.setLevel(logging.INFO)
+
+        # Create console handler and set level to debug
+        ch = logging.StreamHandler()
+        ch.setLevel(logging.DEBUG)
+
+        # Add formatter to ch
+        ch.setFormatter(formatter)
+
+        # Add ch to logger
+        self.logger.addHandler(ch)
+
+    def get_logger(self):
+        return self.logger

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,84 @@
+# 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 os
+
+from typing import Optional
+
+from snowflake.snowpark_checkpoints.utils.constants import (
+    SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH_ENV_VAR,
+)
+
+
+# noinspection DuplicatedCode
+def _get_checkpoint_contract_file_path() -> str:
+    return os.environ.get(SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH_ENV_VAR, os.getcwd())
+
+
+# noinspection DuplicatedCode
+def _get_metadata():
+    try:
+        from snowflake.snowpark_checkpoints_configuration.checkpoint_metadata import (
+            CheckpointMetadata,
+        )
+
+        path = _get_checkpoint_contract_file_path()
+        metadata = CheckpointMetadata(path)
+        return True, metadata
+
+    except ImportError:
+        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
+    else:
+        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
+    else:
+        return None