Deterministically compute parquet page fragments that first see dictionary keys

cpp/src/io/parquet/chunk_dict.cu

@@ -18,6 +18,7 @@
 #include <cooperative_groups/reduce.h>
 #include <cub/block/block_scan.cuh>
 #include <cuco/static_map_ref.cuh>
+#include <cuco/utility/reduction_functors.cuh>
 #include <cuda/atomic>
 #include <cuda/functional>
 #include <cuda/std/bit>
@@ -98,9 +99,6 @@ struct map_insert_fn {
       auto const col                     = chunk->col_desc;
       column_device_view const& data_col = *col->leaf_column;
       __shared__ size_type total_num_dict_entries;
-      __shared__ size_type num_dict_vals;
-      if (t == 0) { num_dict_vals = 0; };
-      __syncthreads();
 
       using equality_fn_type = equality_functor<T>;
       using hash_fn_type     = hash_functor<T>;
@@ -116,8 +114,8 @@ struct map_insert_fn {
                                                cuco::thread_scope_block,
                                                storage_ref};
 
-      // Create a map ref with `cuco::insert` operator
-      auto map_insert_ref = hash_map_ref.rebind_operators(cuco::insert);
+      // Create a map ref with `cuco::insert_or_apply` operator
+      auto map_insert_ref = hash_map_ref.rebind_operators(cuco::insert_or_apply);
 
       // Create atomic refs to the current chunk's num_dict_entries and uniq_data_size
       cuda::atomic_ref<size_type, SCOPE> const chunk_num_dict_entries{chunk->num_dict_entries};
@@ -136,10 +134,9 @@ struct map_insert_fn {
         // Insert fragment index to hash map using a single thread (for best performance for now)
         // and count successful insertions.
         if (is_valid) {
-          // TODO(mh): Here we insert the fragment index of the CAS winner, which may not be the
-          // smallest one (relies on monotonic block scheduling). Switch to static_map's
-          // `insert_or_apply` with `cuco::op::min` for deterministic first-fragment semantics
-          is_unique = map_insert_ref.insert(slot_type{static_cast<key_type>(val_idx), frag_idx});
+          // Insert or ensure this is the smallest fragment index inserting this key
+          is_unique = map_insert_ref.insert_or_apply(
+            slot_type{static_cast<key_type>(val_idx), frag_idx}, cuco::reduce::min{});
           uniq_elem_size = [&]() -> size_type {
             if (not is_unique) { return 0; }
             switch (col->physical_type) {
@@ -175,22 +172,18 @@ struct map_insert_fn {
         auto num_unique = block_reduce(reduce_storage).Sum(is_unique);
         __syncthreads();
         auto uniq_data_size = block_reduce(reduce_storage).Sum(uniq_elem_size);
-        // One thread atomically updates the number and data size of total unique values as well as
-        // the number of unique values inserted by this fragment.
+        // First thread atomically updates the total number and data size of unique values
         if (t == 0) {
           total_num_dict_entries =
             chunk_num_dict_entries.fetch_add(num_unique, cuda::std::memory_order_relaxed);
           total_num_dict_entries += num_unique;
-          num_dict_vals += num_unique;
           chunk_uniq_data_size.fetch_add(uniq_data_size, cuda::std::memory_order_relaxed);
         }
         __syncthreads();
 
         // Check if the num unique values in chunk has already exceeded max dict size and early exit
-        if (total_num_dict_entries > MAX_DICT_SIZE) { break; }
+        if (total_num_dict_entries > MAX_DICT_SIZE) { return; }
       }  // for loop
-      // Flush the number of unique values inserted by this fragment
-      if (t == 0) { frag->num_dict_vals = num_dict_vals; };
     } else {
       CUDF_UNREACHABLE("Unsupported type to insert in map");
     }
@@ -317,19 +310,40 @@ CUDF_KERNEL void __launch_bounds__(block_size)
     auto const& col_frags = frags[col_idx];
     auto const frag_start = static_cast<size_type>(chunk.fragments - col_frags.data());
 
-    // Initialize fragment_offsets with prefix sum (exclusive) of number of dictionary values
-    // inserted by each fragment in this chunk.
+    // fragment_offsets will contain the prefix-sum of number of dictionary values first seen in
+    // each page fragment
     __shared__ size_type fragment_offsets[MAX_FRAGMENTS_PER_CHUNK];
+
+    // Initialize all fragment offsets to 0
+    for (auto idx = t; idx < num_frags; idx += block_size) {
+      fragment_offsets[idx] = 0;
+    }
+    __syncthreads();
+
+    // Iterate over slots and count the number of dict values first seen page fragment
     {
-      using block_scan = cub::BlockScan<size_type, block_size>;
-      __shared__ typename block_scan::TempStorage scan_storage;
+      for (auto slot_idx = t; slot_idx < chunk.dict_map_size; slot_idx += block_size) {
+        auto const* slot = map_storage.data() + chunk.dict_map_offset + slot_idx;
+        if (slot->first != KEY_SENTINEL) {
+          auto const frag_loc = static_cast<size_type>(slot->second) - frag_start;
+          cudf_assert(frag_loc >= 0 && frag_loc < num_frags &&
+                      "fragment index in the slot is out of range of the chunk");
+          atomicAdd(&fragment_offsets[frag_loc], 1);
+        }
+      }
+      __syncthreads();
+    }
 
-      auto base_idx = 0;
+    // Exclusive scan to convert counts to offsets
+    using block_scan = cub::BlockScan<size_type, block_size>;
+    __shared__ typename block_scan::TempStorage scan_storage;
+
+    {
+      auto base_idx = uint32_t{0};
       while (base_idx < num_frags) {
-        auto const idx = base_idx + t;
-        auto const per_thread_count =
-          (idx < num_frags) ? col_frags[frag_start + idx].num_dict_vals : 0;
-        auto per_thread_offset = 0;
+        auto const idx              = base_idx + t;
+        auto const per_thread_count = (idx < num_frags) ? fragment_offsets[t] : 0;
+        auto per_thread_offset      = 0;
         block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset);
         if (idx < num_frags) { fragment_offsets[idx] = per_thread_offset; }
         base_idx += block_size;

cpp/src/io/parquet/page_enc.cu

@@ -153,7 +153,6 @@ void __device__ init_frag_state(frag_init_state_s* const s,
   // smaller. num_rows is fixed but fragment size could be larger if the data is strings or
   // nested.
   s->frag.num_rows = cuda::std::min<size_type>(fragment_size, part_end_row - s->frag.start_row);
-  s->frag.num_dict_vals      = 0;
   s->frag.fragment_data_size = 0;
   s->frag.dict_data_size     = 0;
 

cpp/src/io/parquet/parquet_gpu.hpp

@@ -531,7 +531,6 @@ struct PageFragment {
   uint32_t num_valid;        //<! Number of non-null leaf values
   size_type start_row;       //!< First row in fragment
   size_type num_rows;        //!< Number of rows in fragment
-  size_type num_dict_vals;   //!< Number of unique dictionary entries
   EncColumnChunk* chunk;     //!< The chunk that this fragment belongs to
 };
 

cpp/tests/io/parquet_writer_test.cpp

@@ -1197,14 +1197,11 @@ TEST_F(ParquetWriterTest, VariableBitWidthDictEncoding)
     EXPECT_LE(*max_bits_iter, chunk_wide_max_bits);
 
     // Check expected number of freq and rare pages
-    // TODO(mh): Race-dependent checks. Enable these along with cuDF PR #22323.
-
-    // auto const total_page_count = static_cast<int>(page_dict_bits.size());
-    // auto const freq_page_count  = static_cast<int>(
-    //   std::ranges::count_if(page_dict_bits, [&](int nbits) { return nbits <= frequent_max_bits;
-    //   }));
-    // EXPECT_EQ(freq_page_count, freq_pages);
-    // EXPECT_EQ(total_page_count - freq_page_count, num_pages - freq_pages);
+    auto const total_page_count = static_cast<int>(page_dict_bits.size());
+    auto const freq_page_count  = static_cast<int>(
+      std::ranges::count_if(page_dict_bits, [&](int nbits) { return nbits <= frequent_max_bits; }));
+    EXPECT_EQ(freq_page_count, freq_pages);
+    EXPECT_EQ(total_page_count - freq_page_count, num_pages - freq_pages);
   }
 }