Skip to content

Fix segmented gather for sliced gather map input #23876

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
davidwendt wants to merge 1 commit into
base: main
Choose a base branch
from
+21 −7
Open

Fix segmented gather for sliced gather map input #23876

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions cpp/src/lists/copying/segmented_gather.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,13 @@
#include <cudf/utilities/default_stream.hpp>
#include <cudf/utilities/memory_resource.hpp>

#include <rmm/exec_policy.hpp>

#include <cuda/functional>
#include <cuda/stream>
#include <thrust/binary_search.h>
#include <thrust/execution_policy.h>
#include <thrust/transform.h>

namespace cudf {
namespace lists {
Expand Down Expand Up @@ -95,6 +98,17 @@ std::unique_ptr<column> segmented_gather(lists_column_view const& value_column,
gather_map.offset() + output_offset_view.size(),
0,
stream);
// When gather_map is a sliced view, the copied offsets start at a nonzero base.
// Subtract it so that output_offset[0] == 0.
if (gather_map.offset() > 0) {
auto* output_data = output_offset_view.data<int32_t>();
auto const* base_ptr = gather_map.offsets_begin();
thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
Comment thread
coderabbitai[bot] marked this conversation as resolved.
output_data,
output_data + output_offset_view.size(),
output_data,
[base_ptr] __device__(int32_t val) { return val - *base_ptr; });
}
// Assemble list column & return
auto null_mask = cudf::detail::copy_bitmask(value_column.parent(), stream, mr);
size_type null_count = value_column.null_count();
Expand Down
14 changes: 7 additions & 7 deletions cpp/tests/copying/segmented_gather_list_tests.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/
#include <cudf_test/base_fixture.hpp>
Expand Down Expand Up @@ -520,13 +520,13 @@ TEST_F(SegmentedGatherTestFloat, GatherMapSliced)

auto result0 = cudf::lists::segmented_gather(cudf::lists_column_view{sliced[0]},
cudf::lists_column_view{split_m[0]});
CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(split_e[0], result0->view());
CUDF_TEST_EXPECT_COLUMNS_EQUAL(split_e[0], result0->view());
auto result1 = cudf::lists::segmented_gather(cudf::lists_column_view{sliced[1]},
cudf::lists_column_view{split_m[1]});
CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(split_e[1], result1->view());
CUDF_TEST_EXPECT_COLUMNS_EQUAL(split_e[1], result1->view());
auto result2 = cudf::lists::segmented_gather(cudf::lists_column_view{sliced[2]},
cudf::lists_column_view{split_m[2]});
CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(split_e[2], result2->view());
CUDF_TEST_EXPECT_COLUMNS_EQUAL(split_e[2], result2->view());
}

// List<T>, with out-of-bounds gather indices.
Expand All @@ -546,13 +546,13 @@ TEST_F(SegmentedGatherTestFloat, GatherMapSliced)

auto const result0 = cudf::lists::segmented_gather(
cudf::lists_column_view{sliced[0]}, cudf::lists_column_view{split_m[0]}, NULLIFY);
CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(split_e[0], result0->view());
CUDF_TEST_EXPECT_COLUMNS_EQUAL(split_e[0], result0->view());
auto const result1 = cudf::lists::segmented_gather(
cudf::lists_column_view{sliced[1]}, cudf::lists_column_view{split_m[1]}, NULLIFY);
CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(split_e[1], result1->view());
CUDF_TEST_EXPECT_COLUMNS_EQUAL(split_e[1], result1->view());
auto const result2 = cudf::lists::segmented_gather(
cudf::lists_column_view{sliced[2]}, cudf::lists_column_view{split_m[2]}, NULLIFY);
CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(split_e[2], result2->view());
CUDF_TEST_EXPECT_COLUMNS_EQUAL(split_e[2], result2->view());
}
}

Expand Down
Loading