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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-checkpoints-collectors
-Version: 0.3.0
+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/
@@ -101,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.3.0 → 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.3.0 → 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.3.0 → 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.3.0"
+__version__ = "0.3.1"

{snowpark_checkpoints_collectors-0.3.0 → 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.3.0 → snowpark_checkpoints_collectors-0.3.1}/src/snowflake/snowpark_checkpoints_collector/summary_stats_collector.py

@@ -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,

{snowpark_checkpoints_collectors-0.3.0 → snowpark_checkpoints_collectors-0.3.1}/src/snowflake/snowpark_checkpoints_collector/utils/file_utils.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 import inspect
 import os
+import re
 import tempfile
 
 from typing import Optional
@@ -84,7 +85,10 @@ def get_collection_point_source_file_path() -> str:
 
     """
     try:
-        collection_point_file_path = inspect.stack()[2].filename
+        stack_frame = _get_stack_frame()
+        if not stack_frame:
+            return UNKNOWN_SOURCE_FILE
+        collection_point_file_path = stack_frame.filename
         is_temporal_file_path = _is_temporal_path(collection_point_file_path)
         if is_temporal_file_path:
             ipynb_file_path_collection = _get_ipynb_file_path_collection()
@@ -100,15 +104,18 @@ def get_collection_point_source_file_path() -> str:
 
 
 def get_collection_point_line_of_code() -> int:
-    """Find the line of code of the source file where collection point it is.
+    """Find the line of code of the source file where collection point is.
 
     Returns:
         int: returns the line of code of the source file where collection point it is.
 
     """
     try:
-        collection_point_file_path = inspect.stack()[2].filename
-        collection_point_line_of_code = inspect.stack()[2].lineno
+        stack_frame = _get_stack_frame()
+        if not stack_frame:
+            return UNKNOWN_LINE_OF_CODE
+        collection_point_file_path = stack_frame.filename
+        collection_point_line_of_code = stack_frame.lineno
         is_temporal_file_path = _is_temporal_path(collection_point_file_path)
         if is_temporal_file_path:
             collection_point_line_of_code = UNKNOWN_LINE_OF_CODE
@@ -124,6 +131,21 @@ def _is_temporal_path(path: str) -> bool:
     return is_temporal_path
 
 
+def _get_stack_frame() -> inspect.FrameInfo:
+    batch = inspect.stack()[:7]
+    batch.reverse()
+    collect_frame_regex = r"(collect_dataframe_checkpoint)"
+
+    for frame in batch:
+        if (
+            frame.code_context is not None
+            and len(frame.code_context) >= 0
+            and re.search(collect_frame_regex, frame.code_context[0])
+        ):
+            return frame
+    return None
+
+
 def _get_ipynb_file_path_collection() -> list[str]:
     current_working_directory_path = get_io_file_manager().getcwd()
     cwd_file_name_collection = get_io_file_manager().ls(current_working_directory_path)

{snowpark_checkpoints_collectors-0.3.0 → snowpark_checkpoints_collectors-0.3.1}/test/unit/test_file_utils.py

@@ -54,25 +54,35 @@ def test_get_output_file_path_create():
 
 
 def test_get_collection_point_source_file_path_scenario_python_source_file():
-    collection_point_source_file_path = (
-        file_utils.get_collection_point_source_file_path()
-    )
-    assert collection_point_source_file_path != UNKNOWN_SOURCE_FILE
+    with mock.patch(
+        "snowflake.snowpark_checkpoints_collector.utils.file_utils._get_stack_frame",
+        return_value=mock.MagicMock(filename="abc.py", lineno=1),
+    ):
+        collection_point_source_file_path = (
+            file_utils.get_collection_point_source_file_path()
+        )
+        assert collection_point_source_file_path != UNKNOWN_SOURCE_FILE
 
 
 def test_get_collection_point_source_file_path_scenario_notebook_source_file():
-    with mock.patch(
-        "snowflake.snowpark_checkpoints_collector.utils.file_utils._is_temporal_path",
-        return_value=True,
-    ):
-        with mock.patch(
+    with (
+        mock.patch(
+            "snowflake.snowpark_checkpoints_collector.utils.file_utils._is_temporal_path",
+            return_value=True,
+        ),
+        mock.patch(
+            "snowflake.snowpark_checkpoints_collector.utils.file_utils._get_stack_frame",
+            return_value=mock.MagicMock(filename="abc.ipynb", lineno=1),
+        ),
+        mock.patch(
             "snowflake.snowpark_checkpoints_collector.utils.file_utils._get_ipynb_file_path_collection",
             return_value=["abc.ipynb"],
-        ):
-            collection_point_source_file_path = (
-                file_utils.get_collection_point_source_file_path()
-            )
-            assert collection_point_source_file_path != UNKNOWN_SOURCE_FILE
+        ),
+    ):
+        collection_point_source_file_path = (
+            file_utils.get_collection_point_source_file_path()
+        )
+        assert collection_point_source_file_path != UNKNOWN_SOURCE_FILE
 
 
 def test_get_collection_point_source_file_path_scenario_unknown_source_file():
@@ -102,8 +112,12 @@ def test_get_collection_point_source_file_path_scenario_exception():
 
 
 def test_get_collection_point_line_of_code_scenario_python_source_file():
-    collection_point_line_of_code = file_utils.get_collection_point_line_of_code()
-    assert collection_point_line_of_code != UNKNOWN_LINE_OF_CODE
+    with mock.patch(
+        "snowflake.snowpark_checkpoints_collector.utils.file_utils._get_stack_frame",
+        return_value=mock.MagicMock(filename=__file__, lineno=1),
+    ):
+        collection_point_line_of_code = file_utils.get_collection_point_line_of_code()
+        assert collection_point_line_of_code != UNKNOWN_LINE_OF_CODE
 
 
 def test_get_collection_point_line_of_code_scenario_notebook_source_file():

snowpark_checkpoints_collectors-0.3.1/test/unit/test_summary_stats_collector.py

@@ -0,0 +1,100 @@
+# 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 os
+import tempfile
+
+from datetime import datetime
+from typing import get_type_hints
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from snowflake.snowpark_checkpoints_collector.summary_stats_collector import (
+    collect_dataframe_checkpoint,
+    generate_parquet_for_spark_df,
+    xcollect_dataframe_checkpoint,
+)
+
+
+def test_generate_parquet_for_spark_df_exception():
+    spark = MagicMock()
+    spark_df = MagicMock()
+    spark_df.dtypes = []
+    spark_df.select = MagicMock()
+    spark_df = spark.createDataFrame()
+    parquet_directory = os.path.join(
+        tempfile.gettempdir(),
+        f"test_spark_df_checkpoint_{datetime.now().strftime('%Y%m%d%H%M%S')}",
+    )
+
+    with pytest.raises(Exception, match="No parquet files were generated."):
+        generate_parquet_for_spark_df(spark_df, parquet_directory)
+
+
+def test_collect_dataframe_checkpoint_disabled_checkpoint(
+    caplog: pytest.LogCaptureFixture,
+):
+    """Test that collect_dataframe_checkpoint logs a message when the checkpoint is disabled."""
+    pyspark_df = MagicMock()
+    checkpoint_name = "my_checkpoint"
+    module_name = "snowflake.snowpark_checkpoints_collector.summary_stats_collector"
+    expected_exception_error_msg = "Checkpoint 'my_checkpoint' is disabled. Please enable it in the checkpoints.json file."
+    expected_fix_suggestion_msg = "In case you want to skip it, use the xcollect_dataframe_checkpoint method instead."
+    try:
+        with (
+            caplog.at_level(
+                level=logging.INFO,
+                logger=module_name,
+            ),
+            patch(
+                f"{module_name}.is_checkpoint_enabled",
+                return_value=False,
+            ) as mock_is_checkpoint_enabled,
+        ):
+            collect_dataframe_checkpoint(pyspark_df, checkpoint_name)
+    except Exception as e:
+        mock_is_checkpoint_enabled.assert_called_once_with(checkpoint_name)
+        error_msg = e.args[0]
+        fix_suggestion_msg = e.args[1]
+        assert error_msg == expected_exception_error_msg
+        assert fix_suggestion_msg == expected_fix_suggestion_msg
+
+
+def test_skip_collector_parameters_commutability():
+    collect_hints = get_type_hints(collect_dataframe_checkpoint)
+    x_collect_hints = get_type_hints(xcollect_dataframe_checkpoint)
+
+    collect_params = {
+        name: hint for name, hint in collect_hints.items() if name != "return"
+    }
+    x_collect_params = {
+        name: hint for name, hint in x_collect_hints.items() if name != "return"
+    }
+    assert (
+        collect_params == x_collect_params
+    ), "The parameters of collect_dataframe_checkpoint and xcollect_dataframe_checkpoint must be the same."
+
+
+def test_skip_collector_return_type_commutability():
+    collect_hints = get_type_hints(collect_dataframe_checkpoint)
+    x_collect_hints = get_type_hints(xcollect_dataframe_checkpoint)
+
+    collect_return = collect_hints.get("return")
+    x_collect_return = x_collect_hints.get("return")
+    assert (
+        collect_return == x_collect_return
+    ), "The return type of collect_dataframe_checkpoint and xcollect_dataframe_checkpoint must be the same."

snowpark_checkpoints_collectors-0.3.0/test/unit/test_summary_stats_collector.py

@@ -1,70 +0,0 @@
-# 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 os
-import tempfile
-
-from datetime import datetime
-from unittest.mock import MagicMock, patch
-
-import pytest
-
-from snowflake.snowpark_checkpoints_collector.summary_stats_collector import (
-    collect_dataframe_checkpoint,
-    generate_parquet_for_spark_df,
-)
-
-
-def test_generate_parquet_for_spark_df_exception():
-    spark = MagicMock()
-    spark_df = MagicMock()
-    spark_df.dtypes = []
-    spark_df.select = MagicMock()
-    spark_df = spark.createDataFrame()
-    parquet_directory = os.path.join(
-        tempfile.gettempdir(),
-        f"test_spark_df_checkpoint_{datetime.now().strftime('%Y%m%d%H%M%S')}",
-    )
-
-    with pytest.raises(Exception, match="No parquet files were generated."):
-        generate_parquet_for_spark_df(spark_df, parquet_directory)
-
-
-def test_collect_dataframe_checkpoint_disabled_checkpoint(
-    caplog: pytest.LogCaptureFixture,
-):
-    """Test that collect_dataframe_checkpoint logs a message when the checkpoint is disabled."""
-    pyspark_df = MagicMock()
-    checkpoint_name = "my_checkpoint"
-    module_name = "snowflake.snowpark_checkpoints_collector.summary_stats_collector"
-    expected_log_msg = (
-        f"Checkpoint '{checkpoint_name}' is disabled. Skipping collection."
-    )
-
-    with (
-        caplog.at_level(
-            level=logging.INFO,
-            logger=module_name,
-        ),
-        patch(
-            f"{module_name}.is_checkpoint_enabled",
-            return_value=False,
-        ) as mock_is_checkpoint_enabled,
-    ):
-        collect_dataframe_checkpoint(pyspark_df, checkpoint_name)
-
-    mock_is_checkpoint_enabled.assert_called_once_with(checkpoint_name)
-    assert expected_log_msg in caplog.messages