Segmented apply_boolean_mask for LIST columns

cpp/CMakeLists.txt

@@ -359,6 +359,7 @@ add_library(
   src/join/mixed_join_size_kernel_nulls.cu
   src/join/mixed_join_size_kernels_semi.cu
   src/join/semi_join.cu
+  src/lists/apply_boolean_mask.cu
   src/lists/contains.cu
   src/lists/combine/concatenate_list_elements.cu
   src/lists/combine/concatenate_rows.cu

cpp/include/cudf/lists/detail/stream_compaction.hpp

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) 2022, 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
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cudf/column/column.hpp>
+#include <cudf/lists/lists_column_view.hpp>
+
+#include <rmm/mr/device/device_memory_resource.hpp>
+
+namespace cudf::lists::detail {
+
+/**
+ * @copydoc cudf::lists::apply_boolean_mask(lists_column_view const&, lists_column_view const&,
+ * rmm::mr::device_memory_resource*)
+ *
+ * @param stream CUDA stream used for device memory operations and kernel launches
+ */
+std::unique_ptr<column> apply_boolean_mask(
+  lists_column_view const& input,
+  lists_column_view const& boolean_mask,
+  rmm::cuda_stream_view stream,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
+}  // namespace cudf::lists::detail

cpp/include/cudf/lists/stream_compaction.hpp

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) 2022, 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
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cudf/column/column.hpp>
+#include <cudf/lists/lists_column_view.hpp>
+
+#include <rmm/mr/device/device_memory_resource.hpp>
+
+namespace cudf::lists {
+
+/**
+ * @brief Filters elements in each row of `input` LIST column using `boolean_mask`
+ * LIST of booleans as a mask.
+ *
+ * Given an input `LIST` column and a list-of-bools column, the function produces
+ * a new `LIST` column of the same type as `input`, where each element is copied
+ * from the input row *only* if the corresponding `boolean_mask` is non-null and `true`.
+ *
+ * E.g.
+ * @code{.pseudo}
+ * auto const input    = lcw<int32_t>{ {0,1,2}, {3,4}, {5,6,7}, {8,9} };
+ * auto const selector = lcw<bool>   { {0,1,1}, {1,0}, {1,1,1}, {0,0} };
+ * auto const results  = apply_boolean_mask(lists_column_view{input}, lists_column_view{selector});
+ * results             == { {1,2}, {3}, {5,6,7}, {} };
+ * @endcode
+ *
+ * `input` and `boolean_mask` must have the same number of rows.
+ * The output column has the same number of rows as the input column.
+ * An element is copied to an output row *only* if the corresponding bool selector is `true`.
+ * An output row is invalid only if the input row is invalid.
+ *
+ * @param input The input list column view to be filtered
+ * @param boolean_mask A nullable list of bools column used to filter `input` elements
+ * @param mr Device memory resource used to allocate the returned table's device memory
+ * @return List column of the same type as `input`, containing filtered list rows
+ */
+std::unique_ptr<column> apply_boolean_mask(
+  lists_column_view const& input,
+  lists_column_view const& boolean_mask,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
+}  // namespace cudf::lists

cpp/src/lists/apply_boolean_mask.cu

@@ -0,0 +1,111 @@
+/*
+ * Copyright (c) 2022, 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
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cudf/column/column_factories.hpp>
+#include <cudf/detail/copy.hpp>
+#include <cudf/detail/fill.hpp>
+#include <cudf/detail/iterator.cuh>
+#include <cudf/detail/null_mask.hpp>
+#include <cudf/detail/reduction_functions.hpp>
+#include <cudf/detail/replace.hpp>
+#include <cudf/detail/stream_compaction.hpp>
+#include <cudf/lists/detail/stream_compaction.hpp>
+#include <cudf/lists/stream_compaction.hpp>
+#include <cudf/utilities/bit.hpp>
+
+#include <rmm/exec_policy.hpp>
+
+#include <thrust/reduce.h>
+
+namespace cudf::lists {
+namespace detail {
+
+/**
+ * @copydoc cudf::lists::detail::apply_boolean_mask
+ */
+std::unique_ptr<column> apply_boolean_mask(lists_column_view const& input,
+                                           lists_column_view const& boolean_mask,
+                                           rmm::cuda_stream_view stream,
+                                           rmm::mr::device_memory_resource* mr)
+{
+  CUDF_EXPECTS(boolean_mask.child().type().id() == type_id::BOOL8, "Mask must be of type BOOL8.");
+  CUDF_EXPECTS(input.size() == boolean_mask.size(),
+               "Boolean masks column must have same number of rows as input.");
+  auto const num_rows = input.size();
+
+  if (num_rows == 0) { return cudf::empty_like(input.parent()); }
+
+  auto constexpr offset_data_type = data_type{type_id::INT32};
+
+  auto filtered_child = [&] {
+    auto filtered =
+      cudf::detail::apply_boolean_mask(cudf::table_view{{input.get_sliced_child(stream)}},
+                                       boolean_mask.get_sliced_child(stream),
+                                       stream,
+                                       mr)
+        ->release();
+    return std::move(filtered.front());
+  };
+
+  auto output_offsets = [&] {
+    auto boolean_mask_sliced_offsets =
+      cudf::detail::slice(
+        boolean_mask.offsets(), {boolean_mask.offset(), boolean_mask.size() + 1}, stream)
+        .front();
+    auto const sizes       = cudf::reduction::segmented_sum(boolean_mask.get_sliced_child(stream),
+                                                      boolean_mask_sliced_offsets,
+                                                      offset_data_type,
+                                                      null_policy::EXCLUDE,
+                                                      stream);
+    auto const d_sizes     = column_device_view::create(*sizes, stream);
+    auto const sizes_begin = cudf::detail::make_null_replacement_iterator(*d_sizes, offset_type{0});
+    auto const sizes_end   = sizes_begin + sizes->size();
+    auto output_offsets    = cudf::make_numeric_column(
+      offset_data_type, num_rows + 1, mask_state::UNALLOCATED, stream, mr);
+    auto output_offsets_view = output_offsets->mutable_view();
+
+    // Could have attempted an exclusive_scan(), but it would not compute the last entry.
+    // Instead, inclusive_scan(), followed by writing `0` to the head of the offsets column.
+    thrust::inclusive_scan(rmm::exec_policy(stream),
+                           sizes_begin,
+                           sizes_end,
+                           output_offsets_view.begin<offset_type>() + 1);
+    CUDF_CUDA_TRY(cudaMemsetAsync(
+      output_offsets_view.begin<offset_type>(), 0, sizeof(offset_type), stream.value()));
+    return output_offsets;
+  };
+
+  return cudf::make_lists_column(input.size(),
+                                 output_offsets(),
+                                 filtered_child(),
+                                 input.null_count(),
+                                 cudf::detail::copy_bitmask(input.parent(), stream, mr),
+                                 stream,
+                                 mr);
+}
+}  // namespace detail
+
+/**
+ * @copydoc cudf::lists::apply_boolean_mask
+ */
+std::unique_ptr<column> apply_boolean_mask(lists_column_view const& input,
+                                           lists_column_view const& boolean_mask,
+                                           rmm::mr::device_memory_resource* mr)
+{
+  return detail::apply_boolean_mask(input, boolean_mask, rmm::cuda_stream_default, mr);
+}
+
+}  // namespace cudf::lists

cpp/tests/CMakeLists.txt

@@ -469,6 +469,7 @@ ConfigureTest(AST_TEST ast/transform_tests.cpp)
 # * lists tests ----------------------------------------------------------------------------------
 ConfigureTest(
   LISTS_TEST
+  lists/apply_boolean_mask_test.cpp
   lists/combine/concatenate_list_elements_tests.cpp
   lists/combine/concatenate_rows_tests.cpp
   lists/contains_tests.cpp

cpp/tests/lists/apply_boolean_mask_test.cpp

@@ -0,0 +1,169 @@
+/*
+ * Copyright (c) 2022, 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
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <cudf/column/column_factories.hpp>
+#include <cudf/detail/null_mask.hpp>
+#include <cudf/lists/extract.hpp>
+#include <cudf/lists/stream_compaction.hpp>
+
+#include <cudf_test/base_fixture.hpp>
+#include <cudf_test/column_utilities.hpp>
+#include <cudf_test/column_wrapper.hpp>
+#include <cudf_test/iterator_utilities.hpp>
+#include <cudf_test/type_lists.hpp>
+
+namespace cudf::test {
+
+using namespace iterators;
+using cudf::lists_column_view;
+using cudf::lists::apply_boolean_mask;
+
+template <typename T>
+using lists    = lists_column_wrapper<T, int32_t>;
+using filter_t = lists_column_wrapper<bool, int32_t>;
+
+auto constexpr X = int32_t{0};  // Placeholder for NULL.
+
+struct ApplyBooleanMaskTest : public BaseFixture {
+};
+
+template <typename T>
+struct ApplyBooleanMaskTypedTest : ApplyBooleanMaskTest {
+};
+
+TYPED_TEST_SUITE(ApplyBooleanMaskTypedTest, cudf::test::NumericTypes);
+
+TYPED_TEST(ApplyBooleanMaskTypedTest, StraightLine)
+{
+  using T    = TypeParam;
+  auto input = lists<T>{{0, 1, 2, 3}, {4, 5}, {6, 7, 8, 9}, {0, 1}, {2, 3, 4, 5}, {6, 7}}.release();
+  auto filter = filter_t{{1, 0, 1, 0}, {1, 0}, {1, 0, 1, 0}, {1, 0}, {1, 0, 1, 0}, {1, 0}};
+
+  {
+    // Unsliced.
+    auto filtered = apply_boolean_mask(lists_column_view{*input}, lists_column_view{filter});
+    auto expected = lists<T>{{0, 2}, {4}, {6, 8}, {0}, {2, 4}, {6}};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*filtered, expected);
+  }
+  {
+    // Sliced input: Remove the first row.
+    auto sliced = cudf::slice(*input, {1, input->size()}).front();
+    //           == lists_t {{4, 5}, {6, 7, 8, 9}, {0, 1}, {2, 3, 4, 5}, {6, 7}};
+    auto filter   = filter_t{{0, 1}, {0, 1, 0, 1}, {1, 1}, {0, 1, 0, 1}, {0, 0}};
+    auto filtered = apply_boolean_mask(lists_column_view{sliced}, lists_column_view{filter});
+    auto expected = lists<T>{{5}, {7, 9}, {0, 1}, {3, 5}, {}};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*filtered, expected);
+  }
+}
+
+TYPED_TEST(ApplyBooleanMaskTypedTest, WithNullElements)
+{
+  using T = TypeParam;
+  auto input =
+    lists<T>{
+      {0, 1, 2, 3},
+      lists<T>{{X, 5}, null_at(0)},
+      {6, 7, 8, 9},
+      {0, 1},
+      lists<T>{{X, 3, 4, X}, nulls_at({0, 3})},
+      lists<T>{{X, X}, nulls_at({0, 1})},
+    }
+      .release();
+  auto filter = filter_t{{1, 0, 1, 0}, {1, 0}, {1, 0, 1, 0}, {1, 0}, {1, 0, 1, 0}, {1, 0}};
+
+  {
+    // Unsliced.
+    auto filtered = apply_boolean_mask(lists_column_view{*input}, lists_column_view{filter});
+    auto expected = lists<T>{{0, 2},
+                             lists<T>{{X}, null_at(0)},
+                             {6, 8},
+                             {0},
+                             lists<T>{{X, 4}, null_at(0)},
+                             lists<T>{{X}, null_at(0)}};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*filtered, expected);
+  }
+  {
+    // Sliced input: Remove the first row.
+    auto sliced = cudf::slice(*input, {1, input->size()}).front();
+    //           == lists_t {{X, 5}, {6, 7, 8, 9}, {0, 1}, {X, 3, 4, X}, {X, X}};
+    auto filter   = filter_t{{0, 1}, {0, 1, 0, 1}, {1, 1}, {0, 1, 0, 1}, {0, 0}};
+    auto filtered = apply_boolean_mask(lists_column_view{sliced}, lists_column_view{filter});
+    auto expected = lists<T>{{5}, {7, 9}, {0, 1}, lists<T>{{3, X}, null_at(1)}, {}};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*filtered, expected);
+  }
+}
+
+TYPED_TEST(ApplyBooleanMaskTypedTest, NullInTheInput)
+{
+  using T = TypeParam;
+  auto input =
+    lists<T>{{{0, 1, 2, 3}, {}, {6, 7, 8, 9}, {}, {2, 3, 4, 5}, {6, 7}}, nulls_at({1, 3})}
+      .release();
+  auto filter = filter_t{{1, 0, 1, 0}, {}, {1, 0, 1, 0}, {}, {1, 0, 1, 0}, {1, 0}};
+
+  {
+    // Unsliced.
+    auto filtered = apply_boolean_mask(lists_column_view{*input}, lists_column_view{filter});
+    auto expected = lists<T>{{{0, 2}, {}, {6, 8}, {}, {2, 4}, {6}}, nulls_at({1, 3})};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*filtered, expected);
+  }
+  {
+    // Sliced input: Remove the first row.
+    auto sliced = cudf::slice(*input, {1, input->size()}).front();
+    //           == lists_t{{{}, {6, 7, 8, 9}, {}, {2, 3, 4, 5}, {6, 7}}, nulls_at({0,2})};
+    auto filter   = filter_t{{}, {0, 1, 0, 1}, {}, {0, 1, 0, 1}, {0, 0}};
+    auto filtered = apply_boolean_mask(lists_column_view{sliced}, lists_column_view{filter});
+    auto expected = lists<T>{{{}, {7, 9}, {}, {3, 5}, {}}, nulls_at({0, 2})};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*filtered, expected);
+  }
+  {
+    // Sliced input: Remove the first two rows.
+    auto sliced = cudf::slice(*input, {2, input->size()}).front();
+    //           == lists_t{{{6, 7, 8, 9}, {}, {2, 3, 4, 5}, {6, 7}}, null_at(1)};
+    auto filter   = filter_t{{0, 1, 0, 1}, {}, {0, 1, 0, 1}, {0, 0}};
+    auto filtered = apply_boolean_mask(lists_column_view{sliced}, lists_column_view{filter});
+    auto expected = lists<T>{{{7, 9}, {}, {3, 5}, {}}, null_at(1)};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*filtered, expected);
+  }
+}
+
+TEST_F(ApplyBooleanMaskTest, Trivial)
+{
+  auto const input  = lists<int32_t>{};
+  auto const filter = filter_t{};
+  auto const result = apply_boolean_mask(lists_column_view{input}, lists_column_view{filter});
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*result, lists<int32_t>{});
+}
+
+TEST_F(ApplyBooleanMaskTest, Failure)
+{
+  {
+    // Invalid mask type.
+    auto const input  = lists<int32_t>{{1, 2, 3}, {4, 5, 6}};
+    auto const filter = lists<int32_t>{{0, 0, 0}};
+    CUDF_EXPECT_THROW_MESSAGE(
+      apply_boolean_mask(lists_column_view{input}, lists_column_view{filter}),
+      "Mask must be of type BOOL8.");
+  }
+  {
+    // Mismatched number of rows.
+    auto const input  = lists<int32_t>{{1, 2, 3}, {4, 5, 6}};
+    auto const filter = filter_t{{0, 0, 0}};
+    CUDF_EXPECT_THROW_MESSAGE(
+      apply_boolean_mask(lists_column_view{input}, lists_column_view{filter}),
+      "Boolean masks column must have same number of rows as input.");
+  }
+}
+}  // namespace cudf::test