Merge pull request #413 from capitalone/develop

Author: fdosani

datacompy/__init__.py

@@ -18,7 +18,7 @@
 Then extended to carry that functionality over to Spark Dataframes.
 """
 
-__version__ = "0.16.6"
+__version__ = "0.16.7"
 
 import platform
 from warnings import warn

datacompy/polars.py

@@ -1014,7 +1014,7 @@ def generate_id_within_group(
             dataframe[join_columns]
             .cast(pl.String)
             .fill_null(default_value)
-            .select(rn=pl.col(dataframe.columns[0]).cum_count().over(join_columns))
+            .select(rn=pl.col(join_columns[0]).cum_count().over(join_columns))
             .to_series()
         )
     else:

datacompy/snowflake.py

@@ -31,7 +31,6 @@
 from ordered_set import OrderedSet
 
 from datacompy.base import BaseCompare
-from datacompy.spark.sql import decimal_comparator
 
 LOG = logging.getLogger(__name__)
 
@@ -51,6 +50,29 @@
         trim,
         when,
     )
+    from snowflake.snowpark.types import (
+        ByteType,
+        DateType,
+        DecimalType,
+        DoubleType,
+        FloatType,
+        IntegerType,
+        LongType,
+        ShortType,
+        StringType,
+        TimestampType,
+    )
+
+    NUMERIC_SNOWPARK_TYPES = [
+        ByteType,
+        ShortType,
+        IntegerType,
+        LongType,
+        FloatType,
+        DoubleType,
+        DecimalType,
+    ]
+
 
 except ImportError:
     LOG.warning(
@@ -59,17 +81,6 @@
     )
 
 
-NUMERIC_SNOWPARK_TYPES = [
-    "tinyint",
-    "smallint",
-    "int",
-    "bigint",
-    "float",
-    "double",
-    decimal_comparator(),
-]
-
-
 class SnowflakeCompare(BaseCompare):
     """Comparison class to be used to compare whether two Snowpark dataframes are equal.
 
@@ -490,17 +501,18 @@ def _calculate_column_compare_stats(self, column: str) -> None:
             match_rate = 0
         LOG.info(f"{column}: {match_cnt} / {row_cnt} ({match_rate:.2%}) match")
 
-        col1_dtype, _ = _get_column_dtypes(self.df1, column, column)
-        col2_dtype, _ = _get_column_dtypes(self.df2, column, column)
+        col1_dtype_instance, _ = _get_column_dtypes(self.df1, column, column)
+        col2_dtype_instance, _ = _get_column_dtypes(self.df2, column, column)
+        col1_dtype, col2_dtype = type(col1_dtype_instance), type(col2_dtype_instance)
 
         self.column_stats.append(
             {
                 "column": column,
                 "match_column": col_match,
                 "match_cnt": match_cnt,
                 "unequal_cnt": row_cnt - match_cnt,
-                "dtype1": str(col1_dtype),
-                "dtype2": str(col2_dtype),
+                "dtype1": col1_dtype_instance.simple_string(),
+                "dtype2": col2_dtype_instance.simple_string(),
                 "all_match": all(
                     (
                         col1_dtype == col2_dtype,
@@ -995,7 +1007,10 @@ def columns_equal(
         A column of boolean values are added.  True == the values match, False == the
         values don't match.
     """
-    base_dtype, compare_dtype = _get_column_dtypes(dataframe, col_1, col_2)
+    base_dtype_instance, compare_dtype_instance = _get_column_dtypes(
+        dataframe, col_1, col_2
+    )
+    base_dtype, compare_dtype = type(base_dtype_instance), type(compare_dtype_instance)
     if _is_comparable(base_dtype, compare_dtype):
         if (base_dtype in NUMERIC_SNOWPARK_TYPES) and (
             compare_dtype in NUMERIC_SNOWPARK_TYPES
@@ -1028,7 +1043,7 @@ def columns_equal(
             )
     else:
         LOG.debug(
-            f"Skipping {col_1}({base_dtype}) and {col_2}({compare_dtype}), columns are not comparable"
+            f"Skipping {col_1}({base_dtype_instance.simple_string()}) and {col_2}({compare_dtype_instance.simple_string()}), columns are not comparable"
         )
         dataframe = dataframe.withColumn(col_match, lit(False))
     return dataframe
@@ -1217,8 +1232,14 @@ def _get_column_dtypes(
     Tuple(str, str)
         Tuple of base and compare datatype
     """
-    base_dtype = next(d[1] for d in dataframe.dtypes if d[0] == col_1)
-    compare_dtype = next(d[1] for d in dataframe.dtypes if d[0] == col_2)
+    df_raw_dtypes = [
+        (name, field.datatype)
+        for name, field in zip(
+            dataframe.schema.names, dataframe.schema.fields, strict=False
+        )
+    ]
+    base_dtype = next(d[1] for d in df_raw_dtypes if d[0] == col_1)
+    compare_dtype = next(d[1] for d in df_raw_dtypes if d[0] == col_2)
     return base_dtype, compare_dtype
 
 
@@ -1244,10 +1265,10 @@ def _is_comparable(type1: str, type2: str) -> bool:
     return (
         type1 == type2
         or (type1 in NUMERIC_SNOWPARK_TYPES and type2 in NUMERIC_SNOWPARK_TYPES)
-        or ("string" in type1 and type2 == "date")
-        or (type1 == "date" and "string" in type2)
-        or ("string" in type1 and type2 == "timestamp")
-        or (type1 == "timestamp" and "string" in type2)
+        or (type1 == StringType and type2 == DateType)
+        or (type1 == DateType and type2 == StringType)
+        or (type1 == StringType and type2 == TimestampType)
+        or (type1 == TimestampType and type2 == StringType)
     )
 
 

docs/source/img/benchmarks.png

@@ -43,6 +43,12 @@
 )
 from pandas.testing import assert_frame_equal, assert_series_equal
 from snowflake.snowpark.exceptions import SnowparkSQLException
+from snowflake.snowpark.types import (
+    DecimalType,
+    StringType,
+    StructField,
+    StructType,
+)
 
 logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
 
@@ -189,6 +195,34 @@ def test_date_columns_equal_with_ignore_spaces(snowpark_session):
     assert_series_equal(expect_out, actual_out_rev, check_names=False)
 
 
+def test_columns_equal_same_type_dif_length(snowpark_session):
+    schema = StructType(
+        [
+            StructField("NAME", StringType(length=20)),
+            StructField("DECIMAL_VAL", DecimalType(precision=7, scale=5)),
+            StructField("NAME_COPY", StringType(is_max_size=True)),
+            StructField("DECIMAL_VAL_COPY", DecimalType(precision=20, scale=10)),
+        ]
+    )
+    data = [
+        ["Alice", 10.44556, "Alice", 10.44556],
+        ["Bob", 2.33445, "Bob", 2.33445],
+        ["Charlie", 5.2234, "Charlie", 5.2234],
+    ]
+
+    df = snowpark_session.create_dataframe(data, schema=schema)
+    assert (
+        columns_equal(df, "NAME", "NAME_COPY", "NAME_ACTUAL")
+        .toPandas()["NAME_ACTUAL"]
+        .all()
+    )
+    assert (
+        columns_equal(df, "DECIMAL_VAL", "DECIMAL_VAL_COPY", "DECIMAL_VAL_ACTUAL")
+        .toPandas()["DECIMAL_VAL_ACTUAL"]
+        .all()
+    )
+
+
 def test_date_columns_unequal(snowpark_session):
     """I want datetime fields to match with dates stored as strings"""
     data = [{"A": "2017-01-01", "B": "2017-01-02"}, {"A": "2017-01-01"}]