diff --git a/src/snowflake/snowpark/modin/plugin/extensions/dataframe_overrides.py b/src/snowflake/snowpark/modin/plugin/extensions/dataframe_overrides.py index d6ceffd8c47..ac2a5d02329 100644 --- a/src/snowflake/snowpark/modin/plugin/extensions/dataframe_overrides.py +++ b/src/snowflake/snowpark/modin/plugin/extensions/dataframe_overrides.py @@ -140,7 +140,7 @@ def decorator(base_method: Any): def _map(self, func: PythonFuncType, na_action: str | None = None, **kwargs): # TODO: SNOW-1063346: Modin upgrade - modin.pandas.DataFrame functions if not callable(func): - raise TypeError(f"{func} is not callable") + raise TypeError(f"{func} is not callable") # pragma: no cover return self.__constructor__( query_compiler=self._query_compiler.applymap( func, na_action=na_action, **kwargs diff --git a/tests/integ/modin/frame/test_applymap.py b/tests/integ/modin/frame/test_applymap.py index 8d992ff7536..bf890fbdf21 100644 --- a/tests/integ/modin/frame/test_applymap.py +++ b/tests/integ/modin/frame/test_applymap.py @@ -134,13 +134,13 @@ def test_applymap_na_action_ignore(method): @pytest.mark.parametrize("invalid_input", ["min", [np.min], {"a": np.max}]) @sql_count_checker(query_count=0) -def test_applymap_invalid_input(invalid_input): +def test_applymap_invalid_input(invalid_input, method): snow_df = pd.DataFrame([1]) native_df = native_pd.DataFrame([1]) eval_snowpark_pandas_result( snow_df, native_df, - lambda x: x.applymap(invalid_input), + lambda x: getattr(x, method)(invalid_input), expect_exception=True, expect_exception_match="is not callable", assert_exception_equal=False,