Skip to content

Commit

Permalink
move map to numpy_to_pandas
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-nkumar committed Dec 13, 2024
1 parent 346bcd4 commit 810cd76
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -185,7 +185,6 @@
APPLY_LABEL_COLUMN_QUOTED_IDENTIFIER,
APPLY_VALUE_COLUMN_QUOTED_IDENTIFIER,
DEFAULT_UDTF_PARTITION_SIZE,
NUMPY_FUNCTION_TO_SNOWFLAKE_FUNCTION,
GroupbyApplySortMethod,
check_return_variant_and_get_return_type,
create_udf_for_series_apply,
Expand Down Expand Up @@ -373,6 +372,9 @@
from snowflake.snowpark.modin.plugin.utils.error_message import ErrorMessage
from snowflake.snowpark.modin.plugin.utils.warning_message import WarningMessage
from snowflake.snowpark.modin.utils import MODIN_UNNAMED_SERIES_LABEL
from snowflake.snowpark.modin.plugin.utils.numpy_to_pandas import (
NUMPY_FUNCTION_TO_SNOWFLAKE_FUNCTION,
)
from snowflake.snowpark.session import Session
from snowflake.snowpark.types import (
ArrayType,
Expand Down Expand Up @@ -8760,7 +8762,10 @@ def applymap(
# Check if the function is a known numpy function that can be translated to Snowflake function.
sf_func = NUMPY_FUNCTION_TO_SNOWFLAKE_FUNCTION.get(func)
if sf_func is not None:
return self._apply_snowpark_python_function_to_columns(sf_func, kwargs)
# TODO SNOW-1739034: remove pragma no cover when apply tests are enabled in CI
return self._apply_snowpark_python_function_to_columns(
sf_func, kwargs
) # pragma: no cover

# Currently, NULL values are always passed into the udtf even if strict=True,
# which is a bug on the server side SNOW-880105.
Expand Down
46 changes: 46 additions & 0 deletions src/snowflake/snowpark/modin/plugin/utils/numpy_to_pandas.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import modin.pandas as pd
from modin.pandas.base import BasePandasDataset
from modin.pandas.utils import is_scalar
import numpy as np

from snowflake.snowpark import functions as sp_func
from snowflake.snowpark.modin.plugin.utils.warning_message import WarningMessage
Expand Down Expand Up @@ -288,3 +289,48 @@ def map_to_bools(inputs: Any) -> Any:
sp_func.trunc
), # df.truncate not supported in snowpandas yet
}


NUMPY_FUNCTION_TO_SNOWFLAKE_FUNCTION = {
# Math operations
np.absolute: sp_func.abs,
np.sign: sp_func.sign,
np.negative: sp_func.negate,
np.positive: lambda col: col,
np.sqrt: sp_func.sqrt,
np.square: lambda col: sp_func.builtin("square")(col),
np.cbrt: lambda col: sp_func.builtin("cbrt")(col),
np.reciprocal: lambda col: 1 / col,
np.exp: sp_func.exp,
np.exp2: lambda col: sp_func.pow(2, col),
np.expm1: lambda col: sp_func.exp(col) - 1,
np.log: sp_func.ln,
np.log2: sp_func._log2,
np.log10: sp_func._log10,
np.log1p: lambda col: sp_func.ln(col + 1),
# Aggregate functions translate to identity functions when applied element wise
np.sum: lambda col: col,
np.min: lambda col: col,
np.max: lambda col: col,
# Trigonometric functions
np.sin: sp_func.sin,
np.cos: sp_func.cos,
np.tan: sp_func.tan,
np.sinh: sp_func.sinh,
np.cosh: sp_func.cosh,
np.tanh: sp_func.tanh,
np.arcsin: lambda col: sp_func.builtin("asin")(col),
np.arccos: lambda col: sp_func.builtin("acos")(col),
np.arctan: lambda col: sp_func.builtin("atan")(col),
np.arctan2: lambda col: sp_func.builtin("atan2")(col),
np.arcsinh: lambda col: sp_func.builtin("asinh")(col),
np.arccosh: lambda col: sp_func.builtin("acosh")(col),
np.arctanh: lambda col: sp_func.builtin("atanh")(col),
np.degrees: lambda col: sp_func.builtin("degrees")(col),
np.radians: lambda col: sp_func.builtin("radians")(col),
# Floating functions
np.ceil: sp_func.ceil,
np.floor: sp_func.floor,
np.trunc: sp_func.trunc,
np.isnan: sp_func.is_null,
}

0 comments on commit 810cd76

Please sign in to comment.