Executor for polars logical plans

dependencies.yaml

@@ -603,7 +603,7 @@ dependencies:
     common:
       - output_types: [conda, requirements, pyproject]
         packages:
-          - polars>=0.20.24
+          - polars>=0.20.30
   run_dask_cudf:
     common:
       - output_types: [conda, requirements, pyproject]

python/cudf/cudf/_lib/pylibcudf/types.pyx

@@ -51,6 +51,9 @@ cdef class DataType:
             self.c_obj == (<DataType>other).c_obj
         )
 
+    def __hash__(self):
+        return hash((self.c_obj.id(), self.c_obj.scale()))
+
     @staticmethod
     cdef DataType from_libcudf(data_type dt):
         """Create a DataType from a libcudf data_type.

python/cudf_polars/cudf_polars/callback.py

@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Callback for the polars collect function to execute on device."""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import TYPE_CHECKING
+
+import nvtx
+
+from cudf_polars.dsl.translate import translate_ir
+
+if TYPE_CHECKING:
+    import polars as pl
+
+    from cudf_polars.dsl.ir import IR
+
+__all__: list[str] = ["execute_with_cudf"]
+
+
+def _callback(
+    ir: IR,
+    with_columns: list[str] | None,
+    pyarrow_predicate: str | None,
+    n_rows: int | None,
+) -> pl.DataFrame:
+    assert with_columns is None
+    assert pyarrow_predicate is None
+    assert n_rows is None
+    with nvtx.annotate(message="ExecuteIR", domain="cudf_polars"):
+        return ir.evaluate(cache={}).to_polars()
+
+
+def execute_with_cudf(nt, *, raise_on_fail: bool = False) -> None:
+    """
+    A post optimization callback that attempts to execute the plan with cudf.
+
+    Parameters
+    ----------
+    nt
+        NodeTraverser
+
+    raise_on_fail
+        Should conversion raise an exception rather than continuing
+        without setting a callback.
+
+    The NodeTraverser is mutated if the libcudf executor can handle the plan.
+    """
+    try:
+        with nvtx.annotate(message="ConvertIR", domain="cudf_polars"):
+            nt.set_udf(partial(_callback, translate_ir(nt)))
+    except NotImplementedError:
+        if raise_on_fail:
+            raise

python/cudf_polars/cudf_polars/containers/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Containers of concrete data."""
+
+from __future__ import annotations
+
+__all__: list[str] = ["DataFrame", "Column", "Scalar"]
+
+from cudf_polars.containers.column import Column
+from cudf_polars.containers.dataframe import DataFrame
+from cudf_polars.containers.scalar import Scalar

python/cudf_polars/cudf_polars/containers/column.py

@@ -0,0 +1,119 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""A column, with some properties."""
+
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING
+
+import cudf._lib.pylibcudf as plc
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+__all__: list[str] = ["Column"]
+
+
+class Column:
+    """A column, a name, and sortedness."""
+
+    obj: plc.Column
+    name: str
+    is_sorted: plc.types.Sorted
+    order: plc.types.Order
+    null_order: plc.types.NullOrder
+
+    def __init__(self, column: plc.Column, name: str):
+        self.obj = column
+        self.name = name
+        self.is_sorted = plc.types.Sorted.NO
+        self.order = plc.types.Order.ASCENDING
+        self.null_order = plc.types.NullOrder.BEFORE
+
+    def sorted_like(self, like: Column, /) -> Self:
+        """
+        Copy sortedness properties from a column onto self.
+
+        Parameters
+        ----------
+        like
+            The column to copy sortedness metadata from.
+
+        Returns
+        -------
+        Self with metadata set.
+
+        See Also
+        --------
+        set_sorted
+        """
+        return self.set_sorted(
+            is_sorted=like.is_sorted, order=like.order, null_order=like.null_order
+        )
+
+    def set_sorted(
+        self,
+        *,
+        is_sorted: plc.types.Sorted,
+        order: plc.types.Order,
+        null_order: plc.types.NullOrder,
+    ) -> Self:
+        """
+        Modify sortedness metadata in place.
+
+        Parameters
+        ----------
+        is_sorted
+            Is the column sorted
+        order
+            The order if sorted
+        null_order
+            Where nulls sort, if sorted
+
+        Returns
+        -------
+        Self with metadata set.
+        """
+        self.is_sorted = is_sorted
+        self.order = order
+        self.null_order = null_order
+        return self
+
+    def copy(self, *, new_name: str | None = None) -> Self:
+        """
+        Return a shallow copy of the column.
+
+        Parameters
+        ----------
+        new_name
+            Optional new name for the copied column.
+
+        Returns
+        -------
+        New column sharing data with self.
+        """
+        return type(self)(
+            self.obj, self.name if new_name is None else new_name
+        ).sorted_like(self)
+
+    def mask_nans(self) -> Self:
+        """Return a copy of self with nans masked out."""
+        if self.nan_count > 0:
+            raise NotImplementedError
+        return self.copy()
+
+    @functools.cached_property
+    def nan_count(self) -> int:
+        """Return the number of NaN values in the column."""
+        if self.obj.type().id() not in (plc.TypeId.FLOAT32, plc.TypeId.FLOAT64):
+            return 0
+        return plc.interop.to_arrow(
+            plc.reduce.reduce(
+                plc.unary.is_nan(self.obj),
+                plc.aggregation.sum(),
+                # TODO: pylibcudf needs to have a SizeType DataType singleton
+                plc.DataType(plc.TypeId.INT32),
+            )
+        ).as_py()