Merge remote-tracking branch 'origin/main' into nvtext-mr-tokenize

# Conflicts:
#	cpp/src/text/bpe/byte_pair_encoding.cu

Author: davidwendt

cpp/include/cudf/dictionary/dictionary_factories.hpp

@@ -83,14 +83,14 @@ std::unique_ptr<column> make_dictionary_column(std::unique_ptr<column> keys_colu
  * @brief Construct a dictionary column by taking ownership of the provided keys
  * and indices columns.
  *
- * The `keys_column` must contain no nulls and is assumed to have elements
- * that are unique and are in a strict, total order. Meaning, `keys_column[i]`
- * is ordered before `keys_column[i+1]` for all `i in [0,n-1)` where `n` is the
- * number of keys.
+ * The `keys_column` must contain no nulls. The keys are not required to be sorted
+ * or unique.
  *
  * The `indices_column` can be any integer type and should contain the null-mask
  * to be used for the output column.
  * The indices values must be in the range [0,keys_column.size()).
+ * Unsigned indices are converted to the signed type of the same width, or to the
+ * next wider signed type if the same width cannot hold `keys_column.size() - 1`.
  *
  * @throw cudf::logic_error if keys_column contains nulls
  *

cpp/include/cudf/interop.hpp

@@ -146,9 +146,9 @@ using owned_columns_t = std::vector<std::unique_ptr<cudf::column>>;
  * @brief functor for a custom deleter to a unique_ptr of table_view
  *
  * When converting from an ArrowDeviceArray, there are cases where data can't
- * be zero-copy (i.e. bools or non-UINT32 dictionary indices). This custom deleter
- * is used to maintain ownership over the data allocated since a `cudf::table_view`
- * doesn't hold ownership.
+ * be zero-copy (i.e. bools or unsigned dictionary indices that must be widened to fit
+ * the number of keys). This custom deleter is used to maintain ownership over the data
+ * allocated since a `cudf::table_view` doesn't hold ownership.
  */
 template <typename ViewType>
 struct custom_view_deleter {

cpp/include/cudf_test/nanoarrow_utils.hpp

@@ -290,6 +290,7 @@ DEFINE_NANOARROW_STORAGE(int32_t, INT32);
 DEFINE_NANOARROW_STORAGE(int64_t, INT64);
 DEFINE_NANOARROW_STORAGE(uint8_t, UINT8);
 DEFINE_NANOARROW_STORAGE(uint16_t, UINT16);
+DEFINE_NANOARROW_STORAGE(uint32_t, UINT32);
 DEFINE_NANOARROW_STORAGE(uint64_t, UINT64);
 DEFINE_NANOARROW_STORAGE(bool, BOOL);
 DEFINE_NANOARROW_STORAGE(cudf::duration_D, INT32);

cpp/src/dictionary/dictionary_factories.cu

@@ -13,6 +13,7 @@
 #include <cudf/utilities/memory_resource.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
 
+#include <cuda/std/limits>
 #include <cuda/std/type_traits>
 #include <cuda/stream>
 
@@ -93,16 +94,25 @@ std::unique_ptr<column> make_dictionary_column(std::unique_ptr<column> keys_colu
 namespace {
 
 /**
- * @brief This functor maps unsigned type_ids to signed counterparts
+ * @brief This functor maps unsigned type_ids to the narrowest signed type_id
+ * that can address all the keys
  */
 struct make_signed_fn {
+  size_type num_keys;
+
   template <cudf::concepts::index_type T>
-  constexpr cudf::type_id operator()()
+  constexpr cudf::type_id operator()() const
   {
-    return cudf::type_to_id<cuda::std::make_signed_t<T>>();
+    using SignedType = cuda::std::make_signed_t<T>;
+    if constexpr (cuda::std::is_unsigned_v<T> && sizeof(T) < sizeof(size_type)) {
+      if (num_keys - 1 > cuda::std::numeric_limits<SignedType>::max()) {
+        return sizeof(T) == 1 ? type_id::INT16 : type_id::INT32;
+      }
+    }
+    return cudf::type_to_id<SignedType>();
   }
   template <typename T>
-  constexpr cudf::type_id operator()()
+  constexpr cudf::type_id operator()() const
   {
     return cudf::type_to_id<T>();
   }
@@ -122,21 +132,23 @@ std::unique_ptr<column> make_dictionary_column(std::unique_ptr<column> keys,
   auto const null_count   = indices->null_count();  // before calling release()
   auto contents           = indices->release();
 
-  auto const new_type_id = cudf::type_dispatcher(indices_type, make_signed_fn{});
+  auto const new_type =
+    data_type{cudf::type_dispatcher(indices_type, make_signed_fn{keys->size()})};
 
   // create the dictionary indices: convert to signed and remove nulls if necessary
   auto indices_column = [&] {
-    // If the types match, then just commandeer the column's data buffer.
-    if (new_type_id == indices_type.id()) {
-      return std::make_unique<column>(indices_type,
+    // If the widths match, then just commandeer the column's data buffer.
+    // Every valid index is less than keys->size() so its bits are the same in both types.
+    if (new_type == indices_type || cudf::size_of(new_type) == cudf::size_of(indices_type)) {
+      return std::make_unique<column>(new_type,
                                       indices_size,
                                       std::move(*(contents.data.release())),
                                       cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED),
                                       0);
     }
-    // If the new type does not match, then convert the data.
+    // If the widths do not match, then convert the data.
     cudf::column_view cast_view{indices_type, indices_size, contents.data->data(), nullptr, 0};
-    return cudf::detail::cast(cast_view, data_type{new_type_id}, stream, mr);
+    return cudf::detail::cast(cast_view, new_type, stream, mr);
   }();
 
   return make_dictionary_column(std::move(keys),

cpp/src/interop/arrow_utilities.cpp

@@ -163,6 +163,30 @@ ArrowType id_to_arrow_storage_type(cudf::type_id id)
   }
 }
 
+dictionary_indices_types get_dictionary_indices_types(ArrowType storage_type, int64_t num_keys)
+{
+  auto const same = [](type_id id) {
+    return dictionary_indices_types{data_type{id}, data_type{id}};
+  };
+  auto const max_index = num_keys - 1;
+  switch (storage_type) {
+    case NANOARROW_TYPE_INT8: return same(type_id::INT8);
+    case NANOARROW_TYPE_INT16: return same(type_id::INT16);
+    case NANOARROW_TYPE_INT32: return same(type_id::INT32);
+    case NANOARROW_TYPE_INT64: return same(type_id::INT64);
+    case NANOARROW_TYPE_UINT8:
+      return max_index <= std::numeric_limits<int8_t>::max()
+               ? same(type_id::INT8)
+               : dictionary_indices_types{data_type{type_id::UINT8}, data_type{type_id::INT16}};
+    case NANOARROW_TYPE_UINT16:
+      return max_index <= std::numeric_limits<int16_t>::max()
+               ? same(type_id::INT16)
+               : dictionary_indices_types{data_type{type_id::UINT16}, data_type{type_id::INT32}};
+    case NANOARROW_TYPE_UINT32: return same(type_id::INT32);
+    default: CUDF_FAIL("Unsupported type_id for dictionary indices", cudf::data_type_error);
+  }
+}
+
 int initialize_array(ArrowArray* arr, ArrowType storage_type, cudf::column_view column)
 {
   NANOARROW_RETURN_NOT_OK(ArrowArrayInitFromType(arr, storage_type));

cpp/src/interop/arrow_utilities.hpp

@@ -96,6 +96,30 @@ ArrowType id_to_arrow_type(cudf::type_id id);
  */
 ArrowType id_to_arrow_storage_type(cudf::type_id id);
 
+/**
+ * @brief cudf types used to import the indices of an Arrow dictionary array
+ */
+struct dictionary_indices_types {
+  data_type storage;  ///< Type used to read the Arrow indices buffer
+  data_type output;   ///< Signed type of the resulting cudf dictionary indices
+};
+
+/**
+ * @brief Return the cudf types used to import Arrow dictionary indices
+ *
+ * Signed index types are returned unchanged. Unsigned index types are mapped to the
+ * narrowest signed type that can represent every valid index `[0, num_keys - 1]`:
+ * the same-width signed type when it fits, otherwise the next wider signed type.
+ *
+ * @throw cudf::data_type_error if `storage_type` is not an 8, 16, or 32-bit unsigned or
+ * an 8, 16, 32, or 64-bit signed integer type
+ *
+ * @param storage_type Arrow storage type of the dictionary indices
+ * @param num_keys Number of keys in the dictionary, at most `size_type` max
+ * @return Types to read the Arrow buffer with and to output
+ */
+dictionary_indices_types get_dictionary_indices_types(ArrowType storage_type, int64_t num_keys);
+
 /**
  * @brief Helper to initialize ArrowArray struct
  *

cpp/src/interop/from_arrow_device.cu

@@ -233,26 +233,24 @@ dispatch_tuple_t dispatch_from_arrow_device::operator()<cudf::dictionary32>(
   auto [keys_view, owned_cols] =
     get_column(&keys_schema_view, input->dictionary, keys_type, true, stream, mr);
 
-  auto const dict_indices_type = [&schema]() -> data_type {
-    // cudf dictionary requires a signed type for the indices
-    switch (schema->storage_type) {
-      case NANOARROW_TYPE_INT8: return data_type(type_id::INT8);
-      case NANOARROW_TYPE_INT16: return data_type(type_id::INT16);
-      case NANOARROW_TYPE_INT32: return data_type(type_id::INT32);
-      case NANOARROW_TYPE_INT64: return data_type(type_id::INT64);
-      default: CUDF_FAIL("Unsupported type_id for dictionary indices", cudf::data_type_error);
-    }
-  }();
+  auto const indices_types =
+    get_dictionary_indices_types(schema->storage_type, input->dictionary->length);
 
   size_type const num_rows   = input->length;
   size_type const offset     = input->offset;
   size_type const null_count = input->null_count;
-  column_view indices_view   = column_view{dict_indices_type,
+  column_view indices_view   = column_view{indices_types.storage,
                                          offset + num_rows,
                                          input->buffers[fixed_width_data_buffer_idx],
                                          nullptr,
                                          0,
                                          0};
+  // Widen the full [0, offset + num_rows) range since the parent offset applies to the indices
+  if (indices_types.storage != indices_types.output) {
+    auto widened = cudf::detail::cast(indices_view, indices_types.output, stream, mr);
+    indices_view = widened->view();
+    owned_cols.emplace_back(std::move(widened));
+  }
 
   return std::make_tuple<column_view, owned_columns_t>(
     {type,

cpp/src/interop/from_arrow_host.cu

@@ -257,18 +257,9 @@ std::unique_ptr<column> dispatch_copy_from_arrow_host::operator()<cudf::dictiona
   auto keys_column =
     get_column_copy(&keys_schema_view, input->dictionary, keys_type, true, stream, mr);
 
-  auto const dict_indices_type = [&schema]() -> data_type {
-    // cudf dictionary requires a signed type for the indices
-    switch (schema->storage_type) {
-      case NANOARROW_TYPE_INT8: return data_type(type_id::INT8);
-      case NANOARROW_TYPE_INT16: return data_type(type_id::INT16);
-      case NANOARROW_TYPE_INT32: return data_type(type_id::INT32);
-      case NANOARROW_TYPE_INT64: return data_type(type_id::INT64);
-      default: CUDF_FAIL("Unsupported type_id for dictionary indices", cudf::data_type_error);
-    }
-  }();
-
-  auto indices_column = get_column_copy(schema, input, dict_indices_type, skip_mask, stream, mr);
+  auto const indices_type =
+    get_dictionary_indices_types(schema->storage_type, input->dictionary->length).storage;
+  auto indices_column = get_column_copy(schema, input, indices_type, skip_mask, stream, mr);
 
   return make_dictionary_column(std::move(keys_column), std::move(indices_column), stream, mr);
 }

cpp/src/io/avro/reader_impl.cu

@@ -12,6 +12,7 @@
 
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/structs/utilities.hpp>
+#include <cudf/detail/utilities/buffer_factories.hpp>
 #include <cudf/detail/utilities/cuda_memcpy.hpp>
 #include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/io/datasource.hpp>
@@ -22,7 +23,6 @@
 #include <cudf/utilities/span.hpp>
 #include <cudf/utilities/traits.hpp>
 
-#include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
 
 #include <cuda/buffer>
@@ -297,7 +297,8 @@ cuda::device_buffer<std::uint8_t> decompress_data(
 
     cuda::device_buffer<std::uint8_t> decompressed_data(
       stream, cudf::get_current_device_resource_ref(), uncompressed_data_size, cuda::no_init);
-    rmm::device_uvector<device_span<uint8_t>> decompressed_blocks(num_blocks, stream);
+    cuda::device_buffer<device_span<uint8_t>> decompressed_blocks(
+      stream, cudf::get_current_device_resource_ref(), num_blocks, cuda::no_init);
     thrust::tabulate(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
                      decompressed_blocks.begin(),
                      decompressed_blocks.end(),
@@ -307,7 +308,8 @@ cuda::device_buffer<std::uint8_t> decompress_data(
                        return device_span<uint8_t>{data + off[i], size[i]};
                      });
 
-    rmm::device_uvector<codec_exec_result> decomp_results(num_blocks, stream);
+    cuda::device_buffer<codec_exec_result> decomp_results(
+      stream, cudf::get_current_device_resource_ref(), num_blocks, cuda::no_init);
     thrust::fill(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
                  decomp_results.begin(),
                  decomp_results.end(),
@@ -425,7 +427,7 @@ std::vector<column_buffer> decode_data(metadata& meta,
     }
   }
 
-  auto block_list = cudf::detail::make_device_uvector_async(
+  auto block_list = cudf::detail::make_device_buffer_async(
     meta.block_list, stream, cudf::get_current_device_resource_ref());
 
   schema_desc.host_to_device_async(stream);
@@ -539,8 +541,10 @@ table_with_metadata read_avro(std::unique_ptr<cudf::io::datasource>&& source,
         }
       }
 
-      auto d_global_dict      = rmm::device_uvector<string_index_pair>(0, stream);
-      auto d_global_dict_data = rmm::device_uvector<char>(0, stream);
+      cuda::device_buffer<string_index_pair> d_global_dict(
+        stream, cudf::get_current_device_resource_ref(), 0, cuda::no_init);
+      cuda::device_buffer<char> d_global_dict_data(
+        stream, cudf::get_current_device_resource_ref(), 0, cuda::no_init);
 
       if (total_dictionary_entries > 0) {
         auto h_global_dict =
@@ -566,9 +570,9 @@ table_with_metadata read_avro(std::unique_ptr<cudf::io::datasource>&& source,
           }
         }
 
-        d_global_dict = cudf::detail::make_device_uvector_async(
+        d_global_dict = cudf::detail::make_device_buffer_async(
           h_global_dict, stream, cudf::get_current_device_resource_ref());
-        d_global_dict_data = cudf::detail::make_device_uvector_async(
+        d_global_dict_data = cudf::detail::make_device_buffer_async(
           h_global_dict_data, stream, cudf::get_current_device_resource_ref());
 
         stream.sync();

cpp/src/io/fst/logical_stack.cuh

@@ -6,6 +6,8 @@
 
 #include <cudf_test/print_utilities.cuh>
 
+#include <cudf/detail/utilities/cuda.hpp>
+#include <cudf/detail/utilities/cuda_memcpy.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/error.hpp>
@@ -455,9 +457,12 @@ void sparse_stack_op_to_top_of_stack(StackSymbolItT d_symbols,
   // temp_storage_bytes
   total_temp_storage_bytes = temp_storage.size();
 
-  rmm::device_uvector<SymbolPositionT> d_symbol_position_alt{num_symbols_in, stream};
-  rmm::device_uvector<StackOpT> d_kv_ops_current{num_symbols_in, stream};
-  rmm::device_uvector<StackOpT> d_kv_ops_alt{num_symbols_in, stream};
+  cuda::device_buffer<SymbolPositionT> d_symbol_position_alt(
+    stream, cudf::get_current_device_resource_ref(), num_symbols_in, cuda::no_init);
+  cuda::device_buffer<StackOpT> d_kv_ops_current{
+    stream, cudf::get_current_device_resource_ref(), num_symbols_in, cuda::no_init};
+  cuda::device_buffer<StackOpT> d_kv_ops_alt(
+    stream, cudf::get_current_device_resource_ref(), num_symbols_in, cuda::no_init);
 
   //------------------------------------------------------------------------------
   // ALGORITHM
@@ -476,7 +481,8 @@ void sparse_stack_op_to_top_of_stack(StackSymbolItT d_symbols,
       d_symbols,
       detail::NewlineToResetStackSegmentOp<StackSymbolToStackOpTypeT>{symbol_to_stack_op});
 
-    rmm::device_uvector<StackSegmentT> key_segments{num_symbols_in, stream};
+    cuda::device_buffer<StackSegmentT> key_segments(
+      stream, cudf::get_current_device_resource_ref(), num_symbols_in, cuda::no_init);
     CUDF_CUDA_TRY(cub::DeviceScan::InclusiveSum(
       temp_storage.data(),
       total_temp_storage_bytes,
@@ -508,7 +514,10 @@ void sparse_stack_op_to_top_of_stack(StackSymbolItT d_symbols,
 
   // Check if the last element of d_kv_operations is 0. If not, then we have a problem.
   if (num_symbols_in && !supports_reset_op) {
-    StackOpT last_symbol = d_kv_ops_current.element(num_symbols_in - 1, stream);
+    StackOpT last_symbol;
+    CUDF_CUDA_TRY(cudf::detail::memcpy_async(
+      &last_symbol, d_kv_ops_current.data() + num_symbols_in - 1, sizeof(last_symbol), stream));
+    cudf::detail::sync_stream(stream);
     CUDF_EXPECTS(last_symbol.stack_level == 0, "The logical stack is not empty!");
   }