SIT-2898: Add divnull and nullifzero (#2753)

CHANGELOG.md

@@ -6,7 +6,10 @@
 
 #### New Features
 
-- Added support for function `snowflake_cortex_sentiment` in `functions.py`.
+- Added support for the following functions in `functions.py`
+  - `divnull`
+  - `nullifzero`
+  - `snowflake_cortex_sentiment`
 
 ### Snowpark pandas API Updates
 

docs/source/snowpark/functions.rst

@@ -143,6 +143,7 @@ Functions
     desc_nulls_first
     desc_nulls_last
     div0
+    divnull
     endswith
     equal_nan
     exp
@@ -219,6 +220,7 @@ Functions
     next_day
     not_
     ntile
+    nullifzero
     object_agg
     object_construct
     object_construct_keep_null

src/snowflake/snowpark/functions.py

@@ -2014,6 +2014,47 @@ def div0(
     return builtin("div0", _emit_ast=_emit_ast)(dividend_col, divisor_col)
 
 
+@publicapi
+def divnull(
+    dividend: Union[ColumnOrName, int, float],
+    divisor: Union[ColumnOrName, int, float],
+    _emit_ast: bool = True,
+) -> Column:
+    """Performs division like the division operator (/),
+    but returns NULL when the divisor is 0 (rather then reporting error).
+
+    Example::
+
+        >>> df = session.create_dataframe([1], schema=["a"])
+        >>> df.select(divnull(df["a"], 1).alias("divided_by_one"), divnull(df["a"], 0).alias("divided_by_zero")).collect()
+        [Row(DIVIDED_BY_ONE=Decimal('1.000000'), DIVIDED_BY_ZERO=None)]
+    """
+    dividend_col = (
+        lit(dividend, _emit_ast=False)
+        if isinstance(dividend, (int, float))
+        else _to_col_if_str(dividend, "divnull")
+    )
+    divisor_col = (
+        lit(divisor, _emit_ast=False)
+        if isinstance(divisor, (int, float))
+        else _to_col_if_str(divisor, "divnull")
+    )
+    return dividend_col / nullifzero(divisor_col, _emit_ast=False)
+
+
+@publicapi
+def nullifzero(e: ColumnOrName, _emit_ast: bool = True) -> Column:
+    """Returns NULL if the argument evaluates to 0; otherwise, returns the argument.
+
+    Example::
+        >>> df = session.create_dataframe([0, 1], schema=["a"])
+        >>> df.select(nullifzero(df["a"]).alias("result")).collect()
+        [Row(RESULT=None), Row(RESULT=1)]
+    """
+    c = _to_col_if_str(e, "nullifzero")
+    return builtin("nullifzero", _emit_ast=_emit_ast)(c)
+
+
 @publicapi
 def sqrt(e: ColumnOrName, _emit_ast: bool = True) -> Column:
     """Returns the square-root of a non-negative numeric expression.