snowpark-checkpoints-collectors 0.2.1__tar.gz → 0.3.1__tar.gz

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-checkpoints-collectors
-Version: 0.2.1
+Version: 0.3.1
 Summary: Snowpark column and table statistics collection
 Project-URL: Bug Tracker, https://github.com/snowflakedb/snowpark-checkpoints/issues
 Project-URL: Source code, https://github.com/snowflakedb/snowpark-checkpoints/
@@ -30,6 +30,7 @@ Requires-Dist: pandera[io]==0.20.4
 Requires-Dist: snowflake-connector-python
 Requires-Dist: snowflake-snowpark-python>=1.23.0
 Provides-Extra: development
+Requires-Dist: certifi==2025.1.31; extra == 'development'
 Requires-Dist: coverage>=7.6.7; extra == 'development'
 Requires-Dist: deepdiff>=8.0.0; extra == 'development'
 Requires-Dist: hatchling==1.25.0; extra == 'development'
@@ -100,6 +101,28 @@ def collect_dataframe_checkpoint(
 - `output_path`: The output path to save the checkpoint, defaults to current working directory.
 
 
+### Skip DataFrame Checkpoint Collection
+
+
+
+```python
+from pyspark.sql import DataFrame as SparkDataFrame
+from snowflake.snowpark_checkpoints_collector.collection_common import CheckpointMode
+from typing import Optional
+
+# Signature of the function
+def xcollect_dataframe_checkpoint(
+    df: SparkDataFrame,
+    checkpoint_name: str,
+    sample: Optional[float] = None,
+    mode: Optional[CheckpointMode] = None,
+    output_path: Optional[str] = None,
+) -> None:
+    ...
+```
+
+The signature of the method is the same of `collect_dataframe_checkpoint`.
+
 ## Usage Example
 
 ### Schema mode

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.1}/README.md

@@ -54,6 +54,28 @@ def collect_dataframe_checkpoint(
 - `output_path`: The output path to save the checkpoint, defaults to current working directory.
 
 
+### Skip DataFrame Checkpoint Collection
+
+
+
+```python
+from pyspark.sql import DataFrame as SparkDataFrame
+from snowflake.snowpark_checkpoints_collector.collection_common import CheckpointMode
+from typing import Optional
+
+# Signature of the function
+def xcollect_dataframe_checkpoint(
+    df: SparkDataFrame,
+    checkpoint_name: str,
+    sample: Optional[float] = None,
+    mode: Optional[CheckpointMode] = None,
+    output_path: Optional[str] = None,
+) -> None:
+    ...
+```
+
+The signature of the method is the same of `collect_dataframe_checkpoint`.
+
 ## Usage Example
 
 ### Schema mode

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.1}/pyproject.toml

@@ -60,6 +60,7 @@ development = [
   "pyarrow>=18.0.0",
   "deepdiff>=8.0.0",
   "pyspark>=3.5.0",
+  "certifi==2025.1.31",
 ]
 
 [project.urls]

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.1}/src/snowflake/snowpark_checkpoints_collector/__init__.py

@@ -22,9 +22,14 @@ logging.getLogger(__name__).addHandler(logging.NullHandler())
 
 # ruff: noqa: E402
 
-__all__ = ["collect_dataframe_checkpoint", "CheckpointMode"]
+__all__ = [
+    "collect_dataframe_checkpoint",
+    "CheckpointMode",
+    "xcollect_dataframe_checkpoint",
+]
 
 from snowflake.snowpark_checkpoints_collector.collection_common import CheckpointMode
 from snowflake.snowpark_checkpoints_collector.summary_stats_collector import (
     collect_dataframe_checkpoint,
+    xcollect_dataframe_checkpoint,
 )

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.1}/src/snowflake/snowpark_checkpoints_collector/__version__.py

@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "0.2.1"
+__version__ = "0.3.1"

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.1}/src/snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result.py

@@ -30,6 +30,7 @@ CHECKPOINT_NAME_KEY = "checkpoint_name"
 class CollectionResult(Enum):
     FAIL = "FAIL"
     PASS = "PASS"
+    SKIP = "SKIP"
 
 
 class CollectionPointResult:

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.1}/src/snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result_manager.py

@@ -21,6 +21,9 @@ from typing import Optional
 from snowflake.snowpark_checkpoints_collector.collection_result.model import (
     CollectionPointResult,
 )
+from snowflake.snowpark_checkpoints_collector.io_utils.io_file_manager import (
+    get_io_file_manager,
+)
 from snowflake.snowpark_checkpoints_collector.singleton import Singleton
 from snowflake.snowpark_checkpoints_collector.utils import file_utils
 
@@ -70,5 +73,4 @@ class CollectionPointResultManager(metaclass=Singleton):
     def _save_result(self) -> None:
         result_collection_json = self.to_json()
         LOGGER.info("Saving collection results to '%s'", self.output_file_path)
-        with open(self.output_file_path, "w") as f:
-            f.write(result_collection_json)
+        get_io_file_manager().write(self.output_file_path, result_collection_json)

snowpark_checkpoints_collectors-0.3.1/src/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,
+)

snowpark_checkpoints_collectors-0.3.1/src/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)

snowpark_checkpoints_collectors-0.3.1/src/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.
+
+        """

snowpark_checkpoints_collectors-0.3.1/src/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()

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.1}/src/snowflake/snowpark_checkpoints_collector/snow_connection_model/snow_connection.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import glob
+import io
 import logging
 import os.path
 import time
@@ -25,6 +25,9 @@ from snowflake.snowpark import Session
 from snowflake.snowpark_checkpoints_collector.collection_common import (
     DOT_PARQUET_EXTENSION,
 )
+from snowflake.snowpark_checkpoints_collector.io_utils.io_file_manager import (
+    get_io_file_manager,
+)
 
 
 STAGE_NAME = "CHECKPOINT_STAGE"
@@ -130,11 +133,13 @@ class SnowConnection:
         )
 
         def filter_files(name: str):
-            return os.path.isfile(name) and (filter_func(name) if filter_func else True)
+            return get_io_file_manager().file_exists(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_collection = get_io_file_manager().ls(target_dir, recursive=True)
 
         files = [file for file in files_collection if filter_files(file)]
         files_count = len(files)
@@ -152,17 +157,14 @@ class SnowConnection:
                 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
-            )
+            parquet_file = get_io_file_manager().read_bytes(file_full_path)
+            binary_parquet = io.BytesIO(parquet_file)
             LOGGER.info("Loading file '%s' to %s", file_full_path, stage_file_path)
-            self.session.sql(put_statement).collect()
+            self.session.file.put_stream(binary_parquet, stage_file_path)
 
     def create_table_from_parquet(
         self, table_name: str, stage_directory_path: str

{snowpark_checkpoints_collectors-0.2.1 → snowpark_checkpoints_collectors-0.3.1}/src/snowflake/snowpark_checkpoints_collector/summary_stats_collector.py

@@ -12,12 +12,9 @@
 # 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 json
 import logging
 import os
-import shutil
 
 from typing import Optional
 
@@ -54,6 +51,9 @@ from snowflake.snowpark_checkpoints_collector.column_collection import (
 from snowflake.snowpark_checkpoints_collector.column_pandera_checks import (
     PanderaColumnChecksManager,
 )
+from snowflake.snowpark_checkpoints_collector.io_utils.io_file_manager import (
+    get_io_file_manager,
+)
 from snowflake.snowpark_checkpoints_collector.snow_connection_model import (
     SnowConnection,
 )
@@ -117,11 +117,10 @@ def collect_dataframe_checkpoint(
             "Checkpoint names must only contain alphanumeric characters, underscores and dollar signs."
         )
     if not is_checkpoint_enabled(normalized_checkpoint_name):
-        LOGGER.info(
-            "Checkpoint '%s' is disabled. Skipping collection.",
-            normalized_checkpoint_name,
+        raise Exception(
+            f"Checkpoint '{normalized_checkpoint_name}' is disabled. Please enable it in the checkpoints.json file.",
+            "In case you want to skip it, use the xcollect_dataframe_checkpoint method instead.",
         )
-        return
 
     LOGGER.info("Starting to collect checkpoint '%s'", normalized_checkpoint_name)
     LOGGER.debug("DataFrame size: %s rows", df.count())
@@ -184,6 +183,68 @@ def collect_dataframe_checkpoint(
         collection_point_result_manager.add_result(collection_point_result)
 
 
+@log
+def xcollect_dataframe_checkpoint(
+    df: SparkDataFrame,
+    checkpoint_name: str,
+    sample: Optional[float] = None,
+    mode: Optional[CheckpointMode] = None,
+    output_path: Optional[str] = None,
+) -> None:
+    """Skips the collection of metadata from a Dataframe checkpoint.
+
+    Args:
+        df (SparkDataFrame): The input Spark DataFrame to skip.
+        checkpoint_name (str): The name of the checkpoint.
+        sample (float, optional): Fraction of DataFrame to sample for schema inference.
+            Defaults to 1.0.
+        mode (CheckpointMode): The mode to execution the collection.
+            Defaults to CheckpointMode.Schema
+        output_path (str, optional): The output path to save the checkpoint.
+            Defaults to Current working Directory.
+
+    Raises:
+        Exception: Invalid mode value.
+        Exception: Invalid checkpoint name. Checkpoint names must only contain alphanumeric characters,
+                     underscores and dollar signs.
+
+    """
+    normalized_checkpoint_name = checkpoint_name_utils.normalize_checkpoint_name(
+        checkpoint_name
+    )
+    if normalized_checkpoint_name != checkpoint_name:
+        LOGGER.warning(
+            "Checkpoint name '%s' was normalized to '%s'",
+            checkpoint_name,
+            normalized_checkpoint_name,
+        )
+    is_valid_checkpoint_name = checkpoint_name_utils.is_valid_checkpoint_name(
+        normalized_checkpoint_name
+    )
+    if not is_valid_checkpoint_name:
+        raise Exception(
+            f"Invalid checkpoint name: {normalized_checkpoint_name}. "
+            "Checkpoint names must only contain alphanumeric characters, underscores and dollar signs."
+        )
+
+    LOGGER.warning(
+        "Checkpoint '%s' is disabled. Skipping collection.",
+        normalized_checkpoint_name,
+    )
+
+    collection_point_file_path = file_utils.get_collection_point_source_file_path()
+    collection_point_line_of_code = file_utils.get_collection_point_line_of_code()
+    collection_point_result = CollectionPointResult(
+        collection_point_file_path,
+        collection_point_line_of_code,
+        normalized_checkpoint_name,
+    )
+
+    collection_point_result.result = CollectionResult.SKIP
+    collection_point_result_manager = CollectionPointResultManager(output_path)
+    collection_point_result_manager.add_result(collection_point_result)
+
+
 @report_telemetry(params_list=["column_type_dict"])
 def _collect_dataframe_checkpoint_mode_schema(
     checkpoint_name: str,
@@ -321,8 +382,7 @@ def _generate_json_checkpoint_file(
     output_directory_path = file_utils.get_output_directory_path(output_path)
     checkpoint_file_path = os.path.join(output_directory_path, checkpoint_file_name)
     LOGGER.info("Writing DataFrame JSON schema file to '%s'", checkpoint_file_path)
-    with open(checkpoint_file_path, "w") as f:
-        f.write(dataframe_schema_contract)
+    get_io_file_manager().write(checkpoint_file_path, dataframe_schema_contract)
 
 
 @report_telemetry(params_list=["df"])
@@ -366,17 +426,17 @@ def generate_parquet_for_spark_df(spark_df: SparkDataFrame, output_path: str) ->
     ]
     converted_df = spark_df.select(new_cols)
 
-    if os.path.exists(output_path):
+    if get_io_file_manager().folder_exists(output_path):
         LOGGER.warning(
             "Output directory '%s' already exists. Deleting it...", output_path
         )
-        shutil.rmtree(output_path)
+        get_io_file_manager().remove_dir(output_path)
 
     LOGGER.info("Writing DataFrame to parquet files at '%s'", output_path)
     converted_df.write.parquet(output_path, mode="overwrite")
 
     target_dir = os.path.join(output_path, "**", f"*{DOT_PARQUET_EXTENSION}")
-    parquet_files = glob.glob(target_dir, recursive=True)
+    parquet_files = get_io_file_manager().ls(target_dir, recursive=True)
     parquet_files_count = len(parquet_files)
     if parquet_files_count == 0:
         raise Exception("No parquet files were generated.")