Fix

Author: mhaseeb123

cpp/src/io/parquet/parquet_gpu.hpp

@@ -589,8 +589,8 @@ struct EncColumnChunk {
   size_type num_dict_entries;  //!< Total number of entries in dictionary
   size_type
     uniq_data_size;  //!< Size of dictionary page (set of all unique values) if dict enc is used
-  size_type plain_data_size;  //!< Size of data in this chunk if plain encoding is used
-  size_type* dict_data;       //!< Dictionary data (unique row indices)
+  size_t plain_data_size;  //!< Size of data in this chunk if plain encoding is used
+  size_type* dict_data;    //!< Dictionary data (unique row indices)
   size_type* dict_index;  //!< Index of value in dictionary page. column[dict_data[dict_index[row]]]
   uint8_t dict_rle_bits;  //!< Bit size for encoding dictionary indices
   bool use_dictionary;    //!< True if the chunk uses dictionary encoding

cpp/src/io/parquet/writer_impl.cu

@@ -1369,7 +1369,8 @@ build_chunk_dictionaries(hostdevice_2dvector<EncColumnChunk>& chunks,
       // bitpacking bitsize we efficiently support
       if (nbits > MAX_DICT_BITS) { return {false, 0}; }
 
-      auto rle_byte_size = util::div_rounding_up_safe(ck.num_values * nbits, 8);
+      auto rle_byte_size =
+        util::div_rounding_up_safe<size_t>(static_cast<size_t>(ck.num_values) * nbits, 8);
       auto dict_enc_size = ck.uniq_data_size + rle_byte_size;
       if (ck.plain_data_size <= dict_enc_size) { return {false, 0}; }
 
@@ -1912,9 +1913,10 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta,
             return l + r.num_values;
           });
         ck.plain_data_size = std::accumulate(
-          chunk_fragments.begin(), chunk_fragments.end(), 0, [](int sum, PageFragment frag) {
-            return sum + frag.fragment_data_size;
-          });
+          chunk_fragments.begin(),
+          chunk_fragments.end(),
+          size_t{0},
+          [](auto sum, PageFragment frag) { return sum + frag.fragment_data_size; });
         auto& column_chunk_meta          = row_group.columns[c].meta_data;
         column_chunk_meta.type           = parquet_columns[c].physical_type();
         column_chunk_meta.path_in_schema = parquet_columns[c].get_path_in_schema();

python/pylibcudf/tests/io/test_parquet.py

@@ -659,6 +659,31 @@ def test_write_parquet(
     assert isinstance(result, memoryview)
 
 
+def test_write_large_list_row_group():
+    import cupy as cp
+
+    # 800k rows of list<float32>[1024] exceed the signed 32-bit range. The
+    # writer must retain the correct plain-data size for dictionary selection.
+    rows = 800_000
+    embedding_dim = 1024
+    values = cp.zeros((rows, embedding_dim), dtype=cp.float32)
+    table = plc.Table(
+        [plc.Column.from_cuda_array_interface(values)],
+        num_rows=rows,
+    )
+    metadata = plc.io.types.TableInputMetadata(table)
+    sink = plc.io.SinkInfo([io.BytesIO()])
+    options = plc.io.parquet.ParquetWriterOptions.builder(sink, table).metadata(
+        metadata
+    ).build()
+
+    result = plc.io.parquet.write_parquet(options)
+    parquet_file = pq.ParquetFile(io.BytesIO(result))
+
+    assert parquet_file.metadata.num_rows == rows
+    assert parquet_file.metadata.num_row_groups == 1
+
+
 @pytest.mark.parametrize("use_jit_filter", [False, True])
 @pytest.mark.parametrize(
     "pa_filter,plc_filter",