Merge branch 'branch-24.08' of github.com:rapidsai/cudf into pylibcud…

…f-io-writers

cpp/cmake/thirdparty/get_cucollections.cmake

@@ -1,5 +1,5 @@
 # =============================================================================
-# Copyright (c) 2021-2022, NVIDIA CORPORATION.
+# Copyright (c) 2021-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 # in compliance with the License. You may obtain a copy of the License at
@@ -15,6 +15,11 @@
 # This function finds cuCollections and performs any additional configuration.
 function(find_and_configure_cucollections)
   include(${rapids-cmake-dir}/cpm/cuco.cmake)
+  include(${rapids-cmake-dir}/cpm/package_override.cmake)
+
+  set(cudf_patch_dir "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/patches")
+  rapids_cpm_package_override("${cudf_patch_dir}/cuco_override.json")
+
   if(BUILD_SHARED_LIBS)
     rapids_cpm_cuco(BUILD_EXPORT_SET cudf-exports)
   else()

cpp/cmake/thirdparty/patches/cuco_noexcept.diff

@@ -0,0 +1,227 @@
+diff --git a/include/cuco/aow_storage.cuh b/include/cuco/aow_storage.cuh
+index 7f9de01..5228193 100644
+--- a/include/cuco/aow_storage.cuh
++++ b/include/cuco/aow_storage.cuh
+@@ -81,7 +81,7 @@ class aow_storage : public detail::aow_storage_base<T, WindowSize, Extent> {
+    * @param size Number of windows to (de)allocate
+    * @param allocator Allocator used for (de)allocating device storage
+    */
+-  explicit constexpr aow_storage(Extent size, Allocator const& allocator = {}) noexcept;
++  explicit constexpr aow_storage(Extent size, Allocator const& allocator = {});
+
+   aow_storage(aow_storage&&) = default;  ///< Move constructor
+   /**
+@@ -122,7 +122,7 @@ class aow_storage : public detail::aow_storage_base<T, WindowSize, Extent> {
+    * @param key Key to which all keys in `slots` are initialized
+    * @param stream Stream used for executing the kernel
+    */
+-  void initialize(value_type key, cuda_stream_ref stream = {}) noexcept;
++  void initialize(value_type key, cuda_stream_ref stream = {});
+
+   /**
+    * @brief Asynchronously initializes each slot in the AoW storage to contain `key`.
+diff --git a/include/cuco/detail/open_addressing/open_addressing_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_impl.cuh
+index c2c9c14..8ac4236 100644
+--- a/include/cuco/detail/open_addressing/open_addressing_impl.cuh
++++ b/include/cuco/detail/open_addressing/open_addressing_impl.cuh
+@@ -125,7 +125,7 @@ class open_addressing_impl {
+                                  KeyEqual const& pred,
+                                  ProbingScheme const& probing_scheme,
+                                  Allocator const& alloc,
+-                                 cuda_stream_ref stream) noexcept
++                                 cuda_stream_ref stream)
+     : empty_slot_sentinel_{empty_slot_sentinel},
+       erased_key_sentinel_{this->extract_key(empty_slot_sentinel)},
+       predicate_{pred},
+@@ -233,7 +233,7 @@ class open_addressing_impl {
+    *
+    * @param stream CUDA stream this operation is executed in
+    */
+-  void clear(cuda_stream_ref stream) noexcept { storage_.initialize(empty_slot_sentinel_, stream); }
++  void clear(cuda_stream_ref stream) { storage_.initialize(empty_slot_sentinel_, stream); }
+
+   /**
+    * @brief Asynchronously erases all elements from the container. After this call, `size()` returns
+@@ -599,7 +599,7 @@ class open_addressing_impl {
+    *
+    * @return The number of elements in the container
+    */
+-  [[nodiscard]] size_type size(cuda_stream_ref stream) const noexcept
++  [[nodiscard]] size_type size(cuda_stream_ref stream) const
+   {
+     auto counter =
+       detail::counter_storage<size_type, thread_scope, allocator_type>{this->allocator()};
+diff --git a/include/cuco/detail/static_map/static_map.inl b/include/cuco/detail/static_map/static_map.inl
+index e17a145..3fa1d02 100644
+--- a/include/cuco/detail/static_map/static_map.inl
++++ b/include/cuco/detail/static_map/static_map.inl
+@@ -123,7 +123,7 @@ template <class Key,
+           class Allocator,
+           class Storage>
+ void static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::clear(
+-  cuda_stream_ref stream) noexcept
++  cuda_stream_ref stream)
+ {
+   impl_->clear(stream);
+ }
+@@ -215,7 +215,7 @@ template <class Key,
+           class Storage>
+ template <typename InputIt>
+ void static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::
+-  insert_or_assign(InputIt first, InputIt last, cuda_stream_ref stream) noexcept
++  insert_or_assign(InputIt first, InputIt last, cuda_stream_ref stream)
+ {
+   return this->insert_or_assign_async(first, last, stream);
+   stream.synchronize();
+@@ -465,7 +465,7 @@ template <class Key,
+           class Storage>
+ static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::size_type
+ static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::size(
+-  cuda_stream_ref stream) const noexcept
++  cuda_stream_ref stream) const
+ {
+   return impl_->size(stream);
+ }
+diff --git a/include/cuco/detail/static_multiset/static_multiset.inl b/include/cuco/detail/static_multiset/static_multiset.inl
+index 174f9bc..582926b 100644
+--- a/include/cuco/detail/static_multiset/static_multiset.inl
++++ b/include/cuco/detail/static_multiset/static_multiset.inl
+@@ -97,7 +97,7 @@ template <class Key,
+           class Allocator,
+           class Storage>
+ void static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::clear(
+-  cuda_stream_ref stream) noexcept
++  cuda_stream_ref stream)
+ {
+   impl_->clear(stream);
+ }
+@@ -183,7 +183,7 @@ template <class Key,
+           class Storage>
+ static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::size_type
+ static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::size(
+-  cuda_stream_ref stream) const noexcept
++  cuda_stream_ref stream) const
+ {
+   return impl_->size(stream);
+ }
+diff --git a/include/cuco/detail/static_set/static_set.inl b/include/cuco/detail/static_set/static_set.inl
+index 645013f..d3cece0 100644
+--- a/include/cuco/detail/static_set/static_set.inl
++++ b/include/cuco/detail/static_set/static_set.inl
+@@ -98,7 +98,7 @@ template <class Key,
+           class Allocator,
+           class Storage>
+ void static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::clear(
+-  cuda_stream_ref stream) noexcept
++  cuda_stream_ref stream)
+ {
+   impl_->clear(stream);
+ }
+@@ -429,7 +429,7 @@ template <class Key,
+           class Storage>
+ static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::size_type
+ static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::size(
+-  cuda_stream_ref stream) const noexcept
++  cuda_stream_ref stream) const
+ {
+   return impl_->size(stream);
+ }
+diff --git a/include/cuco/detail/storage/aow_storage.inl b/include/cuco/detail/storage/aow_storage.inl
+index 3547f4c..94b7f98 100644
+--- a/include/cuco/detail/storage/aow_storage.inl
++++ b/include/cuco/detail/storage/aow_storage.inl
+@@ -32,8 +32,8 @@
+ namespace cuco {
+
+ template <typename T, int32_t WindowSize, typename Extent, typename Allocator>
+-constexpr aow_storage<T, WindowSize, Extent, Allocator>::aow_storage(
+-  Extent size, Allocator const& allocator) noexcept
++constexpr aow_storage<T, WindowSize, Extent, Allocator>::aow_storage(Extent size,
++                                                                     Allocator const& allocator)
+   : detail::aow_storage_base<T, WindowSize, Extent>{size},
+     allocator_{allocator},
+     window_deleter_{capacity(), allocator_},
+@@ -64,7 +64,7 @@ aow_storage<T, WindowSize, Extent, Allocator>::ref() const noexcept
+
+ template <typename T, int32_t WindowSize, typename Extent, typename Allocator>
+ void aow_storage<T, WindowSize, Extent, Allocator>::initialize(value_type key,
+-                                                               cuda_stream_ref stream) noexcept
++                                                               cuda_stream_ref stream)
+ {
+   this->initialize_async(key, stream);
+   stream.synchronize();
+diff --git a/include/cuco/static_map.cuh b/include/cuco/static_map.cuh
+index c86e90c..95da423 100644
+--- a/include/cuco/static_map.cuh
++++ b/include/cuco/static_map.cuh
+@@ -269,7 +269,7 @@ class static_map {
+    *
+    * @param stream CUDA stream this operation is executed in
+    */
+-  void clear(cuda_stream_ref stream = {}) noexcept;
++  void clear(cuda_stream_ref stream = {});
+
+   /**
+    * @brief Asynchronously erases all elements from the container. After this call, `size()` returns
+@@ -387,7 +387,7 @@ class static_map {
+    * @param stream CUDA stream used for insert
+    */
+   template <typename InputIt>
+-  void insert_or_assign(InputIt first, InputIt last, cuda_stream_ref stream = {}) noexcept;
++  void insert_or_assign(InputIt first, InputIt last, cuda_stream_ref stream = {});
+
+   /**
+    * @brief For any key-value pair `{k, v}` in the range `[first, last)`, if a key equivalent to `k`
+@@ -690,7 +690,7 @@ class static_map {
+    * @param stream CUDA stream used to get the number of inserted elements
+    * @return The number of elements in the container
+    */
+-  [[nodiscard]] size_type size(cuda_stream_ref stream = {}) const noexcept;
++  [[nodiscard]] size_type size(cuda_stream_ref stream = {}) const;
+
+   /**
+    * @brief Gets the maximum number of elements the hash map can hold.
+diff --git a/include/cuco/static_multiset.cuh b/include/cuco/static_multiset.cuh
+index 0daf103..fbcbc9c 100644
+--- a/include/cuco/static_multiset.cuh
++++ b/include/cuco/static_multiset.cuh
+@@ -235,7 +235,7 @@ class static_multiset {
+    *
+    * @param stream CUDA stream this operation is executed in
+    */
+-  void clear(cuda_stream_ref stream = {}) noexcept;
++  void clear(cuda_stream_ref stream = {});
+
+   /**
+    * @brief Asynchronously erases all elements from the container. After this call, `size()` returns
+@@ -339,7 +339,7 @@ class static_multiset {
+    * @param stream CUDA stream used to get the number of inserted elements
+    * @return The number of elements in the container
+    */
+-  [[nodiscard]] size_type size(cuda_stream_ref stream = {}) const noexcept;
++  [[nodiscard]] size_type size(cuda_stream_ref stream = {}) const;
+
+   /**
+    * @brief Gets the maximum number of elements the multiset can hold.
+diff --git a/include/cuco/static_set.cuh b/include/cuco/static_set.cuh
+index a069939..3517f84 100644
+--- a/include/cuco/static_set.cuh
++++ b/include/cuco/static_set.cuh
+@@ -240,7 +240,7 @@ class static_set {
+    *
+    * @param stream CUDA stream this operation is executed in
+    */
+-  void clear(cuda_stream_ref stream = {}) noexcept;
++  void clear(cuda_stream_ref stream = {});
+
+   /**
+    * @brief Asynchronously erases all elements from the container. After this call, `size()` returns
+@@ -687,7 +687,7 @@ class static_set {
+    * @param stream CUDA stream used to get the number of inserted elements
+    * @return The number of elements in the container
+    */
+-  [[nodiscard]] size_type size(cuda_stream_ref stream = {}) const noexcept;
++  [[nodiscard]] size_type size(cuda_stream_ref stream = {}) const;
+
+   /**
+    * @brief Gets the maximum number of elements the hash set can hold.

cpp/cmake/thirdparty/patches/cuco_override.json

@@ -0,0 +1,14 @@
+
+{
+  "packages" : {
+    "cuco" : {
+      "patches" : [
+        {
+          "file" : "${current_json_dir}/cuco_noexcept.diff",
+          "issue" : "Remove erroneous noexcept clauses on cuco functions that may throw [https://github.com/rapidsai/cudf/issues/16059]",
+          "fixed_in" : ""
+        }
+      ]
+    }
+  }
+}

python/cudf_polars/cudf_polars/dsl/expr.py

@@ -27,11 +27,12 @@
 import cudf._lib.pylibcudf as plc
 
 from cudf_polars.containers import Column, NamedColumn
-from cudf_polars.utils import sorting
+from cudf_polars.utils import dtypes, sorting
 
 if TYPE_CHECKING:
     from collections.abc import Mapping, Sequence
 
+    import polars.polars as plrs
     import polars.type_aliases as pl_types
 
     from cudf_polars.containers import DataFrame
@@ -369,6 +370,29 @@ def do_evaluate(
         return Column(plc.Column.from_scalar(plc.interop.from_arrow(self.value), 1))
 
 
+class LiteralColumn(Expr):
+    __slots__ = ("value",)
+    _non_child = ("dtype", "value")
+    value: pa.Array[Any, Any]
+    children: tuple[()]
+
+    def __init__(self, dtype: plc.DataType, value: plrs.PySeries) -> None:
+        super().__init__(dtype)
+        data = value.to_arrow()
+        self.value = data.cast(dtypes.downcast_arrow_lists(data.type))
+
+    def do_evaluate(
+        self,
+        df: DataFrame,
+        *,
+        context: ExecutionContext = ExecutionContext.FRAME,
+        mapping: Mapping[Expr, Column] | None = None,
+    ) -> Column:
+        """Evaluate this expression given a dataframe for context."""
+        # datatype of pyarrow array is correct by construction.
+        return Column(plc.interop.from_arrow(self.value))
+
+
 class Col(Expr):
     __slots__ = ("name",)
     _non_child = ("dtype", "name")
@@ -1156,6 +1180,12 @@ def __init__(
         super().__init__(dtype)
         self.op = op
         self.children = (left, right)
+        if (
+            op in (plc.binaryop.BinaryOperator.ADD, plc.binaryop.BinaryOperator.SUB)
+            and ({left.dtype.id(), right.dtype.id()}.issubset(dtypes.TIMELIKE_TYPES))
+            and not dtypes.have_compatible_resolution(left.dtype.id(), right.dtype.id())
+        ):
+            raise NotImplementedError("Casting rules for timelike types")
 
     _MAPPING: ClassVar[dict[pl_expr.Operator, plc.binaryop.BinaryOperator]] = {
         pl_expr.Operator.Eq: plc.binaryop.BinaryOperator.EQUAL,

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -29,7 +29,7 @@
 
 import cudf_polars.dsl.expr as expr
 from cudf_polars.containers import DataFrame, NamedColumn
-from cudf_polars.utils import sorting
+from cudf_polars.utils import dtypes, sorting
 
 if TYPE_CHECKING:
     from collections.abc import MutableMapping
@@ -130,6 +130,11 @@ class IR:
     schema: Schema
     """Mapping from column names to their data types."""
 
+    def __post_init__(self):
+        """Validate preconditions."""
+        if any(dtype.id() == plc.TypeId.EMPTY for dtype in self.schema.values()):
+            raise NotImplementedError("Cannot make empty columns.")
+
     def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
         """
         Evaluate the node and return a dataframe.
@@ -292,15 +297,10 @@ def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
         table = pdf.to_arrow()
         schema = table.schema
         for i, field in enumerate(schema):
-            # TODO: Nested types
-            if field.type == pa.large_string():
-                # TODO: goes away when libcudf supports large strings
-                schema = schema.set(i, pa.field(field.name, pa.string()))
-            elif isinstance(field.type, pa.LargeListType):
-                # TODO: goes away when libcudf supports large lists
-                schema = schema.set(
-                    i, pa.field(field.name, pa.list_(field.type.field(0)))
-                )
+            schema = schema.set(
+                i, pa.field(field.name, dtypes.downcast_arrow_lists(field.type))
+            )
+        # No-op if the schema is unchanged.
         table = table.cast(schema)
         df = DataFrame.from_table(
             plc.interop.from_arrow(table), list(self.schema.keys())

python/cudf_polars/cudf_polars/dsl/translate.py

@@ -12,6 +12,7 @@
 import pyarrow as pa
 from typing_extensions import assert_never
 
+import polars.polars as plrs
 from polars.polars import _expr_nodes as pl_expr, _ir_nodes as pl_ir
 
 import cudf._lib.pylibcudf as plc
@@ -383,6 +384,8 @@ def _(node: pl_expr.Window, visitor: NodeTraverser, dtype: plc.DataType) -> expr
 
 @_translate_expr.register
 def _(node: pl_expr.Literal, visitor: NodeTraverser, dtype: plc.DataType) -> expr.Expr:
+    if isinstance(node.value, plrs.PySeries):
+        return expr.LiteralColumn(dtype, node.value)
     value = pa.scalar(node.value, type=plc.interop.to_arrow(dtype))
     return expr.Literal(dtype, value)