From 81c2afb4439a2e6ffb75febd8b3d9c16650cd625 Mon Sep 17 00:00:00 2001
From: David Wendt <dwendt@nvidia.com>
Date: Fri, 28 Aug 2026 11:38:11 -0400
Subject: [PATCH] Fix segmented gather for sliced gather map input

---
 cpp/src/lists/copying/segmented_gather.cu         | 14 ++++++++++++++
 cpp/tests/copying/segmented_gather_list_tests.cpp | 14 +++++++-------
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/cpp/src/lists/copying/segmented_gather.cu b/cpp/src/lists/copying/segmented_gather.cu
index e73b7a48c278..fa5b21f7ee54 100644
--- a/cpp/src/lists/copying/segmented_gather.cu
+++ b/cpp/src/lists/copying/segmented_gather.cu
@@ -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 {
@@ -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()),
+                      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();
diff --git a/cpp/tests/copying/segmented_gather_list_tests.cpp b/cpp/tests/copying/segmented_gather_list_tests.cpp
index 1827275b328e..d0e8d0aa23ba 100644
--- a/cpp/tests/copying/segmented_gather_list_tests.cpp
+++ b/cpp/tests/copying/segmented_gather_list_tests.cpp
@@ -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>
@@ -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.
@@ -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());
   }
 }