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

snowflake/snowpark_checkpoints_collector/column_pandera_checks/__init__.py

@@ -0,0 +1,20 @@
+# 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__ = ["PanderaColumnChecksManager"]
+
+from snowflake.snowpark_checkpoints_collector.column_pandera_checks.pandera_column_checks_manager import (
+    PanderaColumnChecksManager,
+)

snowflake/snowpark_checkpoints_collector/column_pandera_checks/pandera_column_checks_manager.py

@@ -0,0 +1,241 @@
+# 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 pandas as pd
+
+from pandera import Check, Column
+from pyspark.sql import DataFrame as SparkDataFrame
+from pyspark.sql.functions import col as spark_col
+from pyspark.sql.functions import length as spark_length
+from pyspark.sql.functions import max as spark_max
+from pyspark.sql.functions import min as spark_min
+
+from snowflake.snowpark_checkpoints_collector.collection_common import (
+    BETWEEN_CHECK_ERROR_MESSAGE_FORMAT,
+    BOOLEAN_COLUMN_TYPE,
+    BYTE_COLUMN_TYPE,
+    COLUMN_MAX_KEY,
+    COLUMN_MIN_KEY,
+    DAYTIMEINTERVAL_COLUMN_TYPE,
+    DOUBLE_COLUMN_TYPE,
+    FLOAT_COLUMN_TYPE,
+    INTEGER_COLUMN_TYPE,
+    LONG_COLUMN_TYPE,
+    SHORT_COLUMN_TYPE,
+    STRING_COLUMN_TYPE,
+    TIMESTAMP_COLUMN_TYPE,
+    TIMESTAMP_NTZ_COLUMN_TYPE,
+)
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+def collector_register(cls):
+    """Decorate a class with the checks mechanism.
+
+    Args:
+        cls: The class to decorate.
+
+    Returns:
+        The class to decorate.
+
+    """
+    LOGGER.debug("Starting to register checks from class %s", cls.__name__)
+    cls._collectors = {}
+    for method_name in dir(cls):
+        method = getattr(cls, method_name)
+        if hasattr(method, "_column_type"):
+            col_type_collection = method._column_type
+            for col_type in col_type_collection:
+                cls._collectors[col_type] = method_name
+                LOGGER.debug(
+                    "Registered check '%s' for column type '%s'", method_name, col_type
+                )
+    return cls
+
+
+def column_register(*args):
+    """Decorate a method to register it in the checks mechanism based on column type.
+
+    Args:
+        args: the column type to register.
+
+    Returns:
+        The wrapper.
+
+    """
+
+    def wrapper(func):
+        has_arguments = len(args) > 0
+        if has_arguments:
+            func._column_type = args
+        return func
+
+    return wrapper
+
+
+@collector_register
+class PanderaColumnChecksManager:
+
+    """Manage class for Pandera column checks based on type."""
+
+    def add_checks_column(
+        self,
+        clm_name: str,
+        clm_type: str,
+        pyspark_df: SparkDataFrame,
+        pandera_column: Column,
+    ) -> None:
+        """Add checks to Pandera column based on the column type.
+
+        Args:
+            clm_name (str): the name of the column.
+            clm_type (str): the type of the column.
+            pyspark_df (pyspark.sql.DataFrame): the DataFrame.
+            pandera_column (pandera.Column): the Pandera column.
+
+        """
+        if clm_type not in self._collectors:
+            LOGGER.debug(
+                "No Pandera checks found for column '%s' of type '%s'. Skipping checks for this column.",
+                clm_name,
+                clm_type,
+            )
+            return
+
+        func_name = self._collectors[clm_type]
+        func = getattr(self, func_name)
+        LOGGER.debug(
+            "Adding Pandera checks to column '%s' of type '%s'", clm_name, clm_type
+        )
+        func(clm_name, pyspark_df, pandera_column)
+
+    @column_register(BOOLEAN_COLUMN_TYPE)
+    def _add_boolean_type_checks(
+        self, clm_name: str, pyspark_df: SparkDataFrame, pandera_column: Column
+    ) -> None:
+        pandera_column.checks.extend([Check.isin([True, False])])
+
+    @column_register(DAYTIMEINTERVAL_COLUMN_TYPE)
+    def _add_daytimeinterval_type_checks(
+        self, clm_name: str, pyspark_df: SparkDataFrame, pandera_column: Column
+    ) -> None:
+        select_result = pyspark_df.select(
+            spark_min(spark_col(clm_name)).alias(COLUMN_MIN_KEY),
+            spark_max(spark_col(clm_name)).alias(COLUMN_MAX_KEY),
+        ).collect()[0]
+
+        min_value = pd.to_timedelta(select_result[COLUMN_MIN_KEY])
+        max_value = pd.to_timedelta(select_result[COLUMN_MAX_KEY])
+
+        pandera_column.checks.append(
+            Check.between(
+                min_value=min_value,
+                max_value=max_value,
+                include_max=True,
+                include_min=True,
+                title=BETWEEN_CHECK_ERROR_MESSAGE_FORMAT.format(min_value, max_value),
+            )
+        )
+
+    @column_register(
+        BYTE_COLUMN_TYPE,
+        SHORT_COLUMN_TYPE,
+        INTEGER_COLUMN_TYPE,
+        LONG_COLUMN_TYPE,
+        FLOAT_COLUMN_TYPE,
+        DOUBLE_COLUMN_TYPE,
+    )
+    def _add_numeric_type_checks(
+        self, clm_name: str, pyspark_df: SparkDataFrame, pandera_column: Column
+    ) -> None:
+        select_result = pyspark_df.select(
+            spark_min(spark_col(clm_name)).alias(COLUMN_MIN_KEY),
+            spark_max(spark_col(clm_name)).alias(COLUMN_MAX_KEY),
+        ).collect()[0]
+
+        min_value = select_result[COLUMN_MIN_KEY]
+        max_value = select_result[COLUMN_MAX_KEY]
+
+        pandera_column.checks.append(
+            Check.between(
+                min_value=min_value,
+                max_value=max_value,
+                include_max=True,
+                include_min=True,
+                title=BETWEEN_CHECK_ERROR_MESSAGE_FORMAT.format(min_value, max_value),
+            )
+        )
+
+    @column_register(STRING_COLUMN_TYPE)
+    def _add_string_type_checks(
+        self, clm_name: str, pyspark_df: SparkDataFrame, pandera_column: Column
+    ) -> None:
+        select_result = pyspark_df.select(
+            spark_min(spark_length(spark_col(clm_name))).alias(COLUMN_MIN_KEY),
+            spark_max(spark_length(spark_col(clm_name))).alias(COLUMN_MAX_KEY),
+        ).collect()[0]
+
+        min_length = select_result[COLUMN_MIN_KEY]
+        max_length = select_result[COLUMN_MAX_KEY]
+
+        pandera_column.checks.append(Check.str_length(min_length, max_length))
+
+    @column_register(TIMESTAMP_COLUMN_TYPE)
+    def _add_timestamp_type_checks(
+        self, clm_name: str, pyspark_df: SparkDataFrame, pandera_column: Column
+    ) -> None:
+        select_result = pyspark_df.select(
+            spark_min(spark_col(clm_name)).alias(COLUMN_MIN_KEY),
+            spark_max(spark_col(clm_name)).alias(COLUMN_MAX_KEY),
+        ).collect()[0]
+
+        min_value = pd.Timestamp(select_result[COLUMN_MIN_KEY])
+        max_value = pd.Timestamp(select_result[COLUMN_MAX_KEY])
+
+        pandera_column.checks.append(
+            Check.between(
+                min_value=min_value,
+                max_value=max_value,
+                include_max=True,
+                include_min=True,
+                title=BETWEEN_CHECK_ERROR_MESSAGE_FORMAT.format(min_value, max_value),
+            )
+        )
+
+    @column_register(TIMESTAMP_NTZ_COLUMN_TYPE)
+    def _add_timestamp_ntz_type_checks(
+        self, clm_name: str, pyspark_df: SparkDataFrame, pandera_column: Column
+    ) -> None:
+        select_result = pyspark_df.select(
+            spark_min(spark_col(clm_name)).alias(COLUMN_MIN_KEY),
+            spark_max(spark_col(clm_name)).alias(COLUMN_MAX_KEY),
+        ).collect()[0]
+
+        min_value = pd.Timestamp(select_result[COLUMN_MIN_KEY])
+        max_value = pd.Timestamp(select_result[COLUMN_MAX_KEY])
+
+        pandera_column.checks.append(
+            Check.between(
+                min_value=min_value,
+                max_value=max_value,
+                include_max=True,
+                include_min=True,
+                title=BETWEEN_CHECK_ERROR_MESSAGE_FORMAT.format(min_value, max_value),
+            )
+        )

snowflake/snowpark_checkpoints_collector/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_collector.io_utils.io_env_strategy import (
+    EnvStrategy,
+)
+from snowflake.snowpark_checkpoints_collector.io_utils.io_default_strategy import (
+    IODefaultStrategy,
+)
+from snowflake.snowpark_checkpoints_collector.io_utils.io_file_manager import (
+    IOFileManager,
+)

snowflake/snowpark_checkpoints_collector/io_utils/io_default_strategy.py

@@ -0,0 +1,61 @@
+# 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
+import shutil
+
+from pathlib import Path
+from typing import Optional
+
+from snowflake.snowpark_checkpoints_collector.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 remove_dir(self, path: str) -> None:
+        shutil.rmtree(path)
+
+    def telemetry_path_files(self, path: str) -> Path:
+        return Path(path)

snowflake/snowpark_checkpoints_collector/io_utils/io_env_strategy.py

@@ -0,0 +1,142 @@
+# 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 remove_dir(self, path: str) -> None:
+        """Remove a directory and all its contents.
+
+        Args:
+            path: The path to the directory to remove.
+
+        """
+
+    @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_collector/io_utils/io_file_manager.py

@@ -0,0 +1,79 @@
+# 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_collector.io_utils import (
+    EnvStrategy,
+    IODefaultStrategy,
+)
+from snowflake.snowpark_checkpoints_collector.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 remove_dir(self, path: str) -> None:
+        return self.strategy.remove_dir(path)
+
+    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_collector/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_collector/snow_connection_model/__init__.py

@@ -0,0 +1,20 @@
+# 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__ = ["SnowConnection"]
+
+from snowflake.snowpark_checkpoints_collector.snow_connection_model.snow_connection import (
+    SnowConnection,
+)