snowpark-checkpoints-collectors 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

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.2"
+__version__ = "0.4.0"

snowflake/snowpark_checkpoints_collector/collection_common.py

@@ -48,11 +48,12 @@ STRUCT_COLUMN_TYPE = "struct"
 TIMESTAMP_COLUMN_TYPE = "timestamp"
 TIMESTAMP_NTZ_COLUMN_TYPE = "timestamp_ntz"
 
-PANDAS_BOOLEAN_DTYPE = "bool"
+PANDAS_BOOLEAN_DTYPE = "boolean"
 PANDAS_DATETIME_DTYPE = "datetime64[ns]"
 PANDAS_FLOAT_DTYPE = "float64"
-PANDAS_INTEGER_DTYPE = "int64"
+PANDAS_INTEGER_DTYPE = "Int64"
 PANDAS_OBJECT_DTYPE = "object"
+PANDAS_STRING_DTYPE = "string[python]"
 PANDAS_TIMEDELTA_DTYPE = "timedelta64[ns]"
 
 NUMERIC_TYPE_COLLECTION = [
@@ -142,6 +143,8 @@ BACKSLASH_TOKEN = "\\"
 SLASH_TOKEN = "/"
 PYSPARK_NONE_SIZE_VALUE = -1
 PANDAS_LONG_TYPE = "Int64"
+PANDAS_STRING_TYPE = "string"
+PANDAS_FLOAT_TYPE = "float64"
 
 # ENVIRONMENT VARIABLES
 SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH_ENV_VAR = (

snowflake/snowpark_checkpoints_collector/snow_connection_model/snow_connection.py

@@ -22,6 +22,8 @@ from pathlib import Path
 from typing import Callable, Optional
 
 from snowflake.snowpark import Session
+from snowflake.snowpark.functions import col, expr
+from snowflake.snowpark.types import TimestampType
 from snowflake.snowpark_checkpoints_collector.collection_common import (
     DOT_PARQUET_EXTENSION,
 )
@@ -195,9 +197,28 @@ class SnowConnection:
             stage_directory_path,
         )
         dataframe = self.session.read.parquet(path=stage_directory_path)
+        dataframe = convert_timestamps_to_utc_date(dataframe)
         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()
+
+
+def convert_timestamps_to_utc_date(df):
+    """Convert all timestamp columns to UTC normalized timestamps.
+
+    Reading a parquet written by spark from a snowpark session modifies the original timestamps,
+    so this function normalizes timestamps for comparison.
+    """
+    new_cols = []
+    for field in df.schema.fields:
+        if isinstance(field.datatype, TimestampType):
+            utc_normalized_ts = expr(
+                f"convert_timezone('UTC', cast(to_date({field.name}) as timestamp_tz))"
+            ).alias(field.name)
+            new_cols.append(utc_normalized_ts)
+        else:
+            new_cols.append(col(field.name))
+    return df.select(new_cols)

snowflake/snowpark_checkpoints_collector/summary_stats_collector.py

@@ -23,9 +23,15 @@ import pandera as pa
 
 from pyspark.sql import DataFrame as SparkDataFrame
 from pyspark.sql.functions import col
+from pyspark.sql.types import BinaryType as SparkBinaryType
+from pyspark.sql.types import BooleanType as SparkBooleanType
+from pyspark.sql.types import DateType as SparkDateType
 from pyspark.sql.types import DoubleType as SparkDoubleType
+from pyspark.sql.types import FloatType as SparkFloatType
+from pyspark.sql.types import IntegerType as SparkIntegerType
 from pyspark.sql.types import StringType as SparkStringType
-from pyspark.sql.types import StructField
+from pyspark.sql.types import StructField as SparkStructField
+from pyspark.sql.types import TimestampType as SparkTimestampType
 
 from snowflake.snowpark_checkpoints_collector.collection_common import (
     CHECKPOINT_JSON_OUTPUT_FILE_NAME_FORMAT,
@@ -36,8 +42,10 @@ from snowflake.snowpark_checkpoints_collector.collection_common import (
     DOT_PARQUET_EXTENSION,
     INTEGER_TYPE_COLLECTION,
     NULL_COLUMN_TYPE,
+    PANDAS_FLOAT_TYPE,
     PANDAS_LONG_TYPE,
     PANDAS_OBJECT_TYPE_COLLECTION,
+    PANDAS_STRING_TYPE,
     CheckpointMode,
 )
 from snowflake.snowpark_checkpoints_collector.collection_result.model import (
@@ -72,6 +80,16 @@ from snowflake.snowpark_checkpoints_collector.utils.telemetry import report_tele
 
 LOGGER = logging.getLogger(__name__)
 
+default_null_types = {
+    SparkIntegerType(): 0,
+    SparkFloatType(): 0.0,
+    SparkDoubleType(): 0.0,
+    SparkStringType(): "",
+    SparkBooleanType(): False,
+    SparkTimestampType(): None,
+    SparkDateType(): None,
+}
+
 
 @log
 def collect_dataframe_checkpoint(
@@ -253,6 +271,7 @@ def _collect_dataframe_checkpoint_mode_schema(
     column_type_dict: dict[str, any],
     output_path: Optional[str] = None,
 ) -> None:
+    df = normalize_missing_values(df)
     sampled_df = df.sample(sample)
     if sampled_df.isEmpty():
         LOGGER.warning("Sampled DataFrame is empty. Collecting full DataFrame.")
@@ -327,7 +346,16 @@ def _collect_dataframe_checkpoint_mode_schema(
     )
 
 
-def _get_spark_column_types(df: SparkDataFrame) -> dict[str, StructField]:
+def normalize_missing_values(df: SparkDataFrame) -> SparkDataFrame:
+    """Normalize missing values in a PySpark DataFrame to ensure consistent handling of NA values."""
+    for field in df.schema.fields:
+        default_value = default_null_types.get(field.dataType, None)
+        if default_value is not None:
+            df = df.fillna({field.name: default_value})
+    return df
+
+
+def _get_spark_column_types(df: SparkDataFrame) -> dict[str, SparkStructField]:
     schema = df.schema
     column_type_collection = {}
     for field in schema.fields:
@@ -457,14 +485,83 @@ def _to_pandas(sampled_df: SparkDataFrame) -> pandas.DataFrame:
     LOGGER.debug("Converting Spark DataFrame to Pandas DataFrame")
     pandas_df = sampled_df.toPandas()
     for field in sampled_df.schema.fields:
-        has_nan = pandas_df[field.name].isna().any()
         is_integer = field.dataType.typeName() in INTEGER_TYPE_COLLECTION
-        if has_nan and is_integer:
+        is_spark_string = isinstance(field.dataType, SparkStringType)
+        is_spark_binary = isinstance(field.dataType, SparkBinaryType)
+        is_spark_timestamp = isinstance(field.dataType, SparkTimestampType)
+        is_spark_float = isinstance(field.dataType, SparkFloatType)
+        is_spark_boolean = isinstance(field.dataType, SparkBooleanType)
+        is_spark_date = isinstance(field.dataType, SparkDateType)
+        if is_integer:
             LOGGER.debug(
-                "Converting column '%s' to '%s' type",
+                "Converting Spark integer column '%s' to Pandas nullable '%s' type",
                 field.name,
                 PANDAS_LONG_TYPE,
             )
-            pandas_df[field.name] = pandas_df[field.name].astype(PANDAS_LONG_TYPE)
+            pandas_df[field.name] = (
+                pandas_df[field.name].astype(PANDAS_LONG_TYPE).fillna(0)
+            )
+        elif is_spark_string or is_spark_binary:
+            LOGGER.debug(
+                "Converting Spark string column '%s' to Pandas nullable '%s' type",
+                field.name,
+                PANDAS_STRING_TYPE,
+            )
+            pandas_df[field.name] = (
+                pandas_df[field.name].astype(PANDAS_STRING_TYPE).fillna("")
+            )
+        elif is_spark_timestamp:
+            LOGGER.debug(
+                "Converting Spark timestamp column '%s' to UTC naive Pandas datetime",
+                field.name,
+            )
+            pandas_df[field.name] = convert_all_to_utc_naive(
+                pandas_df[field.name]
+            ).fillna(pandas.NaT)
+        elif is_spark_float:
+            LOGGER.debug(
+                "Converting Spark float column '%s' to Pandas nullable float",
+                field.name,
+            )
+            pandas_df[field.name] = (
+                pandas_df[field.name].astype(PANDAS_FLOAT_TYPE).fillna(0.0)
+            )
+        elif is_spark_boolean:
+            LOGGER.debug(
+                "Converting Spark boolean column '%s' to Pandas nullable boolean",
+                field.name,
+            )
+            pandas_df[field.name] = (
+                pandas_df[field.name].astype("boolean").fillna(False)
+            )
+        elif is_spark_date:
+            LOGGER.debug(
+                "Converting Spark date column '%s' to Pandas nullable datetime",
+                field.name,
+            )
+            pandas_df[field.name] = pandas_df[field.name].fillna(pandas.NaT)
 
     return pandas_df
+
+
+def convert_all_to_utc_naive(series: pandas.Series) -> pandas.Series:
+    """Convert all timezone-aware or naive timestamps in a series to UTC naive.
+
+    Naive timestamps are assumed to be in UTC and localized accordingly.
+    Timezone-aware timestamps are converted to UTC and then made naive.
+
+    Args:
+        series (pandas.Series): A Pandas Series of `pd.Timestamp` objects,
+            either naive or timezone-aware.
+
+    Returns:
+        pandas.Series: A Series of UTC-normalized naive timestamps (`tzinfo=None`).
+
+    """
+
+    def convert(ts):
+        if ts.tz is None:
+            ts = ts.tz_localize("UTC")
+        return ts.tz_convert("UTC").tz_localize(None)
+
+    return series.apply(convert)

{snowpark_checkpoints_collectors-0.3.2.dist-info → snowpark_checkpoints_collectors-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-checkpoints-collectors
-Version: 0.3.2
+Version: 0.4.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/

{snowpark_checkpoints_collectors-0.3.2.dist-info → snowpark_checkpoints_collectors-0.4.0.dist-info}/RECORD

@@ -1,8 +1,8 @@
 snowflake/snowpark_checkpoints_collector/__init__.py,sha256=g4NemuA6Mj4O2jkK0yLQ8sEV3owHiiJnBEz_OWvlW1I,1179
-snowflake/snowpark_checkpoints_collector/__version__.py,sha256=1W0aBeLTL5Svy-qrNkZc6gAKtQLDbncpMyN2SlnJhoU,632
-snowflake/snowpark_checkpoints_collector/collection_common.py,sha256=ff5vYffrTRjoJXZQvVQBaOlegAUj_vXBbl1IZidz8Qo,4510
+snowflake/snowpark_checkpoints_collector/__version__.py,sha256=mZG_4eaVJdzo54iJo1tR3khnIA6lKjmN2lUgMoangNY,632
+snowflake/snowpark_checkpoints_collector/collection_common.py,sha256=qHiBWOICEbc1bvpUbfZU_mkmRiy77TB_2eR12mg52oQ,4612
 snowflake/snowpark_checkpoints_collector/singleton.py,sha256=7AgIHQBXVRvPBBCkmBplzkdrrm-xVWf_N8svzA2vF8E,836
-snowflake/snowpark_checkpoints_collector/summary_stats_collector.py,sha256=kRJpVRE9Iy_uqeIPT-__Aan-YLWxQbgSjkJ3w4LpvCc,17214
+snowflake/snowpark_checkpoints_collector/summary_stats_collector.py,sha256=-KhVUcZX9z3_RmFxkcKa-31Ry9PRdcYN_U6O_cPYNhg,20984
 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=XelL7LughZpKl1B_6bJoKOc_PqQg3UleX6zdgVXqTus,2926
 snowflake/snowpark_checkpoints_collector/collection_result/model/collection_point_result_manager.py,sha256=EY6WIIXRbvkTYC4bQn7jFALHh7D2PirVoiLZ5Kq8dNs,2659
@@ -31,13 +31,13 @@ snowflake/snowpark_checkpoints_collector/io_utils/io_default_strategy.py,sha256=
 snowflake/snowpark_checkpoints_collector/io_utils/io_env_strategy.py,sha256=kJMbg2VOKNXXdkGCt_tMMLGEZ2aUl1_nie1qYvx5M-c,3770
 snowflake/snowpark_checkpoints_collector/io_utils/io_file_manager.py,sha256=M17EtANswD5gcgGnmT13OImO_W1uH4K3ewu2CXL9aes,2597
 snowflake/snowpark_checkpoints_collector/snow_connection_model/__init__.py,sha256=kLjZId-aGCljK7lF6yeEw-syEqeTOJDxdXfpv9YxvZA,755
-snowflake/snowpark_checkpoints_collector/snow_connection_model/snow_connection.py,sha256=r3IPnmDMb8151PTgE4YojOhWnxWGPLyBWlgFvvhOfRY,7314
+snowflake/snowpark_checkpoints_collector/snow_connection_model/snow_connection.py,sha256=lM3oqHUHXShALDVVU5ZSuXGREUVfHYHprB5fy1r5T0I,8154
 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=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.2.dist-info/METADATA,sha256=ueYk6-aMlhiKfvH0CZbqjiEjlxUP1VQwKDejX28ju30,6613
-snowpark_checkpoints_collectors-0.3.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-snowpark_checkpoints_collectors-0.3.2.dist-info/licenses/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
-snowpark_checkpoints_collectors-0.3.2.dist-info/RECORD,,
+snowpark_checkpoints_collectors-0.4.0.dist-info/METADATA,sha256=HMpSzXXczuG-5_RuKadoEbA8JjADUvadLj2sQWvu9MY,6613
+snowpark_checkpoints_collectors-0.4.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+snowpark_checkpoints_collectors-0.4.0.dist-info/licenses/LICENSE,sha256=DVQuDIgE45qn836wDaWnYhSdxoLXgpRRKH4RuTjpRZQ,10174
+snowpark_checkpoints_collectors-0.4.0.dist-info/RECORD,,