Skip to content

Commit

Permalink
Expose stream-ordering in column view APIs (#17434)
Browse files Browse the repository at this point in the history
Adds stream parameter to
```
cudf::detail::column_view_base::null_count(begin, end)
cudf::detail::column_view_base::has_nulls(begin, end)
```
Note: Since stream-ordered prefetching is [back-logged](#17434 (comment)), we defer modifying the `get_data` member functions to accept a stream parameter for now. 

Reference: 
1. #13744
2. #16251 (comment)

Authors:
  - Shruti Shivakumar (https://github.com/shrshi)

Approvers:
  - Nghia Truong (https://github.com/ttnghia)
  - Bradley Dice (https://github.com/bdice)

URL: #17434
  • Loading branch information
shrshi authored Jan 16, 2025
1 parent 7f2b2ba commit a4bbd09
Show file tree
Hide file tree
Showing 6 changed files with 129 additions and 15 deletions.
15 changes: 11 additions & 4 deletions cpp/include/cudf/column/column_view.hpp
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2019-2024, NVIDIA CORPORATION.
* Copyright (c) 2019-2025, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -176,9 +176,13 @@ class column_view_base {
*
* @param[in] begin The starting index of the range (inclusive).
* @param[in] end The index of the last element in the range (exclusive).
* @param[in] stream CUDA stream used for device memory operations and kernel launches
* @return The count of null elements in the given range
*/
[[nodiscard]] size_type null_count(size_type begin, size_type end) const;
[[nodiscard]] size_type null_count(
size_type begin,
size_type end,
rmm::cuda_stream_view stream = cudf::get_default_stream()) const;

/**
* @brief Indicates if the column contains null elements,
Expand All @@ -198,12 +202,15 @@ class column_view_base {
*
* @param begin The starting index of the range (inclusive).
* @param end The index of the last element in the range (exclusive).
* @param stream CUDA stream used for device memory operations and kernel launches
* @return true One or more elements are null in the range [begin, end)
* @return false All elements are valid in the range [begin, end)
*/
[[nodiscard]] bool has_nulls(size_type begin, size_type end) const
[[nodiscard]] bool has_nulls(size_type begin,
size_type end,
rmm::cuda_stream_view stream = cudf::get_default_stream()) const
{
return null_count(begin, end) > 0;
return null_count(begin, end, stream) > 0;
}

/**
Expand Down
9 changes: 5 additions & 4 deletions cpp/src/column/column_view.cpp
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2019-2024, NVIDIA CORPORATION.
* Copyright (c) 2019-2025, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -141,13 +141,14 @@ column_view_base::column_view_base(data_type type,
}
}

size_type column_view_base::null_count(size_type begin, size_type end) const
size_type column_view_base::null_count(size_type begin,
size_type end,
rmm::cuda_stream_view stream) const
{
CUDF_EXPECTS((begin >= 0) && (end <= size()) && (begin <= end), "Range is out of bounds.");
return (null_count() == 0)
? 0
: cudf::detail::null_count(
null_mask(), offset() + begin, offset() + end, cudf::get_default_stream());
: cudf::detail::null_count(null_mask(), offset() + begin, offset() + end, stream);
}

bool is_shallow_equivalent(column_view const& lhs, column_view const& rhs)
Expand Down
4 changes: 2 additions & 2 deletions cpp/src/copying/copy_range.cu
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2019-2024, NVIDIA CORPORATION.
* Copyright (c) 2019-2025, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -103,7 +103,7 @@ struct out_of_place_copy_range_dispatch {
rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref())
{
auto p_ret = std::make_unique<cudf::column>(target, stream, mr);
if ((!p_ret->nullable()) && source.has_nulls(source_begin, source_end)) {
if ((!p_ret->nullable()) && source.has_nulls(source_begin, source_end, stream)) {
p_ret->set_null_mask(
cudf::detail::create_null_mask(p_ret->size(), cudf::mask_state::ALL_VALID, stream, mr), 0);
}
Expand Down
12 changes: 7 additions & 5 deletions cpp/src/rolling/grouped_rolling.cu
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020-2024, NVIDIA CORPORATION.
* Copyright (c) 2020-2025, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -419,7 +419,7 @@ template <typename T, CUDF_ENABLE_IF(cudf::column_device_view::has_element_acces
/// at the beginning of the column or at the end.
/// If no null values are founds, null_begin and null_end are 0.
std::tuple<size_type, size_type> get_null_bounds_for_orderby_column(
column_view const& orderby_column)
column_view const& orderby_column, rmm::cuda_stream_view stream)
{
auto const num_rows = orderby_column.size();
auto const num_nulls = orderby_column.null_count();
Expand All @@ -429,7 +429,7 @@ std::tuple<size_type, size_type> get_null_bounds_for_orderby_column(
return std::make_tuple(0, num_nulls);
}

auto const first_row_is_null = orderby_column.null_count(0, 1) == 1;
auto const first_row_is_null = orderby_column.null_count(0, 1, stream) == 1;

return first_row_is_null ? std::make_tuple(0, num_nulls)
: std::make_tuple(num_rows - num_nulls, num_rows);
Expand All @@ -451,7 +451,8 @@ std::unique_ptr<column> range_window_ASC(column_view const& input,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
auto [h_nulls_begin_idx, h_nulls_end_idx] = get_null_bounds_for_orderby_column(orderby_column);
auto [h_nulls_begin_idx, h_nulls_end_idx] =
get_null_bounds_for_orderby_column(orderby_column, stream);
auto const p_orderby_device_view = cudf::column_device_view::create(orderby_column, stream);

auto const preceding_calculator = cuda::proclaim_return_type<size_type>(
Expand Down Expand Up @@ -740,7 +741,8 @@ std::unique_ptr<column> range_window_DESC(column_view const& input,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
auto [h_nulls_begin_idx, h_nulls_end_idx] = get_null_bounds_for_orderby_column(orderby_column);
auto [h_nulls_begin_idx, h_nulls_end_idx] =
get_null_bounds_for_orderby_column(orderby_column, stream);
auto const p_orderby_device_view = cudf::column_device_view::create(orderby_column, stream);

auto const preceding_calculator = cuda::proclaim_return_type<size_type>(
Expand Down
1 change: 1 addition & 0 deletions cpp/tests/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -698,6 +698,7 @@ if(CUDF_BUILD_STREAMS_TEST_UTIL)
endif()

ConfigureTest(STREAM_BINARYOP_TEST streams/binaryop_test.cpp STREAM_MODE testing)
ConfigureTest(STREAM_COLUMN_VIEW_TEST streams/column_view_test.cpp STREAM_MODE testing)
ConfigureTest(STREAM_CONCATENATE_TEST streams/concatenate_test.cpp STREAM_MODE testing)
ConfigureTest(STREAM_COPYING_TEST streams/copying_test.cpp STREAM_MODE testing)
ConfigureTest(STREAM_CSVIO_TEST streams/io/csv_test.cpp STREAM_MODE testing)
Expand Down
103 changes: 103 additions & 0 deletions cpp/tests/streams/column_view_test.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
/*
* Copyright (c) 2025, 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_test/base_fixture.hpp>
#include <cudf_test/column_wrapper.hpp>
#include <cudf_test/default_stream.hpp>
#include <cudf_test/type_lists.hpp>

#include <cudf/column/column_view.hpp>
#include <cudf/null_mask.hpp>
#include <cudf/transform.hpp>

#include <random>
#include <vector>

template <typename T>
struct TypedColumnTest : public cudf::test::BaseFixture {
cudf::data_type type() { return cudf::data_type{cudf::type_to_id<T>()}; }

TypedColumnTest(rmm::cuda_stream_view stream = cudf::test::get_default_stream())
: data{_num_elements * cudf::size_of(type()), stream},
mask{cudf::bitmask_allocation_size_bytes(_num_elements), stream}
{
std::vector<char> h_data(std::max(data.size(), mask.size()));
std::iota(h_data.begin(), h_data.end(), 0);
CUDF_CUDA_TRY(
cudaMemcpyAsync(data.data(), h_data.data(), data.size(), cudaMemcpyDefault, stream.value()));
CUDF_CUDA_TRY(
cudaMemcpyAsync(mask.data(), h_data.data(), mask.size(), cudaMemcpyDefault, stream.value()));
}

cudf::size_type num_elements() { return _num_elements; }

std::random_device r;
std::default_random_engine generator{r()};
std::uniform_int_distribution<cudf::size_type> distribution{200, 1000};
cudf::size_type _num_elements{distribution(generator)};
rmm::device_buffer data{};
rmm::device_buffer mask{};
rmm::device_buffer all_valid_mask{create_null_mask(
num_elements(), cudf::mask_state::ALL_VALID, cudf::test::get_default_stream())};
rmm::device_buffer all_null_mask{
create_null_mask(num_elements(), cudf::mask_state::ALL_NULL, cudf::test::get_default_stream())};
};

TYPED_TEST_SUITE(TypedColumnTest, cudf::test::Types<int32_t>);

/**
* @brief Verifies equality of the properties and data of a `column`'s views.
*
* @param col The `column` to verify
*/
void verify_column_views(cudf::column& col)
{
cudf::column_view view = col;
cudf::mutable_column_view mutable_view = col;
EXPECT_EQ(col.type(), view.type());
EXPECT_EQ(col.type(), mutable_view.type());
EXPECT_EQ(col.size(), view.size());
EXPECT_EQ(col.size(), mutable_view.size());
EXPECT_EQ(col.null_count(), view.null_count());
EXPECT_EQ(col.null_count(), mutable_view.null_count());
EXPECT_EQ(view.null_count(0, col.size(), cudf::test::get_default_stream()),
mutable_view.null_count(0, col.size(), cudf::test::get_default_stream()));
EXPECT_EQ(view.has_nulls(0, col.size(), cudf::test::get_default_stream()),
mutable_view.has_nulls(0, col.size(), cudf::test::get_default_stream()));
EXPECT_EQ(col.null_count(), mutable_view.null_count());
EXPECT_EQ(col.nullable(), view.nullable());
EXPECT_EQ(col.nullable(), mutable_view.nullable());
EXPECT_EQ(col.num_children(), view.num_children());
EXPECT_EQ(col.num_children(), mutable_view.num_children());
EXPECT_EQ(view.head(), mutable_view.head());
EXPECT_EQ(view.data<char>(), mutable_view.data<char>());
EXPECT_EQ(view.offset(), mutable_view.offset());
}

TYPED_TEST(TypedColumnTest, CopyConstructorWithMask)
{
cudf::column original{
this->type(), this->num_elements(), std::move(this->data), std::move(this->all_valid_mask), 0};
cudf::column copy{original, cudf::test::get_default_stream()};
verify_column_views(copy);
CUDF_TEST_EXPECT_COLUMNS_EQUAL(original, copy);

// Verify deep copy
cudf::column_view original_view = original;
cudf::column_view copy_view = copy;
EXPECT_NE(original_view.head(), copy_view.head());
EXPECT_NE(original_view.null_mask(), copy_view.null_mask());
}

0 comments on commit a4bbd09

Please sign in to comment.