Repository: rapidsai/cudf (language: python)

Issue:
[BUG] `cudaErrorIllegalAddress` in `purge_nonempty_nulls` when reading a wide `list<string>` Parquet table.
**Describe the bug**
A `cudaErrorIllegalAddress` is encountered when reading super-wide nullable `list<string>` table. In case the columns are non-nullable, the PQ reader succeeds but the output table isn't valid. This issue is ***not likely*** from #17059 or a problem with Parquet writer as the file can be correctly read with pyarrow. This issue is related to #16985 as it arises when `num_keys > 2billion` so my suspicion is that there is an 32 bit integer overflow at least at 1 place in the reader's pipeline.

For mixed nullable and non-nullable columns, the reader succeeds and we still don't get valid data with error something like the following when compared against the file read with pyarrow.

```python
AssertionError: DataFrame.iloc[:, 30678] (column name="_col30678") are different

DataFrame.iloc[:, 30678] (column name="_col30678") values are different (20.0 %)
[index]: [0, 1, 2, 3, 4]
[left]:  [[ovefylbpvcmgphlwxwgkdjxnpx, None, None, None], None, None, [axpgabkabirwcogoedpslxmdlumvldoeaobltkfm, None, oedps, lxmdlumvldoeaobltkfm, None], None]
[right]: [[ovefylbpvcmgphlwxwgkdjxnpx, None, None, None], None, None, [axpgabkabirwcog, None, oedps, lxmdlumvldoeaobltkfm, None], None]
At positional index 3, first diff: ['axpgabkabirwcogoedpslxmdlumvldoeaobltkfm', None, 'oedps', 'lxmdlumvldoeaobltkfm', None] != ['axpgabkabirwcog' None 'oedps' 'lxmdlumvldoeaobltkfm' None]
```


**Steps/Code to reproduce bug**

Modify `ParquetReaderTest.ListsWideTable` gtest as follows:

```cpp
TEST_F(ParquetReaderTest, ListsWideTable)
{
  auto constexpr num_rows = 5;
  auto constexpr num_cols = 32769;  // fails for 32769+
  auto constexpr seed     = 0xceed;

  std::mt19937 engine{seed};

  auto str_list = make_parquet_string_list_col(engine, num_rows, 5, 32, true);
  // auto str_list = make_parquet_string_list_col(engine, num_rows, 5, 32, false); // < without nulls. reader succeeds assertion fails (different number of rows read)

  std::vector<cudf::column_view> cols(num_cols, str_list->view());
  cudf::table_view expected(cols);

  // Use a host buffer for faster I/O
  auto const filepath = "wide_list_str.parquet";
  auto const out_opts =
    cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, expected).build();
  cudf::io::write_parquet(out_opts);

  cudf::io::parquet_reader_options default_in_opts =
    cudf::io::parquet_reader_options::builder(cudf::io::source_info(filepath));
  auto const [result, _] = cudf::io::read_parquet(default_in_opts);

  CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result->view());
}
```

**Expected behavior**
Wide `list<string>` tables should roundtrip perfectly. 

**Environment overview (please complete the following information)**
cudf `branch-24.12` cuda 12.5 conda-devcontainer on `dgx05` RDS Lab machine

**Environment details**
N/A.

**Additional context**
This issue also showed up while investigating #16985 for wide and deep structs. Reproduced as:

```python
import cudf
import os
import rmm
import pyarrow.parquet as pq
from numba import cuda

cuda.select_device(0)
rmm.mr.set_current_device_resource(rmm.mr.CudaAsyncMemoryResource())
os.environ["KVIKIO_COMPAT_MODE"] = "on"

def write_buf():
    # do this only once to get buffer.json for future
    df = cudf.read_parquet(
        "/datasets/gkimball/spark_json/20241001/part-00000-505e98e9-a5c8-4720-8bb4-d6cc96625744-c000.snappy.parquet"
    )
    print("cudf read input parquet")

    buf = StringIO(df["columnC"].str.cat(sep="\n", na_rep="{}"))

    with open("/home/coder/transcoded/buffer.json", "w") as fd:
        buf.seek(0)
        shutil.copyfileobj(buf, fd)

if __name__ == "__main__":
    # do this only once
    write_buf()

    # read json buffer with pandas
    df = cudf.read_json(
        "/home/coder/transcoded/buffer.json", engine="pandas", lines=True
    )
    df = cudf.DataFrame({"AENBHHGIABBBDDGOEI": df["AENBHHGIABBBDDGOEI"]})

    # roundtrip with pyarrow
    df = cudf.DataFrame.from_arrow(df.to_arrow())

    buf = BytesIO()
    df.to_parquet(buf)

    # read parquet fails while reading list nested as: `col0_697_1`
    _ = cudf.read_parquet(buf)



Discussion / hints:

Fix patch:
diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu
index 8cab68ea721..5138a92ac14 100644
--- a/cpp/src/io/parquet/reader_impl_preprocess.cu
+++ b/cpp/src/io/parquet/reader_impl_preprocess.cu
@@ -44,6 +44,7 @@
 #include <thrust/unique.h>
 
 #include <bitset>
+#include <limits>
 #include <numeric>
 
 namespace cudf::io::parquet::detail {
@@ -1592,36 +1593,68 @@ void reader::impl::allocate_columns(read_mode mode, size_t skip_rows, size_t num
     auto const d_cols_info = cudf::detail::make_device_uvector_async(
       h_cols_info, _stream, cudf::get_current_device_resource_ref());
 
-    auto const num_keys = _input_columns.size() * max_depth * subpass.pages.size();
-    // size iterator. indexes pages by sorted order
-    rmm::device_uvector<size_type> size_input{num_keys, _stream};
-    thrust::transform(
-      rmm::exec_policy(_stream),
-      thrust::make_counting_iterator<size_type>(0),
-      thrust::make_counting_iterator<size_type>(num_keys),
-      size_input.begin(),
-      get_page_nesting_size{
-        d_cols_info.data(), max_depth, subpass.pages.size(), subpass.pages.device_begin()});
-    auto const reduction_keys =
-      cudf::detail::make_counting_transform_iterator(0, get_reduction_key{subpass.pages.size()});
+    // Vector to store page sizes for each column at each depth
     cudf::detail::hostdevice_vector<size_t> sizes{_input_columns.size() * max_depth, _stream};
 
-    // find the size of each column
-    thrust::reduce_by_key(rmm::exec_policy(_stream),
-                          reduction_keys,
-                          reduction_keys + num_keys,
-                          size_input.cbegin(),
-                          thrust::make_discard_iterator(),
-                          sizes.d_begin());
-
-    // for nested hierarchies, compute per-page start offset
-    thrust::exclusive_scan_by_key(
-      rmm::exec_policy(_stream),
-      reduction_keys,
-      reduction_keys + num_keys,
-      size_input.cbegin(),
-      start_offset_output_iterator{
-        subpass.pages.device_begin(), 0, d_cols_info.data(), max_depth, subpass.pages.size()});
+    // Total number of keys to process
+    auto const num_keys = _input_columns.size() * max_depth * subpass.pages.size();
+
+    // Maximum 1 billion keys processed per iteration
+    auto constexpr max_keys_per_iter =
+      static_cast<size_t>(std::numeric_limits<size_type>::max() / 2);
+
+    // Number of keys for per each column
+    auto const num_keys_per_col = max_depth * subpass.pages.size();
+
+    // The largest multiple of `num_keys_per_col` that is <= `num_keys`
+    auto const num_keys_per_iter =
+      num_keys <= max_keys_per_iter
+        ? num_keys
+        : num_keys_per_col * std::max<size_t>(1, max_keys_per_iter / num_keys_per_col);
+
+    // Size iterator. Indexes pages by sorted order
+    rmm::device_uvector<size_type> size_input{num_keys_per_iter, _stream};
+
+    // To keep track of the starting key of an iteration
+    size_t key_start = 0;
+    // Loop until all keys are processed
+    while (key_start < num_keys) {
+      // Number of keys processed in this iteration
+      auto const num_keys_this_iter = std::min<size_t>(num_keys_per_iter, num_keys - key_start);
+      thrust::transform(
+        rmm::exec_policy_nosync(_stream),
+        thrust::make_counting_iterator<size_t>(key_start),
+        thrust::make_counting_iterator<size_t>(key_start + num_keys_this_iter),
+        size_input.begin(),
+        get_page_nesting_size{
+          d_cols_info.data(), max_depth, subpass.pages.size(), subpass.pages.device_begin()});
+
+      // Manually create a int64_t `key_start` compatible counting_transform_iterator to avoid
+      // implicit casting to size_type.
+      auto const reduction_keys = thrust::make_transform_iterator(
+        thrust::make_counting_iterator<size_t>(key_start), get_reduction_key{subpass.pages.size()});
+
+      // Find the size of each column
+      thrust::reduce_by_key(rmm::exec_policy_nosync(_stream),
+                            reduction_keys,
+                            reduction_keys + num_keys_this_iter,
+                            size_input.cbegin(),
+                            thrust::make_discard_iterator(),
+                            sizes.d_begin() + (key_start / subpass.pages.size()));
+
+      // For nested hierarchies, compute per-page start offset
+      thrust::exclusive_scan_by_key(rmm::exec_policy_nosync(_stream),
+                                    reduction_keys,
+                                    reduction_keys + num_keys_this_iter,
+                                    size_input.cbegin(),
+                                    start_offset_output_iterator{subpass.pages.device_begin(),
+                                                                 key_start,
+                                                                 d_cols_info.data(),
+                                                                 max_depth,
+                                                                 subpass.pages.size()});
+      // Increment the key_start
+      key_start += num_keys_this_iter;
+    }
 
     sizes.device_to_host_sync(_stream);
     for (size_type idx = 0; idx < static_cast<size_type>(_input_columns.size()); idx++) {


Test patch:
diff --git a/cpp/tests/io/parquet_reader_test.cpp b/cpp/tests/io/parquet_reader_test.cpp
index 4a5309f3ba7..ab4645c2e25 100644
--- a/cpp/tests/io/parquet_reader_test.cpp
+++ b/cpp/tests/io/parquet_reader_test.cpp
@@ -2724,3 +2724,38 @@ TYPED_TEST(ParquetReaderPredicatePushdownTest, FilterTyped)
   EXPECT_EQ(result_table.num_columns(), expected->num_columns());
   CUDF_TEST_EXPECT_TABLES_EQUAL(expected->view(), result_table);
 }
+
+TEST_F(ParquetReaderTest, ListsWideTable)
+{
+  auto constexpr num_rows = 2;
+  auto constexpr num_cols = 26'755;  // for slightly over 2B keys
+  auto constexpr seed     = 0xceed;
+
+  std::mt19937 engine{seed};
+
+  auto list_list       = make_parquet_list_list_col<int32_t>(0, num_rows, 1, 1, false);
+  auto list_list_nulls = make_parquet_list_list_col<int32_t>(0, num_rows, 1, 1, true);
+
+  // switch between nullable and non-nullable
+  std::vector<cudf::column_view> cols(num_cols);
+  bool with_nulls = false;
+  std::generate_n(cols.begin(), num_cols, [&]() {
+    auto const view = with_nulls ? list_list_nulls->view() : list_list->view();
+    with_nulls      = not with_nulls;
+    return view;
+  });
+
+  cudf::table_view expected(cols);
+
+  // Use a host buffer for faster I/O
+  std::vector<char> buffer;
+  auto const out_opts =
+    cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, expected).build();
+  cudf::io::write_parquet(out_opts);
+
+  cudf::io::parquet_reader_options default_in_opts =
+    cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer.data(), buffer.size()));
+  auto const [result, _] = cudf::io::read_parquet(default_in_opts);
+
+  CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result->view());
+}