NO-SNOW: Add nullifzero function support for local testing (#4036)

Author: mkanetsuna

src/snowflake/snowpark/mock/_functions.py

@@ -114,7 +114,6 @@ def _check_constant_result(self, input_data, args, result):
         return result
 
     def __call__(self, *args, input_data=None, row_number=None, **kwargs):
-
         if self._pass_input_data:
             kwargs["raw_input"] = input_data
         if self._pass_row_index:
@@ -393,7 +392,7 @@ def mock_stddev(column: ColumnEmulator) -> ColumnEmulator:
 
 @patch("approx_percentile_accumulate")
 def mock_approx_percentile_accumulate(
-    column: Union[TableEmulator, ColumnEmulator]
+    column: Union[TableEmulator, ColumnEmulator],
 ) -> ColumnEmulator:
     # TODO SNOW-1800512: Fix, returns dummy of 42 for now.
     _logger.warning("TODO SNOW-1800512: Returns dummy value of 42 now, need to fix.")
@@ -707,6 +706,18 @@ def mock_abs(expr):
         return abs(expr)
 
 
+@patch("nullifzero")
+def mock_nullifzero(expr: ColumnEmulator) -> ColumnEmulator:
+    def convert_zero_to_null(value):
+        if value == 0 or value == 0.0:
+            return None
+        return value
+
+    result = expr.apply(convert_zero_to_null)
+    result.sf_type = ColumnType(expr.sf_type.datatype, nullable=True)
+    return result
+
+
 @patch("to_decimal")
 def mock_to_decimal(
     e: ColumnEmulator,
@@ -974,7 +985,6 @@ def convert_timestamp(row):
         if data is None:
             return None
         try:
-
             datatype = column.sf_type.datatype
             if isinstance(datatype, TimestampType):
                 # data is datetime.datetime type
@@ -1206,7 +1216,10 @@ def convert_char(data, _fmt):
             return try_convert(convert_numeric_to_str, try_cast, data)
         elif isinstance(source_datatype, (DateType, TimeType)):
             default_format = _DEFAULT_OUTPUT_FORMAT.get(type(source_datatype))
-            (format, _,) = convert_snowflake_datetime_format(
+            (
+                format,
+                _,
+            ) = convert_snowflake_datetime_format(
                 _fmt, default_format=default_format, is_input_format=False
             )
             convert_date_time_to_str = (
@@ -1219,7 +1232,10 @@ def convert_char(data, _fmt):
             )
         elif isinstance(source_datatype, TimestampType):
             default_format = _DEFAULT_OUTPUT_FORMAT.get(TimestampType)
-            (format, fractional_seconds,) = convert_snowflake_datetime_format(
+            (
+                format,
+                fractional_seconds,
+            ) = convert_snowflake_datetime_format(
                 _fmt, default_format, is_input_format=False
             )
             # handle 3f, can use str index

tests/mock/test_functions.py

@@ -29,6 +29,7 @@
     lit,
     max,
     min,
+    nullifzero,
     rank,
     row_number,
     sum,
@@ -40,7 +41,6 @@
 from snowflake.snowpark.mock.exceptions import SnowparkLocalTestingException
 from snowflake.snowpark.types import IntegerType
 from snowflake.snowpark.window import Window
-
 from tests.utils import Utils
 
 
@@ -161,6 +161,26 @@ def test_abs(session):
     assert origin_df.select(abs(col("m"))).collect() == [Row(1), Row(1), Row(2)]
 
 
+def test_nullifzero(session):
+    origin_df: DataFrame = session.create_dataframe(
+        [
+            [0],
+            [1],
+            [0.0],
+            [100],
+            [-5],
+        ],
+        schema=["value"],
+    )
+    assert origin_df.select(nullifzero(col("value"))).collect() == [
+        Row(None),
+        Row(1),
+        Row(None),
+        Row(100),
+        Row(-5),
+    ]
+
+
 def test_asc_and_desc(session):
     origin_df: DataFrame = session.create_dataframe(
         [