Merge pull request #160 from icanbwell/sg-HDE-3445

HDE-3445: Adding support for other datatypes in skip_if_column_null
icanbwell · Oct 5, 2023 · 172e48a · 172e48a
2 parents 79ebe68 + d40b739
commit 172e48a
Showing 1 changed file with 31 additions and 6 deletions.
diff --git a/spark_auto_mapper/automappers/with_column_base.py b/spark_auto_mapper/automappers/with_column_base.py
@@ -4,7 +4,17 @@
 
 # noinspection PyUnresolvedReferences
 from pyspark.sql.functions import col, when, lit, size
-from pyspark.sql.types import DataType, StructField, ArrayType
+from pyspark.sql.types import (
+    DataType,
+    StructField,
+    ArrayType,
+    StringType,
+    TimestampType,
+    DateType,
+    MapType,
+    ByteType,
+    BinaryType,
+)
 from pyspark.sql.utils import AnalysisException
 from spark_data_frame_comparer.schema_comparer import SchemaComparer
 
@@ -71,7 +81,24 @@ def get_column_spec(self, source_df: Optional[DataFrame]) -> Column:
                     )
                     column_to_check = f"b.{column}"
                     # wrap column spec in when
-                    if isinstance(column_type, ArrayType):
+                    if isinstance(
+                        column_type,
+                        (StringType, TimestampType, DateType, ByteType, BinaryType),
+                    ):
+                        column_spec = (
+                            when(
+                                col(column_to_check).isNull()
+                                | col(column_to_check).eqNullSafe(""),
+                                lit(None),
+                            )
+                            if is_first_when_case
+                            else column_spec.when(
+                                col(column_to_check).isNull()
+                                | col(column_to_check).eqNullSafe(""),
+                                lit(None),
+                            )
+                        )
+                    elif isinstance(column_type, (ArrayType, MapType)):
                         column_spec = (
                             when(
                                 col(column_to_check).isNull()
@@ -88,14 +115,12 @@ def get_column_spec(self, source_df: Optional[DataFrame]) -> Column:
                     else:
                         column_spec = (
                             when(
-                                col(column_to_check).isNull()
-                                | col(column_to_check).eqNullSafe(""),
+                                col(column_to_check).isNull(),
                                 lit(None),
                             )
                             if is_first_when_case
                             else column_spec.when(
-                                col(column_to_check).isNull()
-                                | col(column_to_check).eqNullSafe(""),
+                                col(column_to_check).isNull(),
                                 lit(None),
                             )
                         )