Raise errors on specific types of fallback in cudf.pandas

python/cudf/cudf/pandas/fast_slow_proxy.py

@@ -16,6 +16,8 @@
 
 import numpy as np
 
+from rmm import RMMError
+
 from ..options import _env_get_bool
 from ..testing import assert_eq
 from .annotation import nvtx
@@ -881,12 +883,36 @@ def _assert_fast_slow_eq(left, right):
         assert_eq(left, right)
 
 
-class ProxyFallbackError(Exception):
+class FallbackError(Exception):
     """Raised when fallback occurs"""
 
     pass
 
 
+class OOMFallbackError(FallbackError):
+    """Warns when cuDF produces a MemoryError or an rmm.RMMError"""
+
+    pass
+
+
+class NotImplementedFallbackError(FallbackError):
+    """Warns cuDF produces a NotImplementedError"""
+
+    pass
+
+
+class AttributeFallbackError(FallbackError):
+    """Warns when cuDF produces an AttributeError"""
+
+    pass
+
+
+class TypeFallbackError(FallbackError):
+    """Warns when cuDF produces a TypeError"""
+
+    pass
+
+
 def _fast_function_call():
     """
     Placeholder fast function for pytest profiling purposes.
@@ -964,7 +990,27 @@ def _fast_slow_function_call(
                         )
     except Exception as err:
         if _env_get_bool("CUDF_PANDAS_FAIL_ON_FALLBACK", False):
-            raise ProxyFallbackError(
+            if isinstance(err, (RMMError, MemoryError)):
+                raise OOMFallbackError(
+                    "Out of Memory Error. Falling back to the slow path. "
+                    f"The exception was {err}."
+                )
+            elif isinstance(err, NotImplementedError):
+                raise NotImplementedFallbackError(
+                    "NotImplementedError. Falling back to the slow path. "
+                    f"The exception was {err}."
+                )
+            elif isinstance(err, AttributeError):
+                raise AttributeFallbackError(
+                    "AttributeError. Falling back to the slow path. "
+                    f"The exception was {err}."
+                )
+            elif isinstance(err, TypeError):
+                raise TypeFallbackError(
+                    "TypeError. Falling back to the slow path. "
+                    f"The exception was {err}."
+                )
+            raise FallbackError(
                 f"The operation failed with cuDF, the reason was {type(err)}: {err}"
             ) from err
         with nvtx.annotate(

python/cudf/cudf_pandas_tests/test_cudf_pandas.py

@@ -30,10 +30,16 @@
 from packaging import version
 from pytz import utc
 
+from rmm import RMMError
+
 from cudf.core._compat import PANDAS_GE_210, PANDAS_GE_220, PANDAS_VERSION
 from cudf.pandas import LOADED, Profiler
 from cudf.pandas.fast_slow_proxy import (
-    ProxyFallbackError,
+    AttributeFallbackError,
+    FallbackError,
+    NotImplementedFallbackError,
+    OOMFallbackError,
+    TypeFallbackError,
     _Unusable,
     is_proxy_object,
 )
@@ -1757,10 +1763,77 @@ def add_one_ufunc(a):
 def test_fallback_raises_error(monkeypatch):
     with monkeypatch.context() as monkeycontext:
         monkeycontext.setenv("CUDF_PANDAS_FAIL_ON_FALLBACK", "True")
-        with pytest.raises(ProxyFallbackError):
+        with pytest.raises(FallbackError):
             pd.Series(range(2)).astype(object)
 
 
+def mock_mean_memory_error(self, *args, **kwargs):
+    raise MemoryError()
+
+
+def mock_mean_rmm_error(self, *args, **kwargs):
+    raise RMMError(1, "error")
+
+
+def mock_mean_not_impl_error(self, *args, **kwargs):
+    raise NotImplementedError()
+
+
+def mock_mean_attr_error(self, *args, **kwargs):
+    raise AttributeError()
+
+
+def mock_mean_type_error(self, *args, **kwargs):
+    raise TypeError()
+
+
+@pytest.mark.parametrize(
+    "mock_mean, err, match_str",
+    [
+        (
+            mock_mean_memory_error,
+            OOMFallbackError,
+            "Out of Memory Error.",
+        ),
+        (
+            mock_mean_rmm_error,
+            OOMFallbackError,
+            "Out of Memory Error.",
+        ),
+        (
+            mock_mean_not_impl_error,
+            NotImplementedFallbackError,
+            "NotImplementedError.",
+        ),
+        (
+            mock_mean_attr_error,
+            AttributeFallbackError,
+            "AttributeError.",
+        ),
+        (
+            mock_mean_type_error,
+            TypeFallbackError,
+            "TypeError.",
+        ),
+    ],
+)
+def test_fallback_raises_specific_error(
+    monkeypatch,
+    mock_mean,
+    err,
+    match_str,
+):
+    with monkeypatch.context() as monkeycontext:
+        monkeypatch.setattr(xpd.Series.mean, "_fsproxy_fast", mock_mean)
+        monkeycontext.setenv("CUDF_PANDAS_FAIL_ON_FALLBACK", "True")
+        s = xpd.Series([1, 2])
+        with pytest.raises(err, match=match_str):
+            assert s.mean() == 1.5
+
+    # Must explicitly undo the patch. Proxy dispatch doesn't work with monkeypatch contexts.
+    monkeypatch.setattr(xpd.Series.mean, "_fsproxy_fast", cudf.Series.mean)
+
+
 @pytest.mark.parametrize(
     "attrs",
     [