snowpark-checkpoints-collectors 0.2.0rc1__py3-none-any.whl → 0.2.1__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/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,
+)

snowflake/snowpark_checkpoints_collector/snow_connection_model/snow_connection.py

@@ -0,0 +1,201 @@
+# 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 logging
+import os.path
+import time
+
+from pathlib import Path
+from typing import Callable, Optional
+
+from snowflake.snowpark import Session
+from snowflake.snowpark_checkpoints_collector.collection_common import (
+    DOT_PARQUET_EXTENSION,
+)
+
+
+STAGE_NAME = "CHECKPOINT_STAGE"
+CREATE_STAGE_STATEMENT_FORMAT = "CREATE TEMP STAGE IF NOT EXISTS {}"
+REMOVE_STAGE_FOLDER_STATEMENT_FORMAT = "REMOVE {}"
+STAGE_PATH_FORMAT = "'@{}/{}'"
+PUT_FILE_IN_STAGE_STATEMENT_FORMAT = "PUT '{}' {} AUTO_COMPRESS=FALSE"
+LOGGER = logging.getLogger(__name__)
+
+
+class SnowConnection:
+
+    """Class for manage the Snowpark Connection.
+
+    Attributes:
+        session (Snowpark.Session): the Snowpark session.
+
+    """
+
+    def __init__(self, session: Optional[Session] = None) -> None:
+        """Init SnowConnection.
+
+        Args:
+            session (Snowpark.Session): the Snowpark session.
+
+        """
+        self.session = (
+            session if session is not None else self._create_snowpark_session()
+        )
+        self.stage_id = int(time.time())
+
+    def create_snowflake_table_from_local_parquet(
+        self,
+        table_name: str,
+        input_path: str,
+        stage_path: Optional[str] = None,
+    ) -> None:
+        """Upload to parquet files from the input path and create a table.
+
+        Args:
+            table_name (str): the name of the table to be created.
+            input_path (str): the input directory path.
+            stage_path: (str, optional): the stage path.
+
+        """
+        input_path = (
+            os.path.abspath(input_path)
+            if not os.path.isabs(input_path)
+            else str(Path(input_path).resolve())
+        )
+        folder = f"table_files_{int(time.time())}"
+        stage_path = stage_path if stage_path else folder
+        stage_name = f"{STAGE_NAME}_{self.stage_id}"
+        stage_directory_path = STAGE_PATH_FORMAT.format(stage_name, stage_path)
+
+        def is_parquet_file(file: str):
+            return file.endswith(DOT_PARQUET_EXTENSION)
+
+        try:
+            self.create_tmp_stage(stage_name)
+            self.load_files_to_stage(
+                stage_name, stage_path, input_path, is_parquet_file
+            )
+            self.create_table_from_parquet(table_name, stage_directory_path)
+        finally:
+            LOGGER.info("Removing stage folder %s", stage_directory_path)
+            self.session.sql(
+                REMOVE_STAGE_FOLDER_STATEMENT_FORMAT.format(stage_directory_path)
+            ).collect()
+
+    def create_tmp_stage(self, stage_name: str) -> None:
+        """Create a temp stage in Snowflake.
+
+        Args:
+            stage_name (str): the name of the stage.
+
+        """
+        create_stage_statement = CREATE_STAGE_STATEMENT_FORMAT.format(stage_name)
+        LOGGER.info("Creating temporal stage '%s'", stage_name)
+        self.session.sql(create_stage_statement).collect()
+
+    def load_files_to_stage(
+        self,
+        stage_name: str,
+        folder_name: str,
+        input_path: str,
+        filter_func: Optional[Callable] = None,
+    ) -> None:
+        """Load files to a stage in Snowflake.
+
+        Args:
+            stage_name (str): the name of the stage.
+            folder_name (str): the folder name.
+            input_path (str): the input directory path.
+            filter_func (Callable): the filter function to apply to the files.
+
+        """
+        LOGGER.info("Starting to load files to '%s'", stage_name)
+        input_path = (
+            os.path.abspath(input_path)
+            if not os.path.isabs(input_path)
+            else str(Path(input_path).resolve())
+        )
+
+        def filter_files(name: str):
+            return os.path.isfile(name) and (filter_func(name) if filter_func else True)
+
+        target_dir = os.path.join(input_path, "**", "*")
+        LOGGER.debug("Searching for files in '%s'", input_path)
+        files_collection = glob.glob(target_dir, recursive=True)
+
+        files = [file for file in files_collection if filter_files(file)]
+        files_count = len(files)
+
+        if files_count == 0:
+            raise Exception(f"No files were found in the input directory: {input_path}")
+
+        LOGGER.debug("Found %s files in '%s'", files_count, input_path)
+
+        for file in files:
+            # if file is relative path, convert to absolute path
+            # if absolute path, then try to resolve as some Win32 paths are not in LPN.
+            file_full_path = (
+                str(os.path.abspath(file))
+                if not os.path.isabs(file)
+                else str(Path(file).resolve())
+            )
+            # Snowflake required URI format for input in the put.
+            normalize_file_path = Path(file_full_path).as_uri()
+            new_file_path = file_full_path.replace(input_path, folder_name)
+            # as Posix to convert Windows dir to posix
+            new_file_path = Path(new_file_path).as_posix()
+            stage_file_path = STAGE_PATH_FORMAT.format(stage_name, new_file_path)
+            put_statement = PUT_FILE_IN_STAGE_STATEMENT_FORMAT.format(
+                normalize_file_path, stage_file_path
+            )
+            LOGGER.info("Loading file '%s' to %s", file_full_path, stage_file_path)
+            self.session.sql(put_statement).collect()
+
+    def create_table_from_parquet(
+        self, table_name: str, stage_directory_path: str
+    ) -> None:
+        """Create a table from a parquet file in Snowflake.
+
+        Args:
+            table_name (str): the name of the table.
+            stage_directory_path (str): the stage directory path.
+
+        Raise:
+            Exception: No parquet files were found in the stage
+
+        """
+        LOGGER.info("Starting to create table '%s' from parquet files", table_name)
+        parquet_files = self.session.sql(
+            f"LIST {stage_directory_path} PATTERN='.*{DOT_PARQUET_EXTENSION}'"
+        ).collect()
+        parquet_files_count = len(parquet_files)
+        if parquet_files_count == 0:
+            raise Exception(
+                f"No parquet files were found in the stage: {stage_directory_path}"
+            )
+
+        LOGGER.info(
+            "Reading %s parquet files from %s",
+            parquet_files_count,
+            stage_directory_path,
+        )
+        dataframe = self.session.read.parquet(path=stage_directory_path)
+        LOGGER.info("Creating table '%s' from parquet files", table_name)
+        dataframe.write.save_as_table(table_name=table_name, mode="overwrite")
+
+    def _create_snowpark_session(self) -> Session:
+        LOGGER.info("Creating a Snowpark session using the default connection")
+        return Session.builder.getOrCreate()