snowpark-checkpoints-validators 0.3.2__tar.gz → 0.3.3__tar.gz

{snowpark_checkpoints_validators-0.3.2 → snowpark_checkpoints_validators-0.3.3}/PKG-INFO

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

{snowpark_checkpoints_validators-0.3.2 → snowpark_checkpoints_validators-0.3.3}/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.3.2"
+__version__ = "0.3.3"

{snowpark_checkpoints_validators-0.3.2 → snowpark_checkpoints_validators-0.3.3}/src/snowflake/snowpark_checkpoints/snowpark_sampler.py

@@ -21,6 +21,10 @@ import pandas
 
 from snowflake.snowpark import DataFrame as SnowparkDataFrame
 from snowflake.snowpark_checkpoints.job_context import SnowparkJobContext
+from snowflake.snowpark_checkpoints.utils.constants import (
+    INTEGER_TYPE_COLLECTION,
+    PANDAS_LONG_TYPE,
+)
 
 
 LOGGER = logging.getLogger(__name__)
@@ -73,17 +77,17 @@ class SamplingAdapter:
                             "Applying random sampling with fraction %s",
                             self.sample_frac,
                         )
-                        df_sample = arg.sample(frac=self.sample_frac).to_pandas()
+                        df_sample = to_pandas(arg.sample(frac=self.sample_frac))
                     else:
                         LOGGER.info(
                             "Applying random sampling with size %s", self.sample_number
                         )
-                        df_sample = arg.sample(n=self.sample_number).to_pandas()
+                        df_sample = to_pandas(arg.sample(n=self.sample_number))
                 else:
                     LOGGER.info(
                         "Applying limit sampling with size %s", self.sample_number
                     )
-                    df_sample = arg.limit(self.sample_number).to_pandas()
+                    df_sample = to_pandas(arg.limit(self.sample_number))
 
                 LOGGER.info(
                     "Successfully sampled the DataFrame. Resulting DataFrame shape: %s",
@@ -122,3 +126,26 @@ class SamplingAdapter:
             else:
                 pyspark_sample_args.append(arg)
         return pyspark_sample_args
+
+
+def to_pandas(sampled_df: SnowparkDataFrame) -> pandas.DataFrame:
+    """Convert a Snowpark DataFrame to a Pandas DataFrame, handling missing values and type conversions."""
+    LOGGER.debug("Converting Snowpark DataFrame to Pandas DataFrame")
+    pandas_df = sampled_df.toPandas()
+    pandas_df = normalize_missing_values_pandas(pandas_df)
+    return pandas_df
+
+
+def normalize_missing_values_pandas(df: pandas.DataFrame) -> pandas.DataFrame:
+    """Normalize missing values in a Pandas DataFrame to ensure consistent handling of NA values."""
+    fill_values = {}
+    for col, dtype in df.dtypes.items():
+        if dtype in INTEGER_TYPE_COLLECTION or str(dtype) in PANDAS_LONG_TYPE:
+            fill_values[col] = 0
+        elif dtype is float or dtype == "float64":
+            fill_values[col] = 0.0
+        elif dtype is bool or dtype == "bool" or dtype == "boolean":
+            fill_values[col] = False
+        elif dtype is object or dtype == "object" or dtype is str:
+            fill_values[col] = ""
+    return df.fillna(value=fill_values)

{snowpark_checkpoints_validators-0.3.2 → snowpark_checkpoints_validators-0.3.3}/src/snowflake/snowpark_checkpoints/utils/constants.py

@@ -133,3 +133,16 @@ VALIDATION_RESULTS_JSON_FILE_NAME: Final[str] = "checkpoint_validation_results.j
 SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH_ENV_VAR: Final[
     str
 ] = "SNOWFLAKE_CHECKPOINT_CONTRACT_FILE_PATH"
+
+BYTE_COLUMN_TYPE = "byte"
+INTEGER_COLUMN_TYPE = "integer"
+LONG_COLUMN_TYPE = "long"
+SHORT_COLUMN_TYPE = "short"
+
+PANDAS_LONG_TYPE = "Int64"
+INTEGER_TYPE_COLLECTION = [
+    BYTE_COLUMN_TYPE,
+    INTEGER_COLUMN_TYPE,
+    LONG_COLUMN_TYPE,
+    SHORT_COLUMN_TYPE,
+]

snowpark_checkpoints_validators-0.3.3/test/unit/test_snowpark_sampler.py

@@ -0,0 +1,96 @@
+# 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.
+
+from snowflake.snowpark_checkpoints.snowpark_sampler import (
+    normalize_missing_values_pandas,
+    to_pandas,
+)
+import pandas as pd
+import numpy as np
+
+
+class DummySnowparkDF:
+    def __init__(self, pandas_df):
+        self._pandas_df = pandas_df
+
+    def toPandas(self):
+        return self._pandas_df
+
+
+def test_normalize_missing_values_pandas_integers():
+    df = pd.DataFrame({"a": [1, None, 3], "b": [None, 2, 3]}, dtype="Int64")
+    result = normalize_missing_values_pandas(df)
+    assert result["a"].iloc[1] == 0
+    assert result["b"].iloc[0] == 0
+
+
+def test_normalize_missing_values_pandas_floats():
+    df = pd.DataFrame({"a": [1.1, np.nan, 3.3], "b": [np.nan, 2.2, 3.3]})
+    result = normalize_missing_values_pandas(df)
+    assert result["a"].iloc[1] == 0.0
+    assert result["b"].iloc[0] == 0.0
+
+
+def test_normalize_missing_values_pandas_bools():
+    df = pd.DataFrame(
+        {"a": [True, None, False], "b": [None, False, True]}, dtype="boolean"
+    )
+    result = normalize_missing_values_pandas(df)
+    assert result["a"].iloc[1] is np.False_
+    assert result["b"].iloc[0] is np.False_
+
+
+def test_normalize_missing_values_pandas_objects_and_strings():
+    df = pd.DataFrame(
+        {"a": ["foo", None, "bar"], "b": [None, "baz", "qux"]}, dtype="object"
+    )
+    result = normalize_missing_values_pandas(df)
+    assert result["a"].iloc[1] == ""
+    assert result["b"].iloc[0] == ""
+
+
+def test_normalize_missing_values_pandas_mixed_types():  ###
+    df = pd.DataFrame(
+        {
+            "int_col": pd.Series([1, None], dtype="Int64"),
+            "float_col": [1.1, np.nan],
+            "bool_col": pd.Series([True, None], dtype="boolean"),
+            "str_col": ["a", None],
+        }
+    )
+    result = normalize_missing_values_pandas(df)
+    assert result["int_col"].iloc[1] == 0
+    assert result["float_col"].iloc[1] == 0.0
+    assert result["bool_col"].iloc[1] is np.False_
+    assert result["str_col"].iloc[1] == ""
+
+
+def test_to_pandas_calls_normalize(monkeypatch):
+    df = pd.DataFrame({"a": [1, None], "b": [None, "foo"]})
+    dummy_snowpark_df = DummySnowparkDF(df)
+    called = {}
+
+    def fake_normalize(df_arg):
+        called["was_called"] = True
+        return df_arg.fillna(0)
+
+    monkeypatch.setattr(
+        "snowflake.snowpark_checkpoints.snowpark_sampler.normalize_missing_values_pandas",
+        fake_normalize,
+    )
+
+    result = to_pandas(dummy_snowpark_df)
+    assert called.get("was_called", False)
+    assert isinstance(result, pd.DataFrame)