SNOW-2913796: [Local Testing] Add support for nullifzero function

[mkanetsuna]

What is the current behavior?

The nullifzero() function is not implemented in the Local Testing Framework. When we attempt to use this function with local_testing=True, they encounter the following error:

from snowflake.snowpark import Session
from snowflake.snowpark.functions import col, nullifzero

session = Session.builder.config("local_testing", True).create()
df = session.create_dataframe([[0], [1], [100]], schema=["value"])
result = df.select(nullifzero(col("value")))
result.show()

error:

NotImplementedError: [Local Testing] Function nullifzero is not implemented.
You can implement and make a patch by using the `snowflake.snowpark.mock.patch` decorator.

What is the desired behavior?

The nullifzero() function should work in Local Testing mode with the following behavior:

• When the input value is 0 (integer or float), return NULL
• When the input value is non-zero, return the original value
• Maintain the original data type while ensuring the result is nullable

This would match the behavior of Snowflake's NULLIFZERO SQL function.

If this is not an existing feature in snowflake-snowpark-python. How would this impact/improve non local testing mode?

This feature does not impact non-local testing mode. The nullifzero() function already works correctly when connected to Snowflake.

References, Other Background

#4036
I'm willing to implement this feature and submit a pull request. The implementation would include:

• mock_nullifzero function in src/snowflake/snowpark/mock/_functions.py
• test_nullifzero test case in tests/mock/test_functions.py

[sfc-gh-sghosh]

Hi @mkanetsuna ,

Thanks for raising the issue and for the PR request.

Yes, its has not been implemented as Build in support like standard functions like col, lit, so you need to either below way or use patch decorator as mentioned in the error message.

session = Session.builder.config('local_testing', True).create()
print("local session created")

df = session.create_dataframe([[0], [1], [100]], schema=["value"])
#result = df.select(nullifzero(col("value")))
result = df.select(F.when(col("value") == 0, None).otherwise(col("value")))
result.show()

output
session created
local session created
------------------------------------------------------
|"CASE  WHEN (""VALUE"" = 0 :: INT) THEN NULL EL...  |
------------------------------------------------------
|NULL                                                |
|1                                                   |
|100                                                 |
------------------------------------------------------


or

from snowflake.snowpark.mock import patch
import snowflake.snowpark.functions as F
from snowflake.snowpark.functions import col

@patch(F.nullifzero)
def mock_nullifzero(column):
    # In Local Testing, 'column' is a ColumnEmulator (Series).
    return column.apply(lambda x: None if x == 0 else x)

df = session.create_dataframe([[0], [1], [100]], schema=["value"])
result = df.select(F.nullifzero(col("value")))
result.show()

---------------------------
|"NULLIFZERO(""VALUE"")"  |
---------------------------
|nan                      |
|1.0                      |
|100.0                    |
---------------------------

The PR request will be reviewed and approve accordingly.

[mkanetsuna]

Hi @sfc-gh-sghosh ,

Thank you for the feedback and for showing me the alternative approaches!
I look forward to your review of the PR.