Use cuda::stream_ref in Java native code

java/src/main/native/include/maps_column_view.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -10,7 +10,7 @@
 #include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream>
 
 namespace cudf {
 
@@ -29,7 +29,7 @@ namespace jni {
 class maps_column_view {
  public:
   maps_column_view(lists_column_view const& lists_of_structs,
-                   rmm::cuda_stream_view stream = cudf::get_default_stream());
+                   cuda::stream_ref stream = cudf::get_default_stream());
 
   // Rule of 5.
   maps_column_view(maps_column_view const& maps_view)  = default;
@@ -74,7 +74,7 @@ class maps_column_view {
    */
   std::unique_ptr<column> get_values_for(
     column_view const& keys,
-    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+    cuda::stream_ref stream           = cudf::get_default_stream(),
     rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const;
 
   /**
@@ -93,7 +93,7 @@ class maps_column_view {
    */
   std::unique_ptr<column> get_values_for(
     scalar const& key,
-    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+    cuda::stream_ref stream           = cudf::get_default_stream(),
     rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const;
 
   /**
@@ -114,7 +114,7 @@ class maps_column_view {
    */
   std::unique_ptr<column> contains(
     scalar const& key,
-    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+    cuda::stream_ref stream           = cudf::get_default_stream(),
     rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const;
 
   /**
@@ -136,7 +136,7 @@ class maps_column_view {
 
   std::unique_ptr<column> contains(
     column_view const& key,
-    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+    cuda::stream_ref stream           = cudf::get_default_stream(),
     rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const;
 
  private:

java/src/main/native/src/ColumnViewJni.cu

@@ -33,7 +33,7 @@
 namespace cudf::jni {
 
 std::unique_ptr<cudf::column> generate_list_offsets(cudf::column_view const& list_length,
-                                                    rmm::cuda_stream_view stream)
+                                                    cuda::stream_ref stream)
 {
   CUDF_EXPECTS(list_length.type().id() == cudf::type_id::INT32,
                "Input column does not have type INT32.");
@@ -73,7 +73,7 @@ __device__ bool list_has_nulls(list_device_view list)
 void post_process_list_overlap(cudf::column_view const& lhs,
                                cudf::column_view const& rhs,
                                std::unique_ptr<cudf::column> const& overlap_result,
-                               rmm::cuda_stream_view stream)
+                               cuda::stream_ref stream)
 {
   // If both of the input columns do not have nulls, we don't need to do anything here.
   if (!lists_column_view{lhs}.child().has_nulls() && !lists_column_view{rhs}.child().has_nulls()) {
@@ -143,7 +143,7 @@ void post_process_list_overlap(cudf::column_view const& lhs,
 }
 
 std::unique_ptr<cudf::column> lists_distinct_by_key(cudf::lists_column_view const& input,
-                                                    rmm::cuda_stream_view stream)
+                                                    cuda::stream_ref stream)
 {
   if (input.is_empty()) { return empty_like(input.parent()); }
 

java/src/main/native/src/ColumnViewJni.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -9,7 +9,7 @@
 #include <cudf/lists/lists_column_view.hpp>
 #include <cudf/utilities/default_stream.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream>
 
 namespace cudf::jni {
 
@@ -25,7 +25,7 @@ namespace cudf::jni {
  * @return The column represents list offsets.
  */
 std::unique_ptr<cudf::column> generate_list_offsets(
-  cudf::column_view const& list_length, rmm::cuda_stream_view stream = cudf::get_default_stream());
+  cudf::column_view const& list_length, cuda::stream_ref stream = cudf::get_default_stream());
 
 /**
  * @brief Perform a special treatment for the results of `cudf::lists::have_overlap` to produce the
@@ -48,7 +48,7 @@ std::unique_ptr<cudf::column> generate_list_offsets(
 void post_process_list_overlap(cudf::column_view const& lhs,
                                cudf::column_view const& rhs,
                                std::unique_ptr<cudf::column> const& overlap_result,
-                               rmm::cuda_stream_view stream = cudf::get_default_stream());
+                               cuda::stream_ref stream = cudf::get_default_stream());
 
 /**
  * @brief Generates lists column by copying elements that are distinct by key from each input list
@@ -64,6 +64,6 @@ void post_process_list_overlap(cudf::column_view const& lhs,
  * @return A new list columns in which the elements in each list are distinct by key.
  */
 std::unique_ptr<cudf::column> lists_distinct_by_key(cudf::lists_column_view const& input,
-                                                    rmm::cuda_stream_view stream);
+                                                    cuda::stream_ref stream);
 
 }  // namespace cudf::jni

java/src/main/native/src/CompiledExpression.cpp

@@ -538,7 +538,7 @@ std::unique_ptr<cudf::jni::ast::compiled_expr> compile_serialized_ast(jni_serial
   if (!jni_ast.at_eof()) { throw std::invalid_argument("Extra bytes at end of serialized AST"); }
 
   // The expression may be handed to a thread with a different default stream.
-  if (jni_expr_ptr->has_literals()) { cudf::get_default_stream().synchronize(); }
+  if (jni_expr_ptr->has_literals()) { cudf::get_default_stream().sync(); }
 
   return jni_expr_ptr;
 }

java/src/main/native/src/RmmJni.cpp

@@ -12,7 +12,6 @@
 #include <cudf/utilities/pinned_memory.hpp>
 
 #include <rmm/aligned.hpp>
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/detail/error.hpp>
 #include <rmm/error.hpp>
 #include <rmm/mr/aligned_resource_adaptor.hpp>
@@ -28,6 +27,7 @@
 #include <rmm/resource_ref.hpp>
 
 #include <cuda/memory_resource>
+#include <cuda/stream>
 #include <cuda_runtime_api.h>
 
 #include <sys/mman.h>
@@ -937,8 +937,8 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_allocInternal(JNIEnv* env,
   {
     cudf::jni::auto_set_device(env);
     rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref();
-    auto c_stream = rmm::cuda_stream_view(reinterpret_cast<cudaStream_t>(stream));
-    void* ret     = mr.allocate(c_stream, size);
+    auto c_stream                     = cuda::stream_ref(reinterpret_cast<cudaStream_t>(stream));
+    void* ret                         = mr.allocate(c_stream, size);
     return reinterpret_cast<jlong>(ret);
   }
   JNI_CATCH(env, 0);
@@ -952,7 +952,7 @@ Java_ai_rapids_cudf_Rmm_free(JNIEnv* env, jclass clazz, jlong ptr, jlong size, j
     cudf::jni::auto_set_device(env);
     rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref();
     void* cptr                        = reinterpret_cast<void*>(ptr);
-    auto c_stream = rmm::cuda_stream_view(reinterpret_cast<cudaStream_t>(stream));
+    auto c_stream                     = cuda::stream_ref(reinterpret_cast<cudaStream_t>(stream));
     mr.deallocate(c_stream, cptr, size);
   }
   JNI_CATCH(env, );

java/src/main/native/src/aggregation128_utils.cu

@@ -71,7 +71,7 @@ namespace cudf::jni {
 std::unique_ptr<cudf::column> extract_chunk32(cudf::column_view const& in_col,
                                               cudf::data_type type,
                                               int chunk_idx,
-                                              rmm::cuda_stream_view stream)
+                                              cuda::stream_ref stream)
 {
   CUDF_EXPECTS(in_col.type().id() == cudf::type_id::DECIMAL128, "not a 128-bit type");
   CUDF_EXPECTS(chunk_idx >= 0 && chunk_idx < 4, "invalid chunk index");
@@ -97,7 +97,7 @@ std::unique_ptr<cudf::column> extract_chunk32(cudf::column_view const& in_col,
 // Reassemble a column of 128-bit values from four 64-bit integer columns with overflow detection.
 std::unique_ptr<cudf::table> assemble128_from_sum(cudf::table_view const& chunks_table,
                                                   cudf::data_type output_type,
-                                                  rmm::cuda_stream_view stream)
+                                                  cuda::stream_ref stream)
 {
   CUDF_EXPECTS(output_type.id() == cudf::type_id::DECIMAL128, "not a 128-bit type");
   CUDF_EXPECTS(chunks_table.num_columns() == 4, "must be 4 column table");

java/src/main/native/src/aggregation128_utils.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -9,7 +9,7 @@
 #include <cudf/table/table.hpp>
 #include <cudf/utilities/default_stream.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream>
 
 #include <memory>
 
@@ -31,11 +31,10 @@ namespace cudf::jni {
  * @param stream    CUDA stream to use
  * @return          A column containing the extracted 32-bit integer values
  */
-std::unique_ptr<cudf::column> extract_chunk32(
-  cudf::column_view const& col,
-  cudf::data_type dtype,
-  int chunk_idx,
-  rmm::cuda_stream_view stream = cudf::get_default_stream());
+std::unique_ptr<cudf::column> extract_chunk32(cudf::column_view const& col,
+                                              cudf::data_type dtype,
+                                              int chunk_idx,
+                                              cuda::stream_ref stream = cudf::get_default_stream());
 
 /**
  * @brief Reassemble a 128-bit column from four 64-bit integer columns with overflow detection.
@@ -60,6 +59,6 @@ std::unique_ptr<cudf::column> extract_chunk32(
 std::unique_ptr<cudf::table> assemble128_from_sum(
   cudf::table_view const& chunks_table,
   cudf::data_type output_type,
-  rmm::cuda_stream_view stream = cudf::get_default_stream());
+  cuda::stream_ref stream = cudf::get_default_stream());
 
 }  // namespace cudf::jni

java/src/main/native/src/check_nvcomp_output_sizes.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #include "check_nvcomp_output_sizes.hpp"
@@ -28,7 +28,7 @@ namespace java {
 bool check_nvcomp_output_sizes(std::size_t const* dev_uncompressed_sizes,
                                std::size_t const* dev_actual_uncompressed_sizes,
                                std::size_t num_chunks,
-                               rmm::cuda_stream_view stream)
+                               cuda::stream_ref stream)
 {
   NVTX3_FUNC_RANGE_IN(java_domain);
   return thrust::equal(rmm::exec_policy_nosync(stream),

java/src/main/native/src/check_nvcomp_output_sizes.hpp

@@ -1,13 +1,14 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/exec_policy.hpp>
 
+#include <cuda/stream>
+
 namespace cudf {
 namespace java {
 
@@ -18,6 +19,6 @@ namespace java {
 bool check_nvcomp_output_sizes(std::size_t const* dev_uncompressed_sizes,
                                std::size_t const* dev_actual_uncompressed_sizes,
                                std::size_t num_chunks,
-                               rmm::cuda_stream_view stream);
+                               cuda::stream_ref stream);
 }  // namespace java
 }  // namespace cudf

java/src/main/native/src/maps_column_view.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -29,7 +29,7 @@ column_view make_lists(column_view const& lists_child, lists_column_view const&
 }  // namespace
 
 maps_column_view::maps_column_view(lists_column_view const& lists_of_structs,
-                                   rmm::cuda_stream_view stream)
+                                   cuda::stream_ref stream)
   : keys_{make_lists(lists_of_structs.child().child(0), lists_of_structs)},
     values_{make_lists(lists_of_structs.child().child(1), lists_of_structs)}
 {
@@ -43,7 +43,7 @@ maps_column_view::maps_column_view(lists_column_view const& lists_of_structs,
 template <typename KeyT>
 std::unique_ptr<column> get_values_for_impl(maps_column_view const& maps_view,
                                             KeyT const& lookup_keys,
-                                            rmm::cuda_stream_view stream,
+                                            cuda::stream_ref stream,
                                             rmm::device_async_resource_ref mr)
 {
   auto const keys_   = maps_view.keys();
@@ -66,7 +66,7 @@ std::unique_ptr<column> get_values_for_impl(maps_column_view const& maps_view,
 }
 
 std::unique_ptr<column> maps_column_view::get_values_for(column_view const& lookup_keys,
-                                                         rmm::cuda_stream_view stream,
+                                                         cuda::stream_ref stream,
                                                          rmm::device_async_resource_ref mr) const
 {
   CUDF_EXPECTS(lookup_keys.size() == size(),
@@ -76,7 +76,7 @@ std::unique_ptr<column> maps_column_view::get_values_for(column_view const& look
 }
 
 std::unique_ptr<column> maps_column_view::get_values_for(cudf::scalar const& lookup_key,
-                                                         rmm::cuda_stream_view stream,
+                                                         cuda::stream_ref stream,
                                                          rmm::device_async_resource_ref mr) const
 {
   return get_values_for_impl(*this, lookup_key, stream, mr);
@@ -85,7 +85,7 @@ std::unique_ptr<column> maps_column_view::get_values_for(cudf::scalar const& loo
 template <typename KeyT>
 std::unique_ptr<column> contains_impl(maps_column_view const& maps_view,
                                       KeyT const& lookup_keys,
-                                      rmm::cuda_stream_view stream,
+                                      cuda::stream_ref stream,
                                       rmm::device_async_resource_ref mr)
 {
   auto const keys = maps_view.keys();
@@ -99,7 +99,7 @@ std::unique_ptr<column> contains_impl(maps_column_view const& maps_view,
 }
 
 std::unique_ptr<column> maps_column_view::contains(column_view const& lookup_keys,
-                                                   rmm::cuda_stream_view stream,
+                                                   cuda::stream_ref stream,
                                                    rmm::device_async_resource_ref mr) const
 {
   CUDF_EXPECTS(lookup_keys.size() == size(),
@@ -109,7 +109,7 @@ std::unique_ptr<column> maps_column_view::contains(column_view const& lookup_key
 }
 
 std::unique_ptr<column> maps_column_view::contains(cudf::scalar const& lookup_key,
-                                                   rmm::cuda_stream_view stream,
+                                                   cuda::stream_ref stream,
                                                    rmm::device_async_resource_ref mr) const
 {
   return contains_impl(*this, lookup_key, stream, mr);