From 524aed4c9742bc025edbac8404191b4fb7932268 Mon Sep 17 00:00:00 2001
From: rjzamora <rzamora217@gmail.com>
Date: Wed, 18 Dec 2024 08:50:32 -0800
Subject: [PATCH] add ValueError for legacy dask-dataframe when dask_cuda is
 imported

---
 dask_cuda/__init__.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/dask_cuda/__init__.py b/dask_cuda/__init__.py
index d6ac9655..d07634f2 100644
--- a/dask_cuda/__init__.py
+++ b/dask_cuda/__init__.py
@@ -7,24 +7,28 @@
 import dask.utils
 import dask.dataframe.core
 import dask.dataframe.shuffle
-import dask.dataframe.multi
-import dask.bag.core
+from .explicit_comms.dataframe.shuffle import patch_shuffle_expression
+from dask.dataframe import DASK_EXPR_ENABLED
 from distributed.protocol.cuda import cuda_deserialize, cuda_serialize
 from distributed.protocol.serialize import dask_deserialize, dask_serialize
 
 from ._version import __git_commit__, __version__
 from .cuda_worker import CUDAWorker
-from .explicit_comms.dataframe.shuffle import (
-    patch_shuffle_expression,
-)
+
 from .local_cuda_cluster import LocalCUDACluster
 from .proxify_device_objects import proxify_decorator, unproxify_decorator
 
 
-# Monkey patching Dask to make use of explicit-comms when `DASK_EXPLICIT_COMMS=True`
-patch_shuffle_expression()
+if not DASK_EXPR_ENABLED:
+    raise ValueError(
+        "Dask-CUDA no longer supports the legacy Dask DataFrame API. "
+        "Please set the 'dataframe.query-planning' config to `True` "
+        "or None, or downgrade RAPIDS to <=24.12."
+    )
 
 
+# Monkey patching Dask to make use of explicit-comms when `DASK_EXPLICIT_COMMS=True`
+patch_shuffle_expression()
 # Monkey patching Dask to make use of proxify and unproxify in compatibility mode
 dask.dataframe.shuffle.shuffle_group = proxify_decorator(
     dask.dataframe.shuffle.shuffle_group