[DNM][WIP] Single-partition Dask executor for cuDF-Polars #17262

rjzamora · 2024-11-07T01:44:37Z

WARNING: This is still very rough.

Description

The goal here is to lay down the initial foundation for dask-based evaluation of IR graphs in cudf-polars. The first pass will only support single-partition workloads. This functionality could be achieved with much less-complicated changes to cudf-polars. However, we do want to build multi-partition support on top of this.

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.

rjzamora · 2024-11-07T01:48:21Z

python/cudf_polars/cudf_polars/experimental/single.py

+
+@_single_partition_node.register
+def _(node: PythonScan, rec) -> SParPythonScan:
+    return SParPythonScan(*node._ctor_arguments(map(rec, node.children)))


@wence- - This PR is still rough, so it may not make sense to debug type-hinting problems yet. That said, I'm getting a ton of errors that I don't know how to address like: error: Argument 1 to "SParPythonScan" has incompatible type "*Sequence[Any | IR]"; expected "Mapping[str, Any]" [arg-type]

Ah, yeah, can you apply:

diff --git a/python/cudf_polars/cudf_polars/dsl/nodebase.py b/python/cudf_polars/cudf_polars/dsl/nodebase.py index 228d300f46..c96414ad3c 100644 --- a/python/cudf_polars/cudf_polars/dsl/nodebase.py +++ b/python/cudf_polars/cudf_polars/dsl/nodebase.py @@ -40,7 +40,7 @@ class Node(Generic[T]): children: tuple[T, ...] _non_child: ClassVar[tuple[str, ...]] = () - def _ctor_arguments(self, children: Sequence[T]) -> Sequence[Any | T]: + def _ctor_arguments(self, children: Sequence[T]) -> Sequence[Any]: return (*(getattr(self, attr) for attr in self._non_child), *children) def reconstruct(

But you should use the reconstruct method:

Then you can do X.reconstruct(map(rec, node.children)) I suspect.

I probably need to fiddle with the types a bit, perhaps this:

diff --git a/python/cudf_polars/cudf_polars/dsl/nodebase.py b/python/cudf_polars/cudf_polars/dsl/nodebase.py index 228d300f46..9fc855a764 100644 --- a/python/cudf_polars/cudf_polars/dsl/nodebase.py +++ b/python/cudf_polars/cudf_polars/dsl/nodebase.py @@ -8,7 +8,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar if TYPE_CHECKING: - from collections.abc import Hashable, Sequence + from collections.abc import Hashable, Iterable, Sequence from typing_extensions import Self @@ -40,12 +40,12 @@ class Node(Generic[T]): children: tuple[T, ...] _non_child: ClassVar[tuple[str, ...]] = () - def _ctor_arguments(self, children: Sequence[T]) -> Sequence[Any | T]: + def _ctor_arguments(self, children: Iterable[T]) -> Sequence[Any]: return (*(getattr(self, attr) for attr in self._non_child), *children) def reconstruct( - self, children: Sequence[T] - ) -> Self: # pragma: no cover; not yet used + self, children: Iterable[T] + ) -> Self: """ Rebuild this node with new children.

rjzamora · 2024-11-07T01:49:14Z

python/cudf_polars/cudf_polars/experimental/single.py

+    return SParHConcat(*node._ctor_arguments(map(rec, node.children)))
+
+
+lower_ir_graph = CachingVisitor(_single_partition_node)


Not sure if I'm using CachingVisitor correctly (or how to deal with the type hinting).

You want:

def lower_ir_graph(node: ir.IR) -> ...: mapper = CachingVisitor(_single_partition_node) return mapper(node)

The cache lives for the lifetime of the mapper, so if you make this module-scope, then you'll hold on to everything forever.

wence- · 2024-11-07T14:42:14Z

python/cudf_polars/cudf_polars/callback.py

+        if os.environ.get("CUDF_POLARS_DASK", "OFF").upper() == "ON":
+            # Use experimental Dask executor
+            from cudf_polars.experimental.parallel import evaluate_dask
+
+            return evaluate_dask(ir).to_polars()


note: You can put arbitrary keyword arguments into the pl.GPUEngine object and handle them here

wence- · 2024-11-07T14:44:14Z

python/cudf_polars/cudf_polars/dsl/ir.py

+
+    def get_hashable(self) -> Hashable:
+        """Hashable representation of the node."""
+        schema_hash = tuple(self.schema.items())
+        return (type(self), schema_hash, self.name, str(self.options))


No, this is wrong, because it doesn't consider the children. What's wrong with the default?

wence- · 2024-11-07T14:44:44Z

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -1514,15 +1521,18 @@ def do_evaluate(cls, name: str, options: Any, df: DataFrame) -> DataFrame:
                        plc.interop.from_arrow(
                            pa.array(
                                pivotees,
-                                type=plc.interop.to_arrow(variable_dtype),
+                                type=plc.interop.to_arrow(schema[variable_name]),


Why, is it you need the schema elsewhere?

wence- · 2024-11-07T14:48:26Z

python/cudf_polars/cudf_polars/experimental/single.py

+    return SParHConcat(*node._ctor_arguments(map(rec, node.children)))
+
+
+lower_ir_graph = CachingVisitor(_single_partition_node)


You want:

def lower_ir_graph(node: ir.IR) -> ...: mapper = CachingVisitor(_single_partition_node) return mapper(node)

The cache lives for the lifetime of the mapper, so if you make this module-scope, then you'll hold on to everything forever.

wence- · 2024-11-07T14:53:15Z

python/cudf_polars/cudf_polars/experimental/single.py

+
+@_single_partition_node.register
+def _(node: PythonScan, rec) -> SParPythonScan:
+    return SParPythonScan(*node._ctor_arguments(map(rec, node.children)))


Ah, yeah, can you apply:

diff --git a/python/cudf_polars/cudf_polars/dsl/nodebase.py b/python/cudf_polars/cudf_polars/dsl/nodebase.py index 228d300f46..c96414ad3c 100644 --- a/python/cudf_polars/cudf_polars/dsl/nodebase.py +++ b/python/cudf_polars/cudf_polars/dsl/nodebase.py @@ -40,7 +40,7 @@ class Node(Generic[T]): children: tuple[T, ...] _non_child: ClassVar[tuple[str, ...]] = () - def _ctor_arguments(self, children: Sequence[T]) -> Sequence[Any | T]: + def _ctor_arguments(self, children: Sequence[T]) -> Sequence[Any]: return (*(getattr(self, attr) for attr in self._non_child), *children) def reconstruct(

But you should use the reconstruct method:

Then you can do X.reconstruct(map(rec, node.children)) I suspect.

wence- · 2024-11-07T14:58:16Z

python/cudf_polars/cudf_polars/experimental/single.py

+
+@_single_partition_node.register
+def _(node: PythonScan, rec) -> SParPythonScan:
+    return SParPythonScan(*node._ctor_arguments(map(rec, node.children)))


I probably need to fiddle with the types a bit, perhaps this:

diff --git a/python/cudf_polars/cudf_polars/dsl/nodebase.py b/python/cudf_polars/cudf_polars/dsl/nodebase.py index 228d300f46..9fc855a764 100644 --- a/python/cudf_polars/cudf_polars/dsl/nodebase.py +++ b/python/cudf_polars/cudf_polars/dsl/nodebase.py @@ -8,7 +8,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar if TYPE_CHECKING: - from collections.abc import Hashable, Sequence + from collections.abc import Hashable, Iterable, Sequence from typing_extensions import Self @@ -40,12 +40,12 @@ class Node(Generic[T]): children: tuple[T, ...] _non_child: ClassVar[tuple[str, ...]] = () - def _ctor_arguments(self, children: Sequence[T]) -> Sequence[Any | T]: + def _ctor_arguments(self, children: Iterable[T]) -> Sequence[Any]: return (*(getattr(self, attr) for attr in self._non_child), *children) def reconstruct( - self, children: Sequence[T] - ) -> Self: # pragma: no cover; not yet used + self, children: Iterable[T] + ) -> Self: """ Rebuild this node with new children.

rjzamora added 2 commits November 6, 2024 14:58

cleanup

a590076

rename to parallel

7f1bec7

rjzamora added 5 - DO NOT MERGE Hold off on merging; see PR for details improvement Improvement / enhancement to an existing function non-breaking Non-breaking change cudf.polars Issues specific to cudf.polars labels Nov 7, 2024

rjzamora self-assigned this Nov 7, 2024

Merge branch 'branch-24.12' into cudf-polars-dask-simple

023e085

github-actions bot added the Python Affects Python cuDF API. label Nov 7, 2024

rjzamora commented Nov 7, 2024

View reviewed changes

rjzamora added 2 - In Progress Currently a work in progress and removed 5 - DO NOT MERGE Hold off on merging; see PR for details labels Nov 7, 2024

Merge branch 'branch-24.12' into cudf-polars-dask-simple

e7a2fce

wence- reviewed Nov 7, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[DNM][WIP] Single-partition Dask executor for cuDF-Polars #17262

[DNM][WIP] Single-partition Dask executor for cuDF-Polars #17262

rjzamora commented Nov 7, 2024

rjzamora Nov 7, 2024

wence- Nov 7, 2024

wence- Nov 7, 2024

rjzamora Nov 7, 2024

wence- Nov 7, 2024

wence- Nov 7, 2024

wence- Nov 7, 2024

wence- Nov 7, 2024

wence- Nov 7, 2024

wence- Nov 7, 2024

wence- Nov 7, 2024

		return SParHConcat(*node._ctor_arguments(map(rec, node.children)))


		lower_ir_graph = CachingVisitor(_single_partition_node)

[DNM][WIP] Single-partition Dask executor for cuDF-Polars #17262

Are you sure you want to change the base?

[DNM][WIP] Single-partition Dask executor for cuDF-Polars #17262

Conversation

rjzamora commented Nov 7, 2024

Description

Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment