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

snowflake/snowpark_checkpoints/io_utils/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+
+__all__ = ["EnvStrategy", "IOFileManager", "IODefaultStrategy"]
+
+from snowflake.snowpark_checkpoints.io_utils.io_env_strategy import (
+    EnvStrategy,
+)
+from snowflake.snowpark_checkpoints.io_utils.io_default_strategy import (
+    IODefaultStrategy,
+)
+from snowflake.snowpark_checkpoints.io_utils.io_file_manager import (
+    IOFileManager,
+)

snowflake/snowpark_checkpoints/io_utils/io_default_strategy.py

@@ -0,0 +1,57 @@
+# 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 glob
+import os
+
+from pathlib import Path
+from typing import Optional
+
+from snowflake.snowpark_checkpoints.io_utils import EnvStrategy
+
+
+class IODefaultStrategy(EnvStrategy):
+    def mkdir(self, path: str, exist_ok: bool = False) -> None:
+        os.makedirs(path, exist_ok=exist_ok)
+
+    def folder_exists(self, path: str) -> bool:
+        return os.path.isdir(path)
+
+    def file_exists(self, path: str) -> bool:
+        return os.path.isfile(path)
+
+    def write(self, file_path: str, file_content: str, overwrite: bool = True) -> None:
+        mode = "w" if overwrite else "x"
+        with open(file_path, mode) as file:
+            file.write(file_content)
+
+    def read(
+        self, file_path: str, mode: str = "r", encoding: Optional[str] = None
+    ) -> str:
+        with open(file_path, mode=mode, encoding=encoding) as file:
+            return file.read()
+
+    def read_bytes(self, file_path: str) -> bytes:
+        with open(file_path, mode="rb") as f:
+            return f.read()
+
+    def ls(self, path: str, recursive: bool = False) -> list[str]:
+        return glob.glob(path, recursive=recursive)
+
+    def getcwd(self) -> str:
+        return os.getcwd()
+
+    def telemetry_path_files(self, path: str) -> Path:
+        return Path(path)

snowflake/snowpark_checkpoints/io_utils/io_env_strategy.py

@@ -0,0 +1,133 @@
+# 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 abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+
+class EnvStrategy(ABC):
+
+    """An abstract base class that defines methods for file and directory operations.
+
+    Subclasses should implement these methods to provide environment-specific behavior.
+    """
+
+    @abstractmethod
+    def mkdir(self, path: str, exist_ok: bool = False) -> None:
+        """Create a directory.
+
+        Args:
+            path: The name of the directory to create.
+            exist_ok: If False, an error is raised if the directory already exists.
+
+        """
+
+    @abstractmethod
+    def folder_exists(self, path: str) -> bool:
+        """Check if a folder exists.
+
+        Args:
+            path: The path to the folder.
+
+        Returns:
+            bool: True if the folder exists, False otherwise.
+
+        """
+
+    @abstractmethod
+    def file_exists(self, path: str) -> bool:
+        """Check if a file exists.
+
+        Args:
+            path: The path to the file.
+
+        Returns:
+            bool: True if the file exists, False otherwise.
+
+        """
+
+    @abstractmethod
+    def write(self, file_path: str, file_content: str, overwrite: bool = True) -> None:
+        """Write content to a file.
+
+        Args:
+            file_path: The name of the file to write to.
+            file_content: The content to write to the file.
+            overwrite: If True, overwrite the file if it exists.
+
+        """
+
+    @abstractmethod
+    def read(
+        self, file_path: str, mode: str = "r", encoding: Optional[str] = None
+    ) -> str:
+        """Read content from a file.
+
+        Args:
+            file_path: The path to the file to read from.
+            mode: The mode in which to open the file.
+            encoding: The encoding to use for reading the file.
+
+        Returns:
+            str: The content of the file.
+
+        """
+
+    @abstractmethod
+    def read_bytes(self, file_path: str) -> bytes:
+        """Read binary content from a file.
+
+        Args:
+            file_path: The path to the file to read from.
+
+        Returns:
+            bytes: The binary content of the file.
+
+        """
+
+    @abstractmethod
+    def ls(self, path: str, recursive: bool = False) -> list[str]:
+        """List the contents of a directory.
+
+        Args:
+            path: The path to the directory.
+            recursive: If True, list the contents recursively.
+
+        Returns:
+            list[str]: A list of the contents of the directory.
+
+        """
+
+    @abstractmethod
+    def getcwd(self) -> str:
+        """Get the current working directory.
+
+        Returns:
+            str: The current working directory.
+
+        """
+
+    @abstractmethod
+    def telemetry_path_files(self, path: str) -> Path:
+        """Get the path to the telemetry files.
+
+        Args:
+            path: The path to the telemetry directory.
+
+        Returns:
+            Path: The path object representing the telemetry files.
+
+        """

snowflake/snowpark_checkpoints/io_utils/io_file_manager.py

@@ -0,0 +1,76 @@
+# 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 pathlib import Path
+from typing import Optional
+
+from snowflake.snowpark_checkpoints.io_utils import (
+    EnvStrategy,
+    IODefaultStrategy,
+)
+from snowflake.snowpark_checkpoints.singleton import Singleton
+
+
+class IOFileManager(metaclass=Singleton):
+    def __init__(self, strategy: Optional[EnvStrategy] = None):
+        self.strategy = strategy or IODefaultStrategy()
+
+    def mkdir(self, path: str, exist_ok: bool = False) -> None:
+        return self.strategy.mkdir(path, exist_ok)
+
+    def folder_exists(self, path: str) -> bool:
+        return self.strategy.folder_exists(path)
+
+    def file_exists(self, path: str) -> bool:
+        return self.strategy.file_exists(path)
+
+    def write(self, file_path: str, file_content: str, overwrite: bool = True) -> None:
+        return self.strategy.write(file_path, file_content, overwrite)
+
+    def read(
+        self, file_path: str, mode: str = "r", encoding: Optional[str] = None
+    ) -> str:
+        return self.strategy.read(file_path, mode, encoding)
+
+    def read_bytes(self, file_path: str) -> bytes:
+        return self.strategy.read_bytes(file_path)
+
+    def ls(self, path: str, recursive: bool = False) -> list[str]:
+        return self.strategy.ls(path, recursive)
+
+    def getcwd(self) -> str:
+        return self.strategy.getcwd()
+
+    def telemetry_path_files(self, path: str) -> Path:
+        return self.strategy.telemetry_path_files(path)
+
+    def set_strategy(self, strategy: EnvStrategy):
+        """Set the strategy for file and directory operations.
+
+        Args:
+            strategy (EnvStrategy): The strategy to use for file and directory operations.
+
+        """
+        self.strategy = strategy
+
+
+def get_io_file_manager():
+    """Get the singleton instance of IOFileManager.
+
+    Returns:
+        IOFileManager: The singleton instance of IOFileManager.
+
+    """
+    return IOFileManager()

snowflake/snowpark_checkpoints/job_context.py

@@ -0,0 +1,128 @@
+# 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 datetime import datetime
+from typing import Optional
+
+import pandas as pd
+
+from pyspark.sql import SparkSession
+
+from snowflake.snowpark import Session
+from snowflake.snowpark_checkpoints.utils.constants import SCHEMA_EXECUTION_MODE
+
+
+LOGGER = logging.getLogger(__name__)
+RESULTS_TABLE = "SNOWPARK_CHECKPOINTS_REPORT"
+
+
+class SnowparkJobContext:
+
+    """Class used to record migration results in Snowflake.
+
+    Args:
+        snowpark_session: A Snowpark session instance.
+        spark_session: A Spark session instance.
+        job_name: The name of the job.
+        log_results: Whether to log the migration results in Snowflake.
+
+    """
+
+    def __init__(
+        self,
+        snowpark_session: Session,
+        spark_session: SparkSession = None,
+        job_name: Optional[str] = None,
+        log_results: Optional[bool] = True,
+    ):
+        self.log_results = log_results
+        self.job_name = job_name
+        self.spark_session = spark_session or self._create_pyspark_session()
+        self.snowpark_session = snowpark_session
+
+    def _mark_fail(
+        self, message, checkpoint_name, data, execution_mode=SCHEMA_EXECUTION_MODE
+    ):
+        if not self.log_results:
+            LOGGER.warning(
+                (
+                    "Recording of migration results into Snowflake is disabled. "
+                    "Failure result for checkpoint '%s' will not be recorded."
+                ),
+                checkpoint_name,
+            )
+            return
+
+        LOGGER.debug(
+            "Marking failure for checkpoint '%s' in '%s' mode with message '%s'",
+            checkpoint_name,
+            execution_mode,
+            message,
+        )
+
+        session = self.snowpark_session
+        df = pd.DataFrame(
+            {
+                "DATE": [datetime.now()],
+                "JOB": [self.job_name],
+                "STATUS": ["fail"],
+                "CHECKPOINT": [checkpoint_name],
+                "MESSAGE": [message],
+                "DATA": [f"{data}"],
+                "EXECUTION_MODE": [execution_mode],
+            }
+        )
+        report_df = session.createDataFrame(df)
+        LOGGER.info("Writing failure result to table: '%s'", RESULTS_TABLE)
+        report_df.write.mode("append").save_as_table(RESULTS_TABLE)
+
+    def _mark_pass(self, checkpoint_name, execution_mode=SCHEMA_EXECUTION_MODE):
+        if not self.log_results:
+            LOGGER.warning(
+                (
+                    "Recording of migration results into Snowflake is disabled. "
+                    "Pass result for checkpoint '%s' will not be recorded."
+                ),
+                checkpoint_name,
+            )
+            return
+
+        LOGGER.debug(
+            "Marking pass for checkpoint '%s' in '%s' mode",
+            checkpoint_name,
+            execution_mode,
+        )
+
+        session = self.snowpark_session
+        df = pd.DataFrame(
+            {
+                "DATE": [datetime.now()],
+                "JOB": [self.job_name],
+                "STATUS": ["pass"],
+                "CHECKPOINT": [checkpoint_name],
+                "MESSAGE": [""],
+                "DATA": [""],
+                "EXECUTION_MODE": [execution_mode],
+            }
+        )
+        report_df = session.createDataFrame(df)
+        LOGGER.info("Writing pass result to table: '%s'", RESULTS_TABLE)
+        report_df.write.mode("append").save_as_table(RESULTS_TABLE)
+
+    def _create_pyspark_session(self) -> SparkSession:
+        LOGGER.info("Creating a PySpark session")
+        return SparkSession.builder.getOrCreate()

snowflake/snowpark_checkpoints/singleton.py

@@ -0,0 +1,23 @@
+# 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.
+
+
+class Singleton(type):
+    _instances = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            cls._instances[cls] = super().__call__(*args, **kwargs)
+        return cls._instances[cls]

snowflake/snowpark_checkpoints/snowpark_sampler.py

@@ -0,0 +1,124 @@
+# 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 Optional
+
+import pandas
+
+from snowflake.snowpark import DataFrame as SnowparkDataFrame
+from snowflake.snowpark_checkpoints.job_context import SnowparkJobContext
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+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
+        LOGGER.info("Processing %s input argument(s) for sampling", len(input_args))
+        for arg in input_args:
+            if isinstance(arg, SnowparkDataFrame):
+                df_count = arg.count()
+                if df_count == 0:
+                    raise SamplingError(
+                        "Input DataFrame is empty. Cannot sample from an empty DataFrame."
+                    )
+
+                LOGGER.info("Sampling a Snowpark DataFrame with %s rows", df_count)
+                if self.sampling_strategy == SamplingStrategy.RANDOM_SAMPLE:
+                    if self.sample_frac:
+                        LOGGER.info(
+                            "Applying random sampling with fraction %s",
+                            self.sample_frac,
+                        )
+                        df_sample = arg.sample(frac=self.sample_frac).to_pandas()
+                    else:
+                        LOGGER.info(
+                            "Applying random sampling with size %s", self.sample_number
+                        )
+                        df_sample = arg.sample(n=self.sample_number).to_pandas()
+                else:
+                    LOGGER.info(
+                        "Applying limit sampling with size %s", self.sample_number
+                    )
+                    df_sample = arg.limit(self.sample_number).to_pandas()
+
+                LOGGER.info(
+                    "Successfully sampled the DataFrame. Resulting DataFrame shape: %s",
+                    df_sample.shape,
+                )
+                self.pandas_sample_args.append(df_sample)
+            else:
+                LOGGER.debug(
+                    "Argument is not a Snowpark DataFrame. No sampling is applied."
+                )
+                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