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

Author: mkanetsuna

src/snowflake/snowpark/mock/_functions.py

@@ -707,6 +707,13 @@ def mock_abs(expr):
         return abs(expr)
 
 
+@patch("nullifzero")
+def mock_nullifzero(expr: ColumnEmulator) -> ColumnEmulator:
+    result = expr.replace(0, None)
+    result.sf_type = ColumnType(expr.sf_type.datatype, nullable=True)
+    return result
+
+
 @patch("to_decimal")
 def mock_to_decimal(
     e: ColumnEmulator,

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,22 @@ 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],
+            [100],
+            [-5],
+            [-0],
+            [None]
+        ],
+        schema=["v"],
+    )
+    expected = [Row(None), Row(1), Row(100), Row(-5), Row(None), Row(None)]
+    assert origin_df.select(nullifzero(col("v"))).collect() == expected
+
+
 def test_asc_and_desc(session):
     origin_df: DataFrame = session.create_dataframe(
         [