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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-checkpoints-collectors
-Version: 0.1.3
+Version: 0.2.0
 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/
@@ -27,19 +27,21 @@ Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: <3.12,>=3.9
 Requires-Dist: pandera[io]==0.20.4
-Requires-Dist: pyspark
 Requires-Dist: snowflake-connector-python
-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.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-collectors
@@ -50,6 +52,18 @@ Description-Content-Type: text/markdown
 ---
 
 **snowpark-checkpoints-collector** package offers a function for extracting information from PySpark dataframes. We can then use that data to validate against the converted Snowpark dataframes to ensure that behavioral equivalence has been achieved.
+
+---
+## Install the library
+```bash
+pip install snowpark-checkpoints-collectors
+```
+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-collectors[pyspark]"
+```
+---
+
 ## Features
 
 - Schema inference collected data mode (Schema): This is the default mode, which leverages Pandera schema inference to obtain the metadata and checks that will be evaluated for the specified dataframe. This mode also collects custom data from columns of the DataFrame based on the PySpark type.

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

@@ -6,6 +6,18 @@
 ---
 
 **snowpark-checkpoints-collector** package offers a function for extracting information from PySpark dataframes. We can then use that data to validate against the converted Snowpark dataframes to ensure that behavioral equivalence has been achieved.
+
+---
+## Install the library
+```bash
+pip install snowpark-checkpoints-collectors
+```
+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-collectors[pyspark]"
+```
+---
+
 ## Features
 
 - Schema inference collected data mode (Schema): This is the default mode, which leverages Pandera schema inference to obtain the metadata and checks that will be evaluated for the specified dataframe. This mode also collects custom data from columns of the DataFrame based on the PySpark type.

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

@@ -26,9 +26,8 @@ classifiers = [
   "Topic :: Scientific/Engineering :: Information Analysis",
 ]
 dependencies = [
-  "snowflake-snowpark-python==1.26.0",
+  "snowflake-snowpark-python>=1.23.0",
   "snowflake-connector-python",
-  "pyspark",
   "pandera[io]==0.20.4",
 ]
 description = "Snowpark column and table statistics collection"
@@ -47,6 +46,9 @@ readme = "README.md"
 requires-python = '>=3.9,<3.12'
 
 [project.optional-dependencies]
+pyspark = [
+  "pyspark>=3.5.0",
+]
 development = [
   "pytest>=8.3.3",
   "pytest-cov>=6.0.0",
@@ -57,6 +59,7 @@ development = [
   "setuptools>=70.0.0",
   "pyarrow>=18.0.0",
   "deepdiff>=8.0.0",
+  "pyspark>=3.5.0",
 ]
 
 [project.urls]

{snowpark_checkpoints_collectors-0.1.3 → snowpark_checkpoints_collectors-0.2.0}/src/snowflake/snowpark_checkpoints_collector/__init__.py

@@ -13,10 +13,18 @@
 # 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
+
 __all__ = ["collect_dataframe_checkpoint", "CheckpointMode"]
 
+from snowflake.snowpark_checkpoints_collector.collection_common import CheckpointMode
 from snowflake.snowpark_checkpoints_collector.summary_stats_collector import (
     collect_dataframe_checkpoint,
 )
-
-from snowflake.snowpark_checkpoints_collector.collection_common import CheckpointMode

{snowpark_checkpoints_collectors-0.1.3 → snowpark_checkpoints_collectors-0.2.0}/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.1.3"
+__version__ = "0.2.0"

{snowpark_checkpoints_collectors-0.1.3 → snowpark_checkpoints_collectors-0.2.0}/src/snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result_manager.py

@@ -12,7 +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 json
+import logging
 
 from typing import Optional
 
@@ -24,6 +26,7 @@ from snowflake.snowpark_checkpoints_collector.utils import file_utils
 
 
 RESULTS_KEY = "results"
+LOGGER = logging.getLogger(__name__)
 
 
 class CollectionPointResultManager(metaclass=Singleton):
@@ -49,6 +52,7 @@ class CollectionPointResultManager(metaclass=Singleton):
 
         """
         result_json = result.get_collection_result_data()
+        LOGGER.debug("Adding a new collection result: %s", result_json)
         self.result_collection.append(result_json)
         self._save_result()
 
@@ -65,5 +69,6 @@ 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)

{snowpark_checkpoints_collectors-0.1.3 → snowpark_checkpoints_collectors-0.2.0}/src/snowflake/snowpark_checkpoints_collector/column_collection/column_collector_manager.py

@@ -12,6 +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 logging
+
 from pyspark.sql import DataFrame as SparkDataFrame
 from pyspark.sql.types import StructField
 
@@ -53,6 +56,9 @@ from snowflake.snowpark_checkpoints_collector.column_collection.model import (
 )
 
 
+LOGGER = logging.getLogger(__name__)
+
+
 def collector_register(cls):
     """Decorate a class with the collection type mechanism.
 
@@ -63,6 +69,7 @@ def collector_register(cls):
         The class to decorate.
 
     """
+    LOGGER.debug("Starting to register collectors from class %s", cls.__name__)
     cls._collectors = {}
     for method_name in dir(cls):
         method = getattr(cls, method_name)
@@ -70,6 +77,11 @@ def collector_register(cls):
             col_type_collection = method._column_type
             for col_type in col_type_collection:
                 cls._collectors[col_type] = method_name
+                LOGGER.debug(
+                    "Registered collector '%s' for column type '%s'",
+                    method_name,
+                    col_type,
+                )
     return cls
 
 
@@ -114,10 +126,21 @@ class ColumnCollectorManager:
         """
         clm_type = struct_field.dataType.typeName()
         if clm_type not in self._collectors:
+            LOGGER.debug(
+                "No collectors found for column '%s' of type '%s'. Skipping collection for this column.",
+                clm_name,
+                clm_type,
+            )
             return {}
 
         func_name = self._collectors[clm_type]
         func = getattr(self, func_name)
+        LOGGER.debug(
+            "Collecting custom data for column '%s' of type '%s' using collector method '%s'",
+            clm_name,
+            clm_type,
+            func_name,
+        )
         data = func(clm_name, struct_field, values)
         return data
 

{snowpark_checkpoints_collectors-0.1.3 → snowpark_checkpoints_collectors-0.2.0}/src/snowflake/snowpark_checkpoints_collector/column_pandera_checks/pandera_column_checks_manager.py

@@ -12,6 +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 logging
+
 import pandas as pd
 
 from pandera import Check, Column
@@ -39,6 +42,9 @@ from snowflake.snowpark_checkpoints_collector.collection_common import (
 )
 
 
+LOGGER = logging.getLogger(__name__)
+
+
 def collector_register(cls):
     """Decorate a class with the checks mechanism.
 
@@ -49,6 +55,7 @@ def collector_register(cls):
         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)
@@ -56,6 +63,9 @@ def collector_register(cls):
             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
 
 
@@ -101,10 +111,18 @@ class PanderaColumnChecksManager:
 
         """
         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)

{snowpark_checkpoints_collectors-0.1.3 → snowpark_checkpoints_collectors-0.2.0}/src/snowflake/snowpark_checkpoints_collector/snow_connection_model/snow_connection.py

@@ -12,7 +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 logging
 import os.path
 import time
 
@@ -30,6 +32,7 @@ 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:
@@ -41,14 +44,16 @@ class SnowConnection:
 
     """
 
-    def __init__(self, session: Session = None) -> None:
+    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 Session.builder.getOrCreate()
+        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(
@@ -84,8 +89,8 @@ class SnowConnection:
                 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()
@@ -98,6 +103,7 @@ class SnowConnection:
 
         """
         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(
@@ -105,7 +111,7 @@ class SnowConnection:
         stage_name: str,
         folder_name: str,
         input_path: str,
-        filter_func: Callable = None,
+        filter_func: Optional[Callable] = None,
     ) -> None:
         """Load files to a stage in Snowflake.
 
@@ -116,6 +122,7 @@ class SnowConnection:
             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)
@@ -126,16 +133,20 @@ class SnowConnection:
             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 len(files) == 0:
+        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.
+            # 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)
@@ -150,6 +161,7 @@ class SnowConnection:
             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(
@@ -165,8 +177,25 @@ class SnowConnection:
             Exception: No parquet files were found in the stage
 
         """
-        files = self.session.sql(f"LIST {stage_directory_path}").collect()
-        if len(files) == 0:
-            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()