diff --git a/CHANGELOG.md b/CHANGELOG.md index 7f519a4169a..182cc2dea17 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/source/snowpark/functions.rst b/docs/source/snowpark/functions.rst index b049c800838..0a46155588f 100644 --- a/docs/source/snowpark/functions.rst +++ b/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 diff --git a/src/snowflake/snowpark/functions.py b/src/snowflake/snowpark/functions.py index 76001e342b1..ca14ae4ca35 100644 --- a/src/snowflake/snowpark/functions.py +++ b/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.