Add context to IR.do_evaluate (#20322)

This adds a keyword-only `context` argument to cudf_polars IR.do_evaluate method. The purpose to provide access to special pieces of data that might be necessary for controlling an IR nodes' execution, but doesn't belong on the IR node itself as a non-child argument. Specifically, we'd like to provide a CUDA `stream` argument as part of #20228, but we generalize that slightly and provide a system for providing arbitrary data.

A few notes on the implementation:

- For now, the context is just an empty dataclass. I suspect its design might change in the future.
- I've opted to push the creation of the context as high as possible. For now it's created in `_callback` and passed into `ir.evaluate` / `evaluate_streaming` and from there to all the methods that require it.
- There's some awkwardness between how our IR nodes and Dask's task graph treat arguments. I've opted to make `context` keyword only in `IR.do_evaluate(..., context)`. However, Dask's task graph doesn't really deal with that. It wants a tuple of `(function, arg1, arg2, ...)`. So that requires using `functools.partial(function, context=context)(arg1, arg2, ...)`.
- After implementing this, I realized that `Expr.evaluate` *also* takes a context, and its a different type `ExecutionContext` :( I can rename the IR variant if we want.

Just a draft for now, and probably not worth reviewing until I have a branch somewhere that combines CUDA streams with this to verify it meets our needs.

Authors:
  - Tom Augspurger (https://github.com/TomAugspurger)

Approvers:
  - Matthew Murray (https://github.com/Matt711)

URL: #20322

Author: TomAugspurger

python/cudf_polars/cudf_polars/callback.py

@@ -23,6 +23,7 @@
 from rmm._cuda import gpu
 
 import cudf_polars.dsl.tracing
+from cudf_polars.dsl.ir import IRExecutionContext
 from cudf_polars.dsl.tracing import CUDF_POLARS_NVTX_DOMAIN
 from cudf_polars.dsl.translate import Translator
 from cudf_polars.utils.config import _env_get_int, get_total_device_memory
@@ -218,14 +219,17 @@ def _callback(
     assert n_rows is None
     if timer is not None:
         assert should_time
+
+    context = IRExecutionContext()
+
     with (
         nvtx.annotate(message="ExecuteIR", domain=CUDF_POLARS_NVTX_DOMAIN),
         # Device must be set before memory resource is obtained.
         set_device(config_options.device),
         set_memory_resource(memory_resource),
     ):
         if config_options.executor.name == "in-memory":
-            df = ir.evaluate(cache={}, timer=timer).to_polars()
+            df = ir.evaluate(cache={}, timer=timer, context=context).to_polars()
             if timer is None:
                 return df
             else:
@@ -243,7 +247,7 @@ def _callback(
                     """)
                 raise NotImplementedError(msg)
 
-            return evaluate_streaming(ir, config_options).to_polars()
+            return evaluate_streaming(ir, config_options, context=context).to_polars()
         assert_never(f"Unknown executor '{config_options.executor}'")
 
 

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -17,6 +17,7 @@
 import json
 import random
 import time
+from dataclasses import dataclass
 from functools import cache
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, overload
@@ -72,6 +73,7 @@
     "GroupBy",
     "HConcat",
     "HStack",
+    "IRExecutionContext",
     "Join",
     "MapFunction",
     "MergeSorted",
@@ -88,6 +90,16 @@
 ]
 
 
+@dataclass(frozen=True)
+class IRExecutionContext:
+    """
+    Runtime context for IR node execution.
+
+    This dataclass holds runtime information and configuration needed
+    during the evaluation of IR nodes.
+    """
+
+
 _BINOPS = {
     plc.binaryop.BinaryOperator.EQUAL,
     plc.binaryop.BinaryOperator.NOT_EQUAL,
@@ -158,7 +170,9 @@ def get_hashable(self) -> Hashable:
         translation phase should fail earlier.
     """
 
-    def evaluate(self, *, cache: CSECache, timer: Timer | None) -> DataFrame:
+    def evaluate(
+        self, *, cache: CSECache, timer: Timer | None, context: IRExecutionContext
+    ) -> DataFrame:
         """
         Evaluate the node (recursively) and return a dataframe.
 
@@ -170,6 +184,8 @@ def evaluate(self, *, cache: CSECache, timer: Timer | None) -> DataFrame:
         timer
             If not None, a Timer object to record timings for the
             evaluation of the node.
+        context
+            The execution context for the node.
 
         Notes
         -----
@@ -188,16 +204,19 @@ def evaluate(self, *, cache: CSECache, timer: Timer | None) -> DataFrame:
             If evaluation fails. Ideally this should not occur, since the
             translation phase should fail earlier.
         """
-        children = [child.evaluate(cache=cache, timer=timer) for child in self.children]
+        children = [
+            child.evaluate(cache=cache, timer=timer, context=context)
+            for child in self.children
+        ]
         if timer is not None:
             start = time.monotonic_ns()
-            result = self.do_evaluate(*self._non_child_args, *children)
+            result = self.do_evaluate(*self._non_child_args, *children, context=context)
             end = time.monotonic_ns()
             # TODO: Set better names on each class object.
             timer.store(start, end, type(self).__name__)
             return result
         else:
-            return self.do_evaluate(*self._non_child_args, *children)
+            return self.do_evaluate(*self._non_child_args, *children, context=context)
 
 
 class ErrorNode(IR):
@@ -587,6 +606,8 @@ def do_evaluate(
         include_file_paths: str | None,
         predicate: expr.NamedExpr | None,
         parquet_options: ParquetOptions,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         stream = get_cuda_stream()
@@ -1111,6 +1132,8 @@ def do_evaluate(
         parquet_options: ParquetOptions,
         options: dict[str, Any],
         df: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Write the dataframe to a file."""
         target = plc.io.SinkInfo([path])
@@ -1164,22 +1187,29 @@ def is_equal(self, other: Self) -> bool:  # noqa: D102
     @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Cache")
     def do_evaluate(
-        cls, key: int, refcount: int | None, df: DataFrame
+        cls,
+        key: int,
+        refcount: int | None,
+        df: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:  # pragma: no cover; basic evaluation never calls this
         """Evaluate and return a dataframe."""
         # Our value has already been computed for us, so let's just
         # return it.
         return df
 
-    def evaluate(self, *, cache: CSECache, timer: Timer | None) -> DataFrame:
+    def evaluate(
+        self, *, cache: CSECache, timer: Timer | None, context: IRExecutionContext
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         # We must override the recursion scheme because we don't want
         # to recurse if we're in the cache.
         try:
             (result, hits) = cache[self.key]
         except KeyError:
             (value,) = self.children
-            result = value.evaluate(cache=cache, timer=timer)
+            result = value.evaluate(cache=cache, timer=timer, context=context)
             cache[self.key] = (result, 0)
             return result
         else:
@@ -1249,6 +1279,8 @@ def do_evaluate(
         schema: Schema,
         df: Any,
         projection: tuple[str, ...] | None,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         if projection is not None:
@@ -1309,6 +1341,8 @@ def do_evaluate(
         exprs: tuple[expr.NamedExpr, ...],
         should_broadcast: bool,  # noqa: FBT001
         df: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         # Handle any broadcasting
@@ -1317,7 +1351,9 @@ def do_evaluate(
             columns = broadcast(*columns)
         return DataFrame(columns, stream=df.stream)
 
-    def evaluate(self, *, cache: CSECache, timer: Timer | None) -> DataFrame:
+    def evaluate(
+        self, *, cache: CSECache, timer: Timer | None, context: IRExecutionContext
+    ) -> DataFrame:
         """
         Evaluate the Select node with special handling for fast count queries.
 
@@ -1329,6 +1365,8 @@ def evaluate(self, *, cache: CSECache, timer: Timer | None) -> DataFrame:
         timer
             If not None, a Timer object to record timings for the
             evaluation of the node.
+        context
+            The execution context for the node.
 
         Returns
         -------
@@ -1364,7 +1402,7 @@ def evaluate(self, *, cache: CSECache, timer: Timer | None) -> DataFrame:
             )
             return DataFrame([col], stream=stream)
 
-        return super().evaluate(cache=cache, timer=timer)
+        return super().evaluate(cache=cache, timer=timer, context=context)
 
 
 class Reduce(IR):
@@ -1394,6 +1432,8 @@ def do_evaluate(
         cls,
         exprs: tuple[expr.NamedExpr, ...],
         df: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:  # pragma: no cover; not exposed by polars yet
         """Evaluate and return a dataframe."""
         columns = broadcast(*(e.evaluate(df) for e in exprs))
@@ -1497,6 +1537,8 @@ def do_evaluate(
         aggs: Sequence[expr.NamedExpr],
         zlice: Zlice | None,
         df: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         keys = broadcast(*(k.evaluate(df) for k in keys_in), target_length=df.num_rows)
@@ -1627,6 +1669,8 @@ def do_evaluate(
         maintain_order: bool,  # noqa: FBT001
         zlice: Zlice | None,
         df: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         keys = broadcast(*(k.evaluate(df) for k in keys_in), target_length=df.num_rows)
@@ -1904,6 +1948,8 @@ def do_evaluate(
         options: tuple,
         left: DataFrame,
         right: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         stream = get_joined_cuda_stream(upstreams=(left.stream, right.stream))
@@ -2187,6 +2233,8 @@ def do_evaluate(
         ],
         left: DataFrame,
         right: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         stream = get_joined_cuda_stream(upstreams=(left.stream, right.stream))
@@ -2358,6 +2406,8 @@ def do_evaluate(
         exprs: Sequence[expr.NamedExpr],
         should_broadcast: bool,  # noqa: FBT001
         df: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         columns = [c.evaluate(df) for c in exprs]
@@ -2426,6 +2476,8 @@ def do_evaluate(
         zlice: Zlice | None,
         stable: bool,  # noqa: FBT001
         df: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         if subset is None:
@@ -2519,6 +2571,8 @@ def do_evaluate(
         stable: bool,  # noqa: FBT001
         zlice: Zlice | None,
         df: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         sort_keys = broadcast(*(k.evaluate(df) for k in by), target_length=df.num_rows)
@@ -2565,7 +2619,9 @@ def __init__(self, schema: Schema, offset: int, length: int | None, df: IR):
     @classmethod
     @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Slice")
-    def do_evaluate(cls, offset: int, length: int, df: DataFrame) -> DataFrame:
+    def do_evaluate(
+        cls, offset: int, length: int, df: DataFrame, *, context: IRExecutionContext
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         return df.slice((offset, length))
 
@@ -2587,7 +2643,9 @@ def __init__(self, schema: Schema, mask: expr.NamedExpr, df: IR):
     @classmethod
     @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Filter")
-    def do_evaluate(cls, mask_expr: expr.NamedExpr, df: DataFrame) -> DataFrame:
+    def do_evaluate(
+        cls, mask_expr: expr.NamedExpr, df: DataFrame, *, context: IRExecutionContext
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         (mask,) = broadcast(mask_expr.evaluate(df), target_length=df.num_rows)
         return df.filter(mask)
@@ -2607,7 +2665,9 @@ def __init__(self, schema: Schema, df: IR):
     @classmethod
     @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Projection")
-    def do_evaluate(cls, schema: Schema, df: DataFrame) -> DataFrame:
+    def do_evaluate(
+        cls, schema: Schema, df: DataFrame, *, context: IRExecutionContext
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         # This can reorder things.
         columns = broadcast(
@@ -2641,7 +2701,9 @@ def __init__(self, schema: Schema, key: str, left: IR, right: IR):
     @classmethod
     @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="MergeSorted")
-    def do_evaluate(cls, key: str, *dfs: DataFrame) -> DataFrame:
+    def do_evaluate(
+        cls, key: str, *dfs: DataFrame, context: IRExecutionContext
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         stream = get_joined_cuda_stream(upstreams=(df.stream for df in dfs))
         left, right = dfs
@@ -2766,7 +2828,13 @@ def get_hashable(self) -> Hashable:
     @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="MapFunction")
     def do_evaluate(
-        cls, schema: Schema, name: str, options: Any, df: DataFrame
+        cls,
+        schema: Schema,
+        name: str,
+        options: Any,
+        df: DataFrame,
+        *,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         if name == "rechunk":
@@ -2872,7 +2940,9 @@ def __init__(self, schema: Schema, zlice: Zlice | None, *children: IR):
     @classmethod
     @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Union")
-    def do_evaluate(cls, zlice: Zlice | None, *dfs: DataFrame) -> DataFrame:
+    def do_evaluate(
+        cls, zlice: Zlice | None, *dfs: DataFrame, context: IRExecutionContext
+    ) -> DataFrame:
         """Evaluate and return a dataframe."""
         stream = get_joined_cuda_stream(upstreams=(df.stream for df in dfs))
 
@@ -2942,6 +3012,7 @@ def do_evaluate(
         cls,
         should_broadcast: bool,  # noqa: FBT001
         *dfs: DataFrame,
+        context: IRExecutionContext,
     ) -> DataFrame:
         """Evaluate and return a dataframe."""
         stream = get_joined_cuda_stream(upstreams=(df.stream for df in dfs))
@@ -2991,7 +3062,9 @@ def __init__(self, schema: Schema):
     @classmethod
     @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="Empty")
-    def do_evaluate(cls, schema: Schema) -> DataFrame:  # pragma: no cover
+    def do_evaluate(
+        cls, schema: Schema, *, context: IRExecutionContext
+    ) -> DataFrame:  # pragma: no cover
         """Evaluate and return a dataframe."""
         stream = get_cuda_stream()
         return DataFrame(

python/cudf_polars/cudf_polars/dsl/tracing.py

@@ -166,9 +166,9 @@ def wrapper(
             log = structlog.get_logger()
 
             # By convention, all non-dataframe arguments (non_child) come first.
-            # Anything remaining is a dataframe.
+            # Anything remaining is a dataframe, except for 'context' kwarg.
             frames: list[cudf_polars.containers.DataFrame] = (
-                list(args) + list(kwargs.values())
+                list(args) + [v for k, v in kwargs.items() if k != "context"]
             )[len(cls._non_child) :]  # type: ignore[assignment]
 
             before_start = time.monotonic_ns()