snowpark-checkpoints-validators 0.1.3__tar.gz → 0.2.0__tar.gz

{snowpark_checkpoints_validators-0.1.3 → snowpark_checkpoints_validators-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-checkpoints-validators
-Version: 0.1.3
+Version: 0.2.0
 Summary: Migration tools for Snowpark
 Project-URL: Bug Tracker, https://github.com/snowflakedb/snowpark-checkpoints/issues
 Project-URL: Source code, https://github.com/snowflakedb/snowpark-checkpoints/
@@ -26,11 +26,9 @@ Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: <3.12,>=3.9
-Requires-Dist: pandera-report==0.1.2
 Requires-Dist: pandera[io]==0.20.4
-Requires-Dist: pyspark
 Requires-Dist: snowflake-connector-python[pandas]
-Requires-Dist: snowflake-snowpark-python==1.26.0
+Requires-Dist: snowflake-snowpark-python>=1.23.0
 Provides-Extra: development
 Requires-Dist: coverage>=7.6.7; extra == 'development'
 Requires-Dist: deepdiff==8.1.1; extra == 'development'
@@ -38,10 +36,13 @@ Requires-Dist: deepdiff>=8.0.0; extra == 'development'
 Requires-Dist: hatchling==1.25.0; extra == 'development'
 Requires-Dist: pre-commit>=4.0.1; extra == 'development'
 Requires-Dist: pyarrow>=18.0.0; extra == 'development'
+Requires-Dist: pyspark>=3.5.0; extra == 'development'
 Requires-Dist: pytest-cov>=6.0.0; extra == 'development'
 Requires-Dist: pytest>=8.3.3; extra == 'development'
 Requires-Dist: setuptools>=70.0.0; extra == 'development'
 Requires-Dist: twine==5.1.1; extra == 'development'
+Provides-Extra: pyspark
+Requires-Dist: pyspark>=3.5.0; extra == 'pyspark'
 Description-Content-Type: text/markdown
 
 # snowpark-checkpoints-validators
@@ -52,6 +53,17 @@ Description-Content-Type: text/markdown
 
 **snowpark-checkpoints-validators** is a package designed to validate Snowpark DataFrames against predefined schemas and checkpoints. This package ensures data integrity and consistency by performing schema and data validation checks at various stages of a Snowpark pipeline.
 
+---
+## Install the library
+```bash
+pip install snowpark-checkpoints-validators
+```
+This package requires PySpark to be installed in the same environment. If you do not have it, you can install PySpark alongside Snowpark Checkpoints by running the following command:
+```bash
+pip install "snowpark-checkpoints-validators[pyspark]"
+```
+---
+
 ## Features
 
 - Validate Snowpark DataFrames against predefined Pandera schemas.

{snowpark_checkpoints_validators-0.1.3 → snowpark_checkpoints_validators-0.2.0}/README.md

@@ -6,6 +6,17 @@
 
 **snowpark-checkpoints-validators** is a package designed to validate Snowpark DataFrames against predefined schemas and checkpoints. This package ensures data integrity and consistency by performing schema and data validation checks at various stages of a Snowpark pipeline.
 
+---
+## Install the library
+```bash
+pip install snowpark-checkpoints-validators
+```
+This package requires PySpark to be installed in the same environment. If you do not have it, you can install PySpark alongside Snowpark Checkpoints by running the following command:
+```bash
+pip install "snowpark-checkpoints-validators[pyspark]"
+```
+---
+
 ## Features
 
 - Validate Snowpark DataFrames against predefined Pandera schemas.

{snowpark_checkpoints_validators-0.1.3 → snowpark_checkpoints_validators-0.2.0}/pyproject.toml

@@ -26,11 +26,9 @@ classifiers = [
   "Topic :: Scientific/Engineering :: Information Analysis",
 ]
 dependencies = [
-  "snowflake-snowpark-python==1.26.0",
+  "snowflake-snowpark-python>=1.23.0",
   "snowflake-connector-python[pandas]",
-  "pyspark",
   "pandera[io]==0.20.4",
-  "pandera-report==0.1.2",
 ]
 description = "Migration tools for Snowpark"
 dynamic = ['version']
@@ -48,6 +46,9 @@ readme = "README.md"
 requires-python = '>=3.9,<3.12'
 
 [project.optional-dependencies]
+pyspark = [
+  "pyspark>=3.5.0",
+]
 development = [
   "deepdiff==8.1.1",
   "pytest>=8.3.3",
@@ -59,6 +60,7 @@ development = [
   "setuptools>=70.0.0",
   "pyarrow>=18.0.0",
   "deepdiff>=8.0.0",
+  "pyspark>=3.5.0",
 ]
 
 [project.urls]

{snowpark_checkpoints_validators-0.1.3 → snowpark_checkpoints_validators-0.2.0}/src/snowflake/snowpark_checkpoints/__init__.py

@@ -13,16 +13,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+
+
+# Add a NullHandler to prevent logging messages from being output to
+# sys.stderr if no logging configuration is provided.
+logging.getLogger(__name__).addHandler(logging.NullHandler())
+
+# ruff: noqa: E402
+
 from snowflake.snowpark_checkpoints.checkpoint import (
     check_dataframe_schema,
-    check_output_schema,
     check_input_schema,
+    check_output_schema,
     validate_dataframe_checkpoint,
 )
 from snowflake.snowpark_checkpoints.job_context import SnowparkJobContext
 from snowflake.snowpark_checkpoints.spark_migration import check_with_spark
 from snowflake.snowpark_checkpoints.utils.constants import CheckpointMode
 
+
 __all__ = [
     "check_with_spark",
     "SnowparkJobContext",

{snowpark_checkpoints_validators-0.1.3 → snowpark_checkpoints_validators-0.2.0}/src/snowflake/snowpark_checkpoints/__version__.py

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

{snowpark_checkpoints_validators-0.1.3 → snowpark_checkpoints_validators-0.2.0}/src/snowflake/snowpark_checkpoints/checkpoint.py

@@ -14,11 +14,14 @@
 # limitations under the License.
 
 # Wrapper around pandera which logs to snowflake
-from typing import Any, Optional, Union
+
+import logging
+
+from typing import Any, Optional, Union, cast
 
 from pandas import DataFrame as PandasDataFrame
-from pandera import Check, DataFrameSchema
-from pandera_report import DataFrameValidator
+from pandera import Check, DataFrameModel, DataFrameSchema
+from pandera.errors import SchemaError, SchemaErrors
 
 from snowflake.snowpark import DataFrame as SnowparkDataFrame
 from snowflake.snowpark_checkpoints.errors import SchemaValidationError
@@ -27,13 +30,13 @@ from snowflake.snowpark_checkpoints.snowpark_sampler import (
     SamplingAdapter,
     SamplingStrategy,
 )
-from snowflake.snowpark_checkpoints.utils.checkpoint_logger import CheckpointLogger
 from snowflake.snowpark_checkpoints.utils.constants import (
     FAIL_STATUS,
     PASS_STATUS,
     CheckpointMode,
 )
 from snowflake.snowpark_checkpoints.utils.extra_config import is_checkpoint_enabled
+from snowflake.snowpark_checkpoints.utils.logging_utils import log
 from snowflake.snowpark_checkpoints.utils.pandera_check_manager import (
     PanderaCheckManager,
 )
@@ -47,6 +50,10 @@ from snowflake.snowpark_checkpoints.utils.utils_checks import (
 )
 
 
+LOGGER = logging.getLogger(__name__)
+
+
+@log
 def validate_dataframe_checkpoint(
     df: SnowparkDataFrame,
     checkpoint_name: str,
@@ -84,31 +91,45 @@ def validate_dataframe_checkpoint(
     """
     checkpoint_name = _replace_special_characters(checkpoint_name)
 
-    if is_checkpoint_enabled(checkpoint_name):
-
-        if mode == CheckpointMode.SCHEMA:
-            return _check_dataframe_schema_file(
-                df,
-                checkpoint_name,
-                job_context,
-                custom_checks,
-                skip_checks,
-                sample_frac,
-                sample_number,
-                sampling_strategy,
-                output_path,
-            )
-        elif mode == CheckpointMode.DATAFRAME:
-            if job_context is None:
-                raise ValueError(
-                    "Connectionless mode is not supported for Parquet validation"
-                )
-            _check_compare_data(df, job_context, checkpoint_name, output_path)
-        else:
+    if not is_checkpoint_enabled(checkpoint_name):
+        LOGGER.warning(
+            "Checkpoint '%s' is disabled. Skipping DataFrame checkpoint validation.",
+            checkpoint_name,
+        )
+        return None
+
+    LOGGER.info(
+        "Starting DataFrame checkpoint validation for checkpoint '%s'", checkpoint_name
+    )
+
+    if mode == CheckpointMode.SCHEMA:
+        result = _check_dataframe_schema_file(
+            df,
+            checkpoint_name,
+            job_context,
+            custom_checks,
+            skip_checks,
+            sample_frac,
+            sample_number,
+            sampling_strategy,
+            output_path,
+        )
+        return result
+
+    if mode == CheckpointMode.DATAFRAME:
+        if job_context is None:
             raise ValueError(
-                """Invalid validation mode.
-                Please use for schema validation use a 1 or for a full data validation use a 2 for schema validation."""
+                "No job context provided. Please provide one when using DataFrame mode validation."
             )
+        _check_compare_data(df, job_context, checkpoint_name, output_path)
+        return None
+
+    raise ValueError(
+        (
+            "Invalid validation mode. "
+            "Please use 1 for schema validation or 2 for full data validation."
+        ),
+    )
 
 
 def _check_dataframe_schema_file(
@@ -156,7 +177,7 @@ def _check_dataframe_schema_file(
 
     schema = _generate_schema(checkpoint_name, output_path)
 
-    return check_dataframe_schema(
+    return _check_dataframe_schema(
         df,
         schema,
         checkpoint_name,
@@ -170,6 +191,7 @@ def _check_dataframe_schema_file(
     )
 
 
+@log
 def check_dataframe_schema(
     df: SnowparkDataFrame,
     pandera_schema: DataFrameSchema,
@@ -212,6 +234,9 @@ def check_dataframe_schema(
 
     """
     checkpoint_name = _replace_special_characters(checkpoint_name)
+    LOGGER.info(
+        "Starting DataFrame schema validation for checkpoint '%s'", checkpoint_name
+    )
 
     if df is None:
         raise ValueError("DataFrame is required")
@@ -219,19 +244,25 @@ def check_dataframe_schema(
     if pandera_schema is None:
         raise ValueError("Schema is required")
 
-    if is_checkpoint_enabled(checkpoint_name):
-        return _check_dataframe_schema(
-            df,
-            pandera_schema,
+    if not is_checkpoint_enabled(checkpoint_name):
+        LOGGER.warning(
+            "Checkpoint '%s' is disabled. Skipping DataFrame schema validation.",
             checkpoint_name,
-            job_context,
-            custom_checks,
-            skip_checks,
-            sample_frac,
-            sample_number,
-            sampling_strategy,
-            output_path,
         )
+        return None
+
+    return _check_dataframe_schema(
+        df,
+        pandera_schema,
+        checkpoint_name,
+        job_context,
+        custom_checks,
+        skip_checks,
+        sample_frac,
+        sample_number,
+        sampling_strategy,
+        output_path,
+    )
 
 
 @report_telemetry(
@@ -259,17 +290,24 @@ def _check_dataframe_schema(
     pandera_schema_upper, sample_df = _process_sampling(
         df, pandera_schema, job_context, sample_frac, sample_number, sampling_strategy
     )
-
-    # Raises SchemaError on validation issues
-    validator = DataFrameValidator()
-    is_valid, validation_result = validator.validate(
-        pandera_schema_upper, sample_df, validity_flag=True
-    )
+    is_valid, validation_result = _validate(pandera_schema_upper, sample_df)
     if is_valid:
+        LOGGER.info(
+            "DataFrame schema validation passed for checkpoint '%s'",
+            checkpoint_name,
+        )
         if job_context is not None:
             job_context._mark_pass(checkpoint_name)
+        else:
+            LOGGER.warning(
+                "No job context provided. Skipping result recording into Snowflake.",
+            )
         _update_validation_result(checkpoint_name, PASS_STATUS, output_path)
     else:
+        LOGGER.error(
+            "DataFrame schema validation failed for checkpoint '%s'",
+            checkpoint_name,
+        )
         _update_validation_result(checkpoint_name, FAIL_STATUS, output_path)
         raise SchemaValidationError(
             "Snowpark DataFrame schema validation error",
@@ -282,6 +320,7 @@ def _check_dataframe_schema(
 
 
 @report_telemetry(params_list=["pandera_schema"])
+@log
 def check_output_schema(
     pandera_schema: DataFrameSchema,
     checkpoint_name: str,
@@ -318,11 +357,8 @@ def check_output_schema(
             function: The decorated function.
 
         """
-        _checkpoint_name = checkpoint_name
-        if checkpoint_name is None:
-            _checkpoint_name = snowpark_fn.__name__
-        _checkpoint_name = _replace_special_characters(_checkpoint_name)
 
+        @log(log_args=False)
         def wrapper(*args, **kwargs):
             """Wrapp a function to validate the schema of the output of a Snowpark function.
 
@@ -334,7 +370,25 @@ def check_output_schema(
                 Any: The result of the Snowpark function.
 
             """
+            _checkpoint_name = checkpoint_name
+            if checkpoint_name is None:
+                LOGGER.warning(
+                    (
+                        "No checkpoint name provided for output schema validation. "
+                        "Using '%s' as the checkpoint name.",
+                    ),
+                    snowpark_fn.__name__,
+                )
+                _checkpoint_name = snowpark_fn.__name__
+            _checkpoint_name = _replace_special_characters(_checkpoint_name)
+            LOGGER.info(
+                "Starting output schema validation for Snowpark function '%s' and checkpoint '%s'",
+                snowpark_fn.__name__,
+                _checkpoint_name,
+            )
+
             # Run the sampled data in snowpark
+            LOGGER.info("Running the Snowpark function '%s'", snowpark_fn.__name__)
             snowpark_results = snowpark_fn(*args, **kwargs)
             sampler = SamplingAdapter(
                 job_context, sample_frac, sample_number, sampling_strategy
@@ -342,22 +396,28 @@ def check_output_schema(
             sampler.process_args([snowpark_results])
             pandas_sample_args = sampler.get_sampled_pandas_args()
 
-            # Raises SchemaError on validation issues
-            validator = DataFrameValidator()
-            is_valid, validation_result = validator.validate(
-                pandera_schema, pandas_sample_args[0], validity_flag=True
+            is_valid, validation_result = _validate(
+                pandera_schema, pandas_sample_args[0]
             )
-            logger = CheckpointLogger().get_logger()
-            logger.info(
-                f"Checkpoint {_checkpoint_name} validation result:\n{validation_result}"
-            )
-
             if is_valid:
+                LOGGER.info(
+                    "Output schema validation passed for Snowpark function '%s' and checkpoint '%s'",
+                    snowpark_fn.__name__,
+                    _checkpoint_name,
+                )
                 if job_context is not None:
                     job_context._mark_pass(_checkpoint_name)
-
+                else:
+                    LOGGER.warning(
+                        "No job context provided. Skipping result recording into Snowflake.",
+                    )
                 _update_validation_result(_checkpoint_name, PASS_STATUS, output_path)
             else:
+                LOGGER.error(
+                    "Output schema validation failed for Snowpark function '%s' and checkpoint '%s'",
+                    snowpark_fn.__name__,
+                    _checkpoint_name,
+                )
                 _update_validation_result(_checkpoint_name, FAIL_STATUS, output_path)
                 raise SchemaValidationError(
                     "Snowpark output schema validation error",
@@ -365,7 +425,6 @@ def check_output_schema(
                     _checkpoint_name,
                     validation_result,
                 )
-
             return snowpark_results
 
         return wrapper
@@ -374,6 +433,7 @@ def check_output_schema(
 
 
 @report_telemetry(params_list=["pandera_schema"])
+@log
 def check_input_schema(
     pandera_schema: DataFrameSchema,
     checkpoint_name: str,
@@ -414,11 +474,8 @@ def check_input_schema(
             Callable: A wrapper function that performs schema validation before executing the original function.
 
         """
-        _checkpoint_name = checkpoint_name
-        if checkpoint_name is None:
-            _checkpoint_name = snowpark_fn.__name__
-        _checkpoint_name = _replace_special_characters(_checkpoint_name)
 
+        @log(log_args=False)
         def wrapper(*args, **kwargs):
             """Wrapp a function to validate the schema of the input of a Snowpark function.
 
@@ -429,6 +486,23 @@ def check_input_schema(
                 Any: The result of the original function after input validation.
 
             """
+            _checkpoint_name = checkpoint_name
+            if checkpoint_name is None:
+                LOGGER.warning(
+                    (
+                        "No checkpoint name provided for input schema validation. "
+                        "Using '%s' as the checkpoint name."
+                    ),
+                    snowpark_fn.__name__,
+                )
+                _checkpoint_name = snowpark_fn.__name__
+            _checkpoint_name = _replace_special_characters(_checkpoint_name)
+            LOGGER.info(
+                "Starting input schema validation for Snowpark function '%s' and checkpoint '%s'",
+                snowpark_fn.__name__,
+                _checkpoint_name,
+            )
+
             # Run the sampled data in snowpark
             sampler = SamplingAdapter(
                 job_context, sample_frac, sample_number, sampling_strategy
@@ -436,47 +510,71 @@ def check_input_schema(
             sampler.process_args(args)
             pandas_sample_args = sampler.get_sampled_pandas_args()
 
+            LOGGER.info(
+                "Validating %s input argument(s) against a Pandera schema",
+                len(pandas_sample_args),
+            )
             # Raises SchemaError on validation issues
-            for arg in pandas_sample_args:
-                if isinstance(arg, PandasDataFrame):
-
-                    validator = DataFrameValidator()
-                    is_valid, validation_result = validator.validate(
-                        pandera_schema,
-                        arg,
-                        validity_flag=True,
+            for index, arg in enumerate(pandas_sample_args, start=1):
+                if not isinstance(arg, PandasDataFrame):
+                    LOGGER.info(
+                        "Arg %s: Skipping schema validation for non-DataFrame argument",
+                        index,
                     )
+                    continue
 
-                    logger = CheckpointLogger().get_logger()
-                    logger.info(
-                        f"Checkpoint {checkpoint_name} validation result:\n{validation_result}"
+                is_valid, validation_result = _validate(
+                    pandera_schema,
+                    arg,
+                )
+                if is_valid:
+                    LOGGER.info(
+                        "Arg %s: Input schema validation passed",
+                        index,
                     )
-
-                    if is_valid:
-                        if job_context is not None:
-                            job_context._mark_pass(
-                                _checkpoint_name,
-                            )
-
-                        _update_validation_result(
-                            _checkpoint_name,
-                            PASS_STATUS,
-                            output_path,
-                        )
-                    else:
-                        _update_validation_result(
-                            _checkpoint_name,
-                            FAIL_STATUS,
-                            output_path,
-                        )
-                        raise SchemaValidationError(
-                            "Snowpark input schema validation error",
-                            job_context,
+                    if job_context is not None:
+                        job_context._mark_pass(
                             _checkpoint_name,
-                            validation_result,
                         )
+                    _update_validation_result(
+                        _checkpoint_name,
+                        PASS_STATUS,
+                        output_path,
+                    )
+                else:
+                    LOGGER.error(
+                        "Arg %s: Input schema validation failed",
+                        index,
+                    )
+                    _update_validation_result(
+                        _checkpoint_name,
+                        FAIL_STATUS,
+                        output_path,
+                    )
+                    raise SchemaValidationError(
+                        "Snowpark input schema validation error",
+                        job_context,
+                        _checkpoint_name,
+                        validation_result,
+                    )
             return snowpark_fn(*args, **kwargs)
 
         return wrapper
 
     return check_input_with_decorator
+
+
+def _validate(
+    schema: Union[type[DataFrameModel], DataFrameSchema],
+    df: PandasDataFrame,
+    lazy: bool = True,
+) -> tuple[bool, PandasDataFrame]:
+    if not isinstance(schema, DataFrameSchema):
+        schema = schema.to_schema()
+    is_valid = True
+    try:
+        df = schema.validate(df, lazy=lazy)
+    except (SchemaErrors, SchemaError) as schema_errors:
+        df = cast(PandasDataFrame, schema_errors.failure_cases)
+        is_valid = False
+    return is_valid, df