Fix segmented_gather() for null LIST rows (#9537)

`segmented_gather()` currently assumes that null LIST rows also have
a `0` size (as defined by the difference of adjacent offsets.)

This might not hold, for example, for LIST columns that are members
of STRUCT columns whose parent null masks are superimposed on its
children. This would cause a non-empty list row to be marked null,
without compaction. This leads to errors in fetching elements of a
list row as seen in NVIDIA/spark-rapids/pull/3770.

This commit adds the handling of uncompacted LIST rows in
`segmented_gather()`.

Authors:
  - MithunR (https://github.com/mythrocks)

Approvers:
  - Conor Hoekstra (https://github.com/codereport)
  - Nghia Truong (https://github.com/ttnghia)
  - David Wendt (https://github.com/davidwendt)

URL: #9537

cpp/src/lists/copying/segmented_gather.cu

@@ -45,14 +45,16 @@ std::unique_ptr<column> segmented_gather(lists_column_view const& value_column,
   auto const gather_index_begin      = gather_map.offsets_begin() + 1;
   auto const gather_index_end        = gather_map.offsets_end();
   auto const value_offsets           = value_column.offsets_begin();
+  auto const value_device_view       = column_device_view::create(value_column.parent(), stream);
   auto const map_begin =
     cudf::detail::indexalator_factory::make_input_iterator(gather_map_sliced_child);
   auto const out_of_bounds = [] __device__(auto const index, auto const list_size) {
     return index >= list_size || (index < 0 && -index > list_size);
   };
 
   // Calculate Flattened gather indices  (value_offset[row]+sub_index
-  auto transformer = [value_offsets,
+  auto transformer = [values_lists_view = *value_device_view,
+                      value_offsets,
                       map_begin,
                       gather_index_begin,
                       gather_index_end,
@@ -64,8 +66,9 @@ std::unique_ptr<column> segmented_gather(lists_column_view const& value_column,
         thrust::seq, gather_index_begin, gather_index_end, gather_index_begin[-1] + index) -
       gather_index_begin;
     // Get each sub_index in list in each row of gather_map.
-    auto sub_index          = map_begin[index];
-    auto list_size          = value_offsets[offset_idx + 1] - value_offsets[offset_idx];
+    auto sub_index    = map_begin[index];
+    auto list_is_null = values_lists_view.is_null(offset_idx);
+    auto list_size = list_is_null ? 0 : (value_offsets[offset_idx + 1] - value_offsets[offset_idx]);
     auto wrapped_sub_index  = sub_index < 0 ? sub_index + list_size : sub_index;
     auto constexpr null_idx = cuda::std::numeric_limits<cudf::size_type>::max();
     // Add sub_index to value_column offsets, to get gather indices of child of value_column

cpp/tests/copying/segmented_gather_list_tests.cpp

@@ -16,6 +16,7 @@
 #include <cudf/column/column_view.hpp>
 #include <cudf/copying.hpp>
 #include <cudf/detail/iterator.cuh>
+#include <cudf/detail/null_mask.hpp>
 #include <cudf/lists/detail/gather.cuh>
 #include <cudf/lists/lists_column_view.hpp>
 
@@ -34,19 +35,14 @@ using FixedWidthTypesNotBool = cudf::test::Concat<cudf::test::IntegralTypesNotBo
                                                   cudf::test::TimestampTypes>;
 TYPED_TEST_SUITE(SegmentedGatherTest, FixedWidthTypesNotBool);
 
-class SegmentedGatherTestList : public cudf::test::BaseFixture {
-};
-
 // to disambiguate between {} == 0 and {} == List{0}
 // Also, see note about compiler issues when declaring nested
 // empty lists in lists_column_wrapper documentation
 template <typename T>
 using LCW = cudf::test::lists_column_wrapper<T, int32_t>;
 using cudf::lists_column_view;
 using cudf::lists::detail::segmented_gather;
-using cudf::test::iterators::no_nulls;
-using cudf::test::iterators::null_at;
-using cudf::test::iterators::nulls_at;
+using namespace cudf::test::iterators;
 auto constexpr NULLIFY = cudf::out_of_bounds_policy::NULLIFY;
 
 TYPED_TEST(SegmentedGatherTest, Gather)
@@ -300,6 +296,25 @@ TYPED_TEST(SegmentedGatherTest, GatherNegatives)
   }
 }
 
+TYPED_TEST(SegmentedGatherTest, GatherOnNonCompactedNullLists)
+{
+  using T          = TypeParam;
+  auto constexpr X = -1;  // Signifies null value.
+
+  // List<T>
+  auto list = LCW<T>{{{1, 2, 3, 4}, {5}, {6, 7}, {8, 9, 0}, {}, {1, 2}, {3, 4, 5}}, no_nulls()};
+  auto const input = list.release();
+
+  // Set non-empty list row at index 5 to null.
+  cudf::detail::set_null_mask(input->mutable_view().null_mask(), 5, 6, false);
+
+  auto const gather_map = LCW<int>{{-1, 2, 1, -4}, {0}, {-2, 1}, {0, 2, 1}, {}, {0}, {1, 2}};
+  auto const expected =
+    LCW<T>{{{4, 3, 2, 1}, {5}, {6, 7}, {8, 0, 9}, {}, {{X}, all_nulls()}, {4, 5}}, null_at(5)};
+  auto const results = segmented_gather(lists_column_view{*input}, lists_column_view{gather_map});
+  CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(results->view(), expected);
+}
+
 TYPED_TEST(SegmentedGatherTest, GatherNestedNulls)
 {
   using T = TypeParam;

cpp/tests/lists/extract_tests.cpp

@@ -15,6 +15,7 @@
  */
 
 #include <cudf/column/column_factories.hpp>
+#include <cudf/detail/null_mask.hpp>
 #include <cudf/lists/extract.hpp>
 
 #include <tests/strings/utilities.h>
@@ -213,6 +214,34 @@ TYPED_TEST(ListsExtractNumericsTest, ExtractElementNestedLists)
   }
 }
 
+TYPED_TEST(ListsExtractNumericsTest, ExtractElementsFromNonCompactedNullLists)
+{
+  using namespace cudf::test::iterators;
+  using indices       = cudf::test::fixed_width_column_wrapper<cudf::size_type>;
+  using lcw           = cudf::test::lists_column_wrapper<TypeParam, int32_t>;
+  using result_column = cudf::test::fixed_width_column_wrapper<TypeParam, int32_t>;
+  auto constexpr X    = -1;  // Value indicating null.
+
+  auto input =
+    lcw{{{1, 2, 3}, {4, 5, 6}, {}, {7, 8, 9}, {0, 1, 2}, {}, {3, 4, 5}}, nulls_at({2, 5})}
+      .release();
+
+  // Set null at index 4.
+  cudf::detail::set_null_mask(input->mutable_view().null_mask(), 4, 5, false);
+
+  {
+    auto result   = cudf::lists::extract_list_element(cudf::lists_column_view{*input}, 0);
+    auto expected = result_column{{1, 4, X, 7, X, X, 3}, nulls_at({2, 4, 5})};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, *result);
+  }
+  {
+    auto index    = indices{0, 1, 2, 0, 1, 2, 0};
+    auto result   = cudf::lists::extract_list_element(cudf::lists_column_view{*input}, index);
+    auto expected = result_column{{1, 5, X, 7, X, X, 3}, nulls_at({2, 4, 5})};
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, *result);
+  }
+}
+
 TEST_F(ListsExtractTest, ExtractElementEmpty)
 {
   using LCW = cudf::test::lists_column_wrapper<cudf::string_view>;