From 98e9e5c7f56ba1db02da60ee7f590b89b3920679 Mon Sep 17 00:00:00 2001
From: masaaya <156692516+mkanetsuna@users.noreply.github.com>
Date: Wed, 17 Dec 2025 04:07:07 +0900
Subject: [PATCH] SNOW-2913796: Add nullifzero function support for local
 testing (#4036)

---
 src/snowflake/snowpark/mock/_functions.py |  7 +++++++
 tests/mock/test_functions.py              | 17 +++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/src/snowflake/snowpark/mock/_functions.py b/src/snowflake/snowpark/mock/_functions.py
index 42f18acea8..b3c50c450a 100644
--- a/src/snowflake/snowpark/mock/_functions.py
+++ b/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,
diff --git a/tests/mock/test_functions.py b/tests/mock/test_functions.py
index df70994af6..4386915fce 100644
--- a/tests/mock/test_functions.py
+++ b/tests/mock/test_functions.py
@@ -29,6 +29,7 @@
     lit,
     max,
     min,
+    nullifzero,
     rank,
     row_number,
     sum,
@@ -161,6 +162,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(
         [