Add stream and mr parameters for `structs::detail::flatten_nested…

…_columns` (#12892)

The internal API  `structs::detail::flatten_nested_columns` currently always use default stream and device resources. This PR adds `stream` and `mr` parameters for it, allowing to pass in the current `stream` and `mr` values.

Closes #12349.

Authors:
  - Nghia Truong (https://github.com/ttnghia)

Approvers:
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Mike Wilson (https://github.com/hyperbolic2346)
  - Divye Gala (https://github.com/divyegala)
  - Karthikeyan (https://github.com/karthikeyann)

URL: #12892

cpp/include/cudf/detail/structs/utilities.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -164,14 +164,18 @@ class flattened_table {
  * @param null_precedence null order for input table
  * @param nullability force output to have nullability columns even if input columns
  * are all valid
+ * @param stream CUDA stream used for device memory operations and kernel launches
+ * @param mr Device memory resource used to allocate new device memory
  * @return `flatten_result` with flattened table, flattened column order, flattened null precedence,
  * alongside the supporting columns and device_buffers for the flattened table.
  */
 [[nodiscard]] flattened_table flatten_nested_columns(
   table_view const& input,
   std::vector<order> const& column_order,
   std::vector<null_order> const& null_precedence,
-  column_nullability nullability = column_nullability::MATCH_INCOMING);
+  column_nullability nullability,
+  rmm::cuda_stream_view stream,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
  * @brief Superimpose nulls from a given null mask into the input column, using bitwise AND.

cpp/include/cudf/detail/unary.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2018-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -74,6 +74,16 @@ std::unique_ptr<cudf::column> unary_operation(
   rmm::cuda_stream_view stream,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
+/**
+ * @copydoc cudf::is_valid
+ *
+ * @param stream CUDA stream used for device memory operations and kernel launches.
+ */
+std::unique_ptr<cudf::column> is_valid(
+  cudf::column_view const& input,
+  rmm::cuda_stream_view stream,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
 /**
  * @copydoc cudf::cast
  *

cpp/src/join/hash_join.cu

@@ -299,7 +299,7 @@ hash_join<Hasher>::hash_join(cudf::table_view const& build,
 
   // need to store off the owning structures for some of the views in _build
   _flattened_build_table = structs::detail::flatten_nested_columns(
-    build, {}, {}, structs::detail::column_nullability::FORCE);
+    build, {}, {}, structs::detail::column_nullability::FORCE, stream);
   _build = _flattened_build_table;
 
   if (_is_empty) { return; }
@@ -357,7 +357,7 @@ std::size_t hash_join<Hasher>::inner_join_size(cudf::table_view const& probe,
   if (_is_empty) { return 0; }
 
   auto flattened_probe = structs::detail::flatten_nested_columns(
-    probe, {}, {}, structs::detail::column_nullability::FORCE);
+    probe, {}, {}, structs::detail::column_nullability::FORCE, stream);
   auto const flattened_probe_table = flattened_probe.flattened_columns();
 
   auto build_table_ptr           = cudf::table_device_view::create(_build, stream);
@@ -382,7 +382,7 @@ std::size_t hash_join<Hasher>::left_join_size(cudf::table_view const& probe,
   if (_is_empty) { return probe.num_rows(); }
 
   auto flattened_probe = structs::detail::flatten_nested_columns(
-    probe, {}, {}, structs::detail::column_nullability::FORCE);
+    probe, {}, {}, structs::detail::column_nullability::FORCE, stream);
   auto const flattened_probe_table = flattened_probe.flattened_columns();
 
   auto build_table_ptr           = cudf::table_device_view::create(_build, stream);
@@ -408,7 +408,7 @@ std::size_t hash_join<Hasher>::full_join_size(cudf::table_view const& probe,
   if (_is_empty) { return probe.num_rows(); }
 
   auto flattened_probe = structs::detail::flatten_nested_columns(
-    probe, {}, {}, structs::detail::column_nullability::FORCE);
+    probe, {}, {}, structs::detail::column_nullability::FORCE, stream);
   auto const flattened_probe_table = flattened_probe.flattened_columns();
 
   auto build_table_ptr           = cudf::table_device_view::create(_build, stream);
@@ -475,7 +475,7 @@ hash_join<Hasher>::compute_hash_join(cudf::table_view const& probe,
                "Probe column size is too big for hash join");
 
   auto flattened_probe = structs::detail::flatten_nested_columns(
-    probe, {}, {}, structs::detail::column_nullability::FORCE);
+    probe, {}, {}, structs::detail::column_nullability::FORCE, stream);
   auto const flattened_probe_table = flattened_probe.flattened_columns();
 
   CUDF_EXPECTS(_build.num_columns() == flattened_probe_table.num_columns(),

cpp/src/reductions/struct_minmax_util.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -98,7 +98,11 @@ class comparison_binop_generator {
 
   comparison_binop_generator(column_view const& input, rmm::cuda_stream_view stream, bool is_min_op)
     : flattened_input{cudf::structs::detail::flatten_nested_columns(
-        table_view{{input}}, {}, std::vector<null_order>{DEFAULT_NULL_ORDER})},
+        table_view{{input}},
+        {},
+        std::vector<null_order>{DEFAULT_NULL_ORDER},
+        cudf::structs::detail::column_nullability::MATCH_INCOMING,
+        stream)},
       d_flattened_input_ptr{table_device_view::create(flattened_input, stream)},
       is_min_op(is_min_op),
       has_nulls{has_nested_nulls(table_view{{input}})},

cpp/src/search/contains_table.cu

@@ -326,9 +326,9 @@ rmm::device_uvector<bool> contains_without_lists_or_nans(table_view const& hayst
                                      ? structs::detail::column_nullability::FORCE
                                      : structs::detail::column_nullability::MATCH_INCOMING;
   auto const haystack_flattened_tables =
-    structs::detail::flatten_nested_columns(haystack, {}, {}, flatten_nullability);
+    structs::detail::flatten_nested_columns(haystack, {}, {}, flatten_nullability, stream);
   auto const needles_flattened_tables =
-    structs::detail::flatten_nested_columns(needles, {}, {}, flatten_nullability);
+    structs::detail::flatten_nested_columns(needles, {}, {}, flatten_nullability, stream);
   auto const haystack_flattened = haystack_flattened_tables.flattened_columns();
   auto const needles_flattened  = needles_flattened_tables.flattened_columns();
   auto const haystack_tdv_ptr   = table_device_view::create(haystack_flattened, stream);

cpp/src/structs/utilities.cpp

@@ -18,11 +18,11 @@
 #include <cudf/detail/copy.hpp>
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/structs/utilities.hpp>
+#include <cudf/detail/unary.hpp>
 #include <cudf/structs/structs_column_view.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/types.hpp>
-#include <cudf/unary.hpp>
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/span.hpp>
 #include <cudf/utilities/traits.hpp>
@@ -87,22 +87,29 @@ bool is_or_has_nested_lists(cudf::column_view const& col)
  */
 struct table_flattener {
   table_view input;
-  // reference variables
   std::vector<order> const& column_order;
   std::vector<null_order> const& null_precedence;
-  // output
-  std::vector<std::unique_ptr<column>> validity_as_column;
+  column_nullability nullability;
+  rmm::cuda_stream_view stream;
+  rmm::mr::device_memory_resource* mr;
+
   temporary_nullable_data nullable_data;
+  std::vector<std::unique_ptr<column>> validity_as_column;
   std::vector<column_view> flat_columns;
   std::vector<order> flat_column_order;
   std::vector<null_order> flat_null_precedence;
-  column_nullability nullability;
 
   table_flattener(table_view const& input,
                   std::vector<order> const& column_order,
                   std::vector<null_order> const& null_precedence,
-                  column_nullability nullability)
-    : column_order(column_order), null_precedence(null_precedence), nullability(nullability)
+                  column_nullability nullability,
+                  rmm::cuda_stream_view stream,
+                  rmm::mr::device_memory_resource* mr)
+    : column_order{column_order},
+      null_precedence{null_precedence},
+      nullability{nullability},
+      stream{stream},
+      mr{mr}
   {
     superimpose_nulls(input);
     fail_if_unsupported_types(input);
@@ -114,7 +121,7 @@ struct table_flattener {
    */
   void superimpose_nulls(table_view const& input_table)
   {
-    auto [table, tmp_nullable_data] = push_down_nulls(input_table, cudf::get_default_stream());
+    auto [table, tmp_nullable_data] = push_down_nulls(input_table, stream, mr);
     this->input                     = std::move(table);
     this->nullable_data             = std::move(tmp_nullable_data);
   }
@@ -143,10 +150,10 @@ struct table_flattener {
     // sure the flattening results are tables having the same number of columns.
 
     if (nullability == column_nullability::FORCE || col.has_nulls()) {
-      validity_as_column.push_back(cudf::is_valid(col));
+      validity_as_column.push_back(cudf::detail::is_valid(col, stream, mr));
       if (col.has_nulls()) {
         // copy bitmask is needed only if the column has null
-        validity_as_column.back()->set_null_mask(copy_bitmask(col));
+        validity_as_column.back()->set_null_mask(cudf::detail::copy_bitmask(col, stream, mr));
       }
       flat_columns.push_back(validity_as_column.back()->view());
       if (not column_order.empty()) { flat_column_order.push_back(col_order); }  // doesn't matter.
@@ -197,12 +204,14 @@ struct table_flattener {
 flattened_table flatten_nested_columns(table_view const& input,
                                        std::vector<order> const& column_order,
                                        std::vector<null_order> const& null_precedence,
-                                       column_nullability nullability)
+                                       column_nullability nullability,
+                                       rmm::cuda_stream_view stream,
+                                       rmm::mr::device_memory_resource* mr)
 {
   auto const has_struct = std::any_of(input.begin(), input.end(), is_struct);
   if (not has_struct) { return flattened_table{input, column_order, null_precedence, {}, {}}; }
 
-  return table_flattener{input, column_order, null_precedence, nullability}();
+  return table_flattener{input, column_order, null_precedence, nullability, stream, mr}();
 }
 
 namespace {

cpp/tests/structs/utilities_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,6 +27,7 @@
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/structs/utilities.hpp>
 #include <cudf/null_mask.hpp>
+#include <cudf/utilities/default_stream.hpp>
 
 template <typename T>
 using nums = cudf::test::fixed_width_column_wrapper<T, int32_t>;
@@ -54,9 +55,10 @@ TYPED_TEST(TypedStructUtilitiesTest, ListsAtTopLevel)
 
   auto table = cudf::table_view{{lists_col, nums_col}};
 
-  CUDF_TEST_EXPECT_TABLES_EQUAL(table,
-                                cudf::structs::detail::flatten_nested_columns(
-                                  table, {}, {}, cudf::structs::detail::column_nullability::FORCE));
+  CUDF_TEST_EXPECT_TABLES_EQUAL(
+    table,
+    cudf::structs::detail::flatten_nested_columns(
+      table, {}, {}, cudf::structs::detail::column_nullability::FORCE, cudf::get_default_stream()));
 }
 
 TYPED_TEST(TypedStructUtilitiesTest, NestedListsUnsupported)
@@ -74,7 +76,8 @@ TYPED_TEST(TypedStructUtilitiesTest, NestedListsUnsupported)
                  cudf::table_view{{nums_col, structs_col}},
                  {},
                  {},
-                 cudf::structs::detail::column_nullability::FORCE),
+                 cudf::structs::detail::column_nullability::FORCE,
+                 cudf::get_default_stream()),
                cudf::logic_error);
 }
 
@@ -90,9 +93,10 @@ TYPED_TEST(TypedStructUtilitiesTest, NoStructs)
 
   auto table = cudf::table_view{{nums_col, strings_col, nuther_nums_col}};
 
-  CUDF_TEST_EXPECT_TABLES_EQUAL(table,
-                                cudf::structs::detail::flatten_nested_columns(
-                                  table, {}, {}, cudf::structs::detail::column_nullability::FORCE));
+  CUDF_TEST_EXPECT_TABLES_EQUAL(
+    table,
+    cudf::structs::detail::flatten_nested_columns(
+      table, {}, {}, cudf::structs::detail::column_nullability::FORCE, cudf::get_default_stream()));
 }
 
 TYPED_TEST(TypedStructUtilitiesTest, SingleLevelStruct)
@@ -116,9 +120,10 @@ TYPED_TEST(TypedStructUtilitiesTest, SingleLevelStruct)
   auto expected = cudf::table_view{
     {expected_nums_col_1, expected_structs_col, expected_nums_col_2, expected_strings_col}};
 
-  CUDF_TEST_EXPECT_TABLES_EQUAL(expected,
-                                cudf::structs::detail::flatten_nested_columns(
-                                  table, {}, {}, cudf::structs::detail::column_nullability::FORCE));
+  CUDF_TEST_EXPECT_TABLES_EQUAL(
+    expected,
+    cudf::structs::detail::flatten_nested_columns(
+      table, {}, {}, cudf::structs::detail::column_nullability::FORCE, cudf::get_default_stream()));
 }
 
 TYPED_TEST(TypedStructUtilitiesTest, SingleLevelStructWithNulls)
@@ -144,9 +149,10 @@ TYPED_TEST(TypedStructUtilitiesTest, SingleLevelStructWithNulls)
   auto expected = cudf::table_view{
     {expected_nums_col_1, expected_structs_col, expected_nums_col_2, expected_strings_col}};
 
-  CUDF_TEST_EXPECT_TABLES_EQUAL(expected,
-                                cudf::structs::detail::flatten_nested_columns(
-                                  table, {}, {}, cudf::structs::detail::column_nullability::FORCE));
+  CUDF_TEST_EXPECT_TABLES_EQUAL(
+    expected,
+    cudf::structs::detail::flatten_nested_columns(
+      table, {}, {}, cudf::structs::detail::column_nullability::FORCE, cudf::get_default_stream()));
 }
 
 TYPED_TEST(TypedStructUtilitiesTest, StructOfStruct)
@@ -183,9 +189,10 @@ TYPED_TEST(TypedStructUtilitiesTest, StructOfStruct)
                                     expected_nums_col_3,
                                     expected_strings_col}};
 
-  CUDF_TEST_EXPECT_TABLES_EQUAL(expected,
-                                cudf::structs::detail::flatten_nested_columns(
-                                  table, {}, {}, cudf::structs::detail::column_nullability::FORCE));
+  CUDF_TEST_EXPECT_TABLES_EQUAL(
+    expected,
+    cudf::structs::detail::flatten_nested_columns(
+      table, {}, {}, cudf::structs::detail::column_nullability::FORCE, cudf::get_default_stream()));
 }
 
 TYPED_TEST(TypedStructUtilitiesTest, StructOfStructWithNullsAtLeafLevel)
@@ -223,9 +230,10 @@ TYPED_TEST(TypedStructUtilitiesTest, StructOfStructWithNullsAtLeafLevel)
                                     expected_nums_col_3,
                                     expected_strings_col}};
 
-  CUDF_TEST_EXPECT_TABLES_EQUAL(expected,
-                                cudf::structs::detail::flatten_nested_columns(
-                                  table, {}, {}, cudf::structs::detail::column_nullability::FORCE));
+  CUDF_TEST_EXPECT_TABLES_EQUAL(
+    expected,
+    cudf::structs::detail::flatten_nested_columns(
+      table, {}, {}, cudf::structs::detail::column_nullability::FORCE, cudf::get_default_stream()));
 }
 
 TYPED_TEST(TypedStructUtilitiesTest, StructOfStructWithNullsAtTopLevel)
@@ -264,9 +272,10 @@ TYPED_TEST(TypedStructUtilitiesTest, StructOfStructWithNullsAtTopLevel)
                                     expected_nums_col_3,
                                     expected_strings_col}};
 
-  CUDF_TEST_EXPECT_TABLES_EQUAL(expected,
-                                cudf::structs::detail::flatten_nested_columns(
-                                  table, {}, {}, cudf::structs::detail::column_nullability::FORCE));
+  CUDF_TEST_EXPECT_TABLES_EQUAL(
+    expected,
+    cudf::structs::detail::flatten_nested_columns(
+      table, {}, {}, cudf::structs::detail::column_nullability::FORCE, cudf::get_default_stream()));
 }
 
 TYPED_TEST(TypedStructUtilitiesTest, StructOfStructWithNullsAtAllLevels)
@@ -305,9 +314,10 @@ TYPED_TEST(TypedStructUtilitiesTest, StructOfStructWithNullsAtAllLevels)
                                     expected_nums_col_3,
                                     expected_strings_col}};
 
-  CUDF_TEST_EXPECT_TABLES_EQUAL(expected,
-                                cudf::structs::detail::flatten_nested_columns(
-                                  table, {}, {}, cudf::structs::detail::column_nullability::FORCE));
+  CUDF_TEST_EXPECT_TABLES_EQUAL(
+    expected,
+    cudf::structs::detail::flatten_nested_columns(
+      table, {}, {}, cudf::structs::detail::column_nullability::FORCE, cudf::get_default_stream()));
 }
 
 TYPED_TEST(TypedStructUtilitiesTest, ListsAreUnsupported)
@@ -327,7 +337,8 @@ TYPED_TEST(TypedStructUtilitiesTest, ListsAreUnsupported)
                  cudf::table_view{{structs_with_lists_col}},
                  {},
                  {},
-                 cudf::structs::detail::column_nullability::FORCE),
+                 cudf::structs::detail::column_nullability::FORCE,
+                 cudf::get_default_stream()),
                cudf::logic_error);
 }