pass a class that has a NodeTraverser instead

rapidsai · Oct 9, 2024 · 3175a7e · 3175a7e
1 parent ded4dd2
commit 3175a7e
Show file tree

Hide file tree

Showing 2 changed files with 217 additions and 112 deletions.
diff --git a/python/cudf_polars/cudf_polars/callback.py b/python/cudf_polars/cudf_polars/callback.py
@@ -18,7 +18,7 @@
 import rmm
 from rmm._cuda import gpu
 
-from cudf_polars.dsl.translate import translate_ir
+from cudf_polars.dsl.translate import translate_ir, Translator
 
 if TYPE_CHECKING:
     from collections.abc import Generator
@@ -174,16 +174,22 @@ def execute_with_cudf(
     device = config.device
     memory_resource = config.memory_resource
     raise_on_fail = config.config.get("raise_on_fail", False)
-    if unsupported := (config.config.keys() - {"raise_on_fail"}):
+    debug_mode = config.config.get("debug_mode", False)
+    if unsupported := (config.config.keys() - {"raise_on_fail", "debug_mode"}):
         raise ValueError(
             f"Engine configuration contains unsupported settings {unsupported}"
         )
     try:
         with nvtx.annotate(message="ConvertIR", domain="cudf_polars"):
+            translator = Translator(nt, debug_mode=debug_mode)
+            translation = translate_ir(translator)
+            if debug_mode:
+                print(set(translator.errors))
+                raise NotImplementedError("Query contained unsupported operations")
             nt.set_udf(
                 partial(
                     _callback,
-                    translate_ir(nt),
+                    translation,
                     device=device,
                     memory_resource=memory_resource,
                 )