snowpark-checkpoints-collectors 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

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

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"

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:

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,

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.dist-info → snowpark_checkpoints_collectors-0.3.1.dist-info}/METADATA

@@ -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.dist-info → snowpark_checkpoints_collectors-0.3.1.dist-info}/RECORD

@@ -1,10 +1,10 @@
-snowflake/snowpark_checkpoints_collector/__init__.py,sha256=GIESlH2W6g_qdcnyRqw9yjsvEkt0aniFvGixKlF4K7A,1096
-snowflake/snowpark_checkpoints_collector/__version__.py,sha256=kbbDnlkY7JOLNHvfWYkCO_mOBOV9GniMGdxYoQpLhyg,632
+snowflake/snowpark_checkpoints_collector/__init__.py,sha256=g4NemuA6Mj4O2jkK0yLQ8sEV3owHiiJnBEz_OWvlW1I,1179
+snowflake/snowpark_checkpoints_collector/__version__.py,sha256=uSRs7fRupFeQ-z3PtU_6qh6ry8YBaSAnEIAvLhJKUR8,632
 snowflake/snowpark_checkpoints_collector/collection_common.py,sha256=ff5vYffrTRjoJXZQvVQBaOlegAUj_vXBbl1IZidz8Qo,4510
 snowflake/snowpark_checkpoints_collector/singleton.py,sha256=7AgIHQBXVRvPBBCkmBplzkdrrm-xVWf_N8svzA2vF8E,836
-snowflake/snowpark_checkpoints_collector/summary_stats_collector.py,sha256=SD5MRF7zSDKXpxekMWdg5gO7ZcZr6Y548vkkKpG_jZs,14745
+snowflake/snowpark_checkpoints_collector/summary_stats_collector.py,sha256=kRJpVRE9Iy_uqeIPT-__Aan-YLWxQbgSjkJ3w4LpvCc,17214
 snowflake/snowpark_checkpoints_collector/collection_result/model/__init__.py,sha256=jZzx29WzrjH7C_6ZsBGoe4PxbW_oM4uIjySS1axIM34,1000
-snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result.py,sha256=8xD9zGnFJ7Rz9RUXIys7JnV3kQD4mk8QwNOTxAihSjQ,2908
+snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result.py,sha256=XelL7LughZpKl1B_6bJoKOc_PqQg3UleX6zdgVXqTus,2926
 snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result_manager.py,sha256=EY6WIIXRbvkTYC4bQn7jFALHh7D2PirVoiLZ5Kq8dNs,2659
 snowflake/snowpark_checkpoints_collector/column_collection/__init__.py,sha256=hpTh1V7hqBSHxNUqISwfxdz-NLD-7oZEMLXDUuRsoOU,783
 snowflake/snowpark_checkpoints_collector/column_collection/column_collector_manager.py,sha256=Vav_vbiipHFIAdHxeQG4ZK1BAmWTi_18hBnVeIeXFRs,9670
@@ -34,10 +34,10 @@ snowflake/snowpark_checkpoints_collector/snow_connection_model/__init__.py,sha25
 snowflake/snowpark_checkpoints_collector/snow_connection_model/snow_connection.py,sha256=r3IPnmDMb8151PTgE4YojOhWnxWGPLyBWlgFvvhOfRY,7314
 snowflake/snowpark_checkpoints_collector/utils/checkpoint_name_utils.py,sha256=Xc4k3JU6A96-79VFRR8NrNAUPeO3V1DEAhngg-hLlU4,1787
 snowflake/snowpark_checkpoints_collector/utils/extra_config.py,sha256=3kVf6WVA-EuyMpTO3ycTlXMSCHtytGtT6wkV4U2Hyjw,5195
-snowflake/snowpark_checkpoints_collector/utils/file_utils.py,sha256=C1gZmQHvLMgHMVc5kTTpcCaUPw5PtpajY_Uu18mMy6c,4515
+snowflake/snowpark_checkpoints_collector/utils/file_utils.py,sha256=5ztlNCv9GdSktUvtdfydv86cCFcmSXCdD4axZXJrOQQ,5125
 snowflake/snowpark_checkpoints_collector/utils/logging_utils.py,sha256=yyi6X5DqKeTg0HRhvsH6ymYp2P0wbnyKIzI2RzrQS7k,2278
 snowflake/snowpark_checkpoints_collector/utils/telemetry.py,sha256=ueN9vM8j5YNax7jMcnEj_UrgGkoeMv_hJHVKjN7hiJE,32161
-snowpark_checkpoints_collectors-0.3.0.dist-info/METADATA,sha256=4nXrRjc1glZUTrb9J8brIHPzyrE43GRKNu7lrqfGMZU,6061
-snowpark_checkpoints_collectors-0.3.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-snowpark_checkpoints_collectors-0.3.0.dist-info/licenses/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
-snowpark_checkpoints_collectors-0.3.0.dist-info/RECORD,,
+snowpark_checkpoints_collectors-0.3.1.dist-info/METADATA,sha256=dPrRK0GO-5U3bQUntCpYyRfG3J_z-v6bV9hX2s8CxdM,6613
+snowpark_checkpoints_collectors-0.3.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+snowpark_checkpoints_collectors-0.3.1.dist-info/licenses/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
+snowpark_checkpoints_collectors-0.3.1.dist-info/RECORD,,