Support concurrent aggregate calls on one streaming_groupby

cpp/include/cudf/groupby.hpp

@@ -519,6 +519,13 @@ class streaming_groupby {
    * are updated atomically. The input `data` table is not referenced after this
    * call returns.
    *
+   * This function may be called concurrently from multiple host threads on the same object,
+   * and each call may supply a different stream. Callers do not need to serialize the calls or
+   * synchronize between them. Key insertion is serialized internally, on the host and across
+   * streams, because a batch's newly discovered keys are held in a transient encoding that is
+   * only valid while that one insertion is in flight. The aggregation that follows each
+   * insertion updates every group atomically, so those phases overlap freely across streams.
+   *
    * @param data Table containing both key and value columns
    * @param stream CUDA stream used for device memory operations and kernel launches
    *
@@ -532,6 +539,10 @@ class streaming_groupby {
    *
    * Extracts the other object's accumulated intermediate state and merges it into this
    * object's persistent hash table. The other object is not modified.
+   *
+   * This function shares the insertion path with `aggregate()` and is serialized against it, so
+   * it is safe to call while other host threads are calling `aggregate()` on this object. The
+   * source object must not be mutated concurrently.
    * Both objects must have been constructed with compatible aggregation requests,
    * and this object must have had at least one `aggregate()` call.
    *

cpp/src/groupby/streaming_groupby/aggregate.cu

@@ -17,15 +17,14 @@
 #include <thrust/for_each.h>
 
 #include <limits>
+#include <mutex>
 #include <string>
 
 namespace cudf::groupby {
 
 void streaming_groupby::impl::do_aggregate(table_view const& data, cuda::stream_ref stream)
 {
-  CUDF_EXPECTS(!_invalidated,
-               "streaming_groupby is in an invalidated state from a prior failure; "
-               "no further aggregate()/merge() is allowed.  finalize() may still be called.");
+  ensure_not_invalidated();
 
   auto const batch_size = data.num_rows();
   if (batch_size == 0) { return; }
@@ -40,15 +39,29 @@ void streaming_groupby::impl::do_aggregate(table_view const& data, cuda::stream_
                "Transient key encoding (max_distinct_keys + batch_size) would overflow size_type.",
                std::invalid_argument);
 
-  if (!_initialized) { initialize(data, stream); }
+  // The transient key encoding is only valid while a single insertion is in flight, so
+  // insertion is serialized across concurrent callers on the host and, via the event, on the
+  // device.  The aggregation below is per-group atomic and runs unserialized.
+  auto const result = [&] {
+    std::lock_guard const lock{_insert_mutex};
 
-  auto const batch_keys = data.select(_key_indices);
+    // Re-check under the lock: another caller may have invalidated the object since the
+    // fail-fast check above.
+    ensure_not_invalidated();
 
-  update_nullable_state(batch_keys);
+    if (!_initialized) { initialize(data, stream); }
 
-  if (!_key_set) { create_key_set(stream); }
+    auto const batch_keys = data.select(_key_indices);
 
-  auto result = probe_and_insert(batch_keys, stream);
+    update_nullable_state(batch_keys);
+
+    if (!_key_set) { create_key_set(stream); }
+
+    _insert_done.wait(stream);
+    auto inserted = probe_and_insert(batch_keys, stream);
+    _insert_done.record(stream);
+    return inserted;
+  }();
 
   auto const values_view = data.select(_value_col_indices);
   auto const d_values    = table_device_view::create(values_view, stream);

cpp/src/groupby/streaming_groupby/common.cuh

@@ -26,11 +26,40 @@
 #include <cuda/std/utility>
 #include <cuda/stream>
 
+#include <atomic>
 #include <memory>
+#include <mutex>
 #include <vector>
 
 namespace cudf::groupby {
 
+/*
+ * Minimal owning wrapper around a CUDA event, used to order the insertion phase of
+ * `aggregate()` / `merge()` calls that overlap on different streams.  Only the insertion
+ * phase needs this ordering; the aggregation phase updates each group with atomics and is
+ * safe to overlap.
+ */
+class insert_order_event {
+ public:
+  insert_order_event() { CUDF_CUDA_TRY(cudaEventCreateWithFlags(&_event, cudaEventDisableTiming)); }
+  ~insert_order_event() { cudaEventDestroy(_event); }
+  insert_order_event(insert_order_event const&)            = delete;
+  insert_order_event& operator=(insert_order_event const&) = delete;
+
+  /// Makes `stream` wait for the most recently recorded insertion.  No-op before the first
+  /// `record()`, which is exactly the behavior the first call needs.
+  void wait(cuda::stream_ref stream) const
+  {
+    CUDF_CUDA_TRY(cudaStreamWaitEvent(stream.get(), _event));
+  }
+
+  /// Records completion of the insertion just enqueued on `stream`.
+  void record(cuda::stream_ref stream) { CUDF_CUDA_TRY(cudaEventRecord(_event, stream.get())); }
+
+ private:
+  cudaEvent_t _event{};
+};
+
 /*
  * Companion location for a stored dense ID: which compacted batch table the key
  * lives in (`first`) and the row index within that table (`second`).  Packed into
@@ -250,17 +279,27 @@ struct streaming_groupby::impl {
   null_policy _null_handling;
   cuda::mr::any_resource<cuda::mr::device_accessible> _mr;
 
+  /*
+   * Serializes the insertion phase of `aggregate()` and `merge()`.  Callers may invoke those
+   * from multiple host threads; everything they mutate on the host, and the transient key
+   * encoding they place in the hash set, is guarded here.
+   */
+  std::mutex _insert_mutex;
+  /// Orders the insertion phase across calls that supply different streams.
+  insert_order_event _insert_done;
+
   bool _initialized{false};
   /// Set true once an `aggregate()` / `merge()` call has thrown after touching the
   /// hash set.  Subsequent `aggregate()` / `merge()` calls fail fast; only
-  /// `finalize()` may still be called to recover partial results.
-  bool _invalidated{false};
+  /// `finalize()` may still be called to recover partial results.  Atomic so the
+  /// fail-fast check in `do_aggregate` can run ahead of `_insert_mutex`.
+  std::atomic<bool> _invalidated{false};
   /*
    * Number of distinct keys accumulated so far.  Also serves as the high-water
    * mark of dense IDs in the persistent hash set: stored slot values are in
    * [0, _distinct_keys).
    */
-  size_type _distinct_keys{0};
+  std::atomic<size_type> _distinct_keys{0};
   bool _has_nullable_keys{false};
   bool _has_nested_keys{false};
 
@@ -301,7 +340,12 @@ struct streaming_groupby::impl {
   std::unique_ptr<streaming_set_t> _key_set;
 
   [[nodiscard]] size_type num_keys() const { return static_cast<size_type>(_key_indices.size()); }
-  [[nodiscard]] bool has_state() const { return _initialized && _distinct_keys > 0; }
+  void ensure_not_invalidated() const
+  {
+    CUDF_EXPECTS(!_invalidated.load(std::memory_order_relaxed),
+                 "streaming_groupby is in an invalidated state from a prior failure; "
+                 "no further aggregate()/merge() is allowed.  finalize() may still be called.");
+  }
 
   impl(host_span<size_type const> key_indices,
        host_span<streaming_aggregation_request const> requests,

cpp/src/groupby/streaming_groupby/impl.cu

@@ -243,7 +243,9 @@ std::unique_ptr<table> streaming_groupby::impl::gather_agg_results(
   // The results we care about are dense in `[0, _distinct_keys)` and can be extracted by
   // slice+copy.
   auto const sliced =
-    cudf::detail::slice(_agg_results->view(), {0, _distinct_keys}, stream).front();
+    cudf::detail::slice(
+      _agg_results->view(), {0, _distinct_keys.load(std::memory_order_relaxed)}, stream)
+      .front();
   return std::make_unique<table>(sliced, stream, mr);
 }
 
@@ -378,7 +380,10 @@ std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> streaming_gro
   return _impl->do_finalize(stream, mr);
 }
 
-size_type streaming_groupby::distinct_keys() const noexcept { return _impl->_distinct_keys; }
+size_type streaming_groupby::distinct_keys() const noexcept
+{
+  return _impl->_distinct_keys.load(std::memory_order_relaxed);
+}
 
 bool is_streaming_groupby_supported(data_type values_type, aggregation::Kind kind)
 {

cpp/src/groupby/streaming_groupby/insert.cuh

@@ -96,9 +96,10 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_
     // Bound check: the hash set has already been written above (transient slot values),
     // so on failure the object is left invalidated; further aggregate()/merge() calls
     // will throw immediately while finalize() can still recover partial results.
-    if (_distinct_keys + new_distinct_keys > _max_distinct_keys) {
+    auto const distinct_so_far = _distinct_keys.load(std::memory_order_relaxed);
+    if (distinct_so_far + new_distinct_keys > _max_distinct_keys) {
       _invalidated = true;
-      CUDF_FAIL("Distinct key count (" + std::to_string(_distinct_keys + new_distinct_keys) +
+      CUDF_FAIL("Distinct key count (" + std::to_string(distinct_so_far + new_distinct_keys) +
                 ") would exceed max_distinct_keys (" + std::to_string(_max_distinct_keys) + ").");
     }
 
@@ -115,7 +116,7 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_
 
     // Store the compacted batch.
     auto const new_batch_id    = static_cast<size_type>(_compacted_batches.size());
-    auto const dense_id_offset = _distinct_keys;
+    auto const dense_id_offset = distinct_so_far;
     _compacted_batches.push_back(std::move(compacted));
     _preprocessed_batches.push_back(preprocessed_compacted);
 
@@ -138,7 +139,7 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_
                        update_transient_target_indices_fn{
                          base, slot_offsets.data(), _max_distinct_keys, target_indices.data()});
 
-    _distinct_keys += new_distinct_keys;
+    _distinct_keys.fetch_add(new_distinct_keys, std::memory_order_relaxed);
   }
   // If new_distinct_keys == 0, target_indices is already final from Pass 1 — every
   // slot held a dense ID at probe time, so *iter was already the correct dense ID.

cpp/src/groupby/streaming_groupby/merge.cu

@@ -18,6 +18,7 @@
 #include <cuda/stream>
 #include <thrust/for_each.h>
 
+#include <mutex>
 #include <string>
 
 namespace cudf::groupby {
@@ -90,17 +91,20 @@ struct merge_single_pass_aggs_fn {
 
 void streaming_groupby::impl::do_merge(impl const& other, cuda::stream_ref stream)
 {
-  CUDF_EXPECTS(!_invalidated,
-               "streaming_groupby is in an invalidated state from a prior failure; "
-               "no further aggregate()/merge() is allowed.  finalize() may still be called.");
-  CUDF_EXPECTS(!other._invalidated, "Cannot merge from an invalidated streaming_groupby.");
+  // `other` is only read from, so a single lock on this object's insertion state is enough.
+  std::lock_guard const lock{_insert_mutex};
 
-  if (!other._initialized || !other.has_state()) { return; }
+  ensure_not_invalidated();
+  CUDF_EXPECTS(!other._invalidated.load(std::memory_order_relaxed),
+               "Cannot merge from an invalidated streaming_groupby.");
+
+  auto const other_distinct_keys = other._distinct_keys.load(std::memory_order_relaxed);
+  if (!other._initialized || other_distinct_keys == 0) { return; }
   CUDF_EXPECTS(_initialized,
                "Cannot merge into an uninitialized streaming_groupby. "
                "Call aggregate() at least once before merge().");
-  CUDF_EXPECTS(other._distinct_keys <= _max_distinct_keys,
-               "Merge source distinct keys (" + std::to_string(other._distinct_keys) +
+  CUDF_EXPECTS(other_distinct_keys <= _max_distinct_keys,
+               "Merge source distinct keys (" + std::to_string(other_distinct_keys) +
                  ") exceeds max_distinct_keys (" + std::to_string(_max_distinct_keys) + ").",
                std::invalid_argument);
   CUDF_EXPECTS(other._agg_kinds == _agg_kinds,
@@ -115,16 +119,16 @@ void streaming_groupby::impl::do_merge(impl const& other, cuda::stream_ref strea
 
   auto const mr = cudf::get_current_device_resource_ref();
 
-  auto other_keys                = other.gather_distinct_keys(stream, mr);
-  auto const other_key_view      = other_keys->view();
-  auto const other_distinct_keys = other._distinct_keys;
-  if (other_distinct_keys == 0) { return; }
+  auto other_keys           = other.gather_distinct_keys(stream, mr);
+  auto const other_key_view = other_keys->view();
 
   update_nullable_state(other_key_view);
 
   if (!_key_set) { create_key_set(stream); }
 
+  _insert_done.wait(stream);
   auto result = probe_and_insert(other_key_view, stream);
+  _insert_done.record(stream);
 
   // Merge aggregation values using dense target indices.  We only read from
   // `other._agg_results`; no need to deep-copy the source rows like keys.

cpp/tests/groupby/streaming_groupby_test.cpp

@@ -19,8 +19,12 @@
 #include <cudf/unary.hpp>
 #include <cudf/utilities/traits.hpp>
 
+#include <rmm/cuda_device.hpp>
+#include <rmm/cuda_stream.hpp>
 #include <rmm/mr/statistics_resource_adaptor.hpp>
 
+#include <atomic>
+#include <thread>
 #include <vector>
 
 static std::vector<cudf::size_type> const KEY_COL{0};
@@ -325,6 +329,79 @@ TEST_F(StreamingGroupbyTest, MergeTwoObjects)
   check(keys, results, cudf::table_view{{ek}}, {ev});
 }
 
+TEST_F(StreamingGroupbyTest, ConcurrentAggregate)
+{
+  using K = int32_t;
+  using V = int32_t;
+
+  constexpr int num_batches = 8;
+
+  // Every batch re-hits keys 0 and 1 and introduces one key of its own, so concurrent calls
+  // both collide on existing groups and discover new keys at the same time.
+  std::vector<cudf::test::fixed_width_column_wrapper<K>> keys;
+  std::vector<cudf::test::fixed_width_column_wrapper<V>> vals;
+  keys.reserve(num_batches);
+  vals.reserve(num_batches);
+  for (int i = 0; i < num_batches; ++i) {
+    keys.emplace_back(std::initializer_list<K>{0, 1, static_cast<K>(i + 2)});
+    vals.emplace_back(std::initializer_list<V>{1, 10, 100});
+  }
+
+  std::vector<cudf::table_view> batches;
+  batches.reserve(num_batches);
+  for (int i = 0; i < num_batches; ++i) {
+    batches.push_back(cudf::table_view{{keys[i], vals[i]}});
+  }
+
+  std::vector<std::unique_ptr<rmm::cuda_stream>> streams;
+  streams.reserve(num_batches);
+  for (int i = 0; i < num_batches; ++i) {
+    streams.push_back(std::make_unique<rmm::cuda_stream>());
+  }
+
+  auto reqs = single_agg_req(1, cudf::make_sum_aggregation<cudf::groupby_aggregation>());
+  cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS);
+
+  auto const device = rmm::get_current_cuda_device();
+  std::vector<std::thread> threads;
+  std::vector<std::exception_ptr> errors(num_batches);
+  // `ready` lets the main thread wait until every worker is spinning, and `start` then releases
+  // them together, so the aggregate() calls actually overlap.
+  std::atomic<int> ready{0};
+  std::atomic<bool> start{false};
+  threads.reserve(num_batches);
+  for (int i = 0; i < num_batches; ++i) {
+    threads.emplace_back([&, i] {
+      rmm::cuda_set_device_raii const device_guard{device};
+      ready.fetch_add(1, std::memory_order_relaxed);
+      while (!start.load(std::memory_order_acquire)) {
+        std::this_thread::yield();
+      }
+      try {
+        streaming_agg.aggregate(batches[i], streams[i]->view());
+      } catch (...) {
+        errors[i] = std::current_exception();
+      }
+    });
+  }
+  while (ready.load(std::memory_order_relaxed) != num_batches) {
+    std::this_thread::yield();
+  }
+  start.store(true, std::memory_order_release);
+  for (auto& thread : threads) {
+    thread.join();
+  }
+  for (auto const& error : errors) {
+    EXPECT_FALSE(error);
+  }
+  for (auto const& stream : streams) {
+    stream->synchronize();
+  }
+
+  auto [out_keys, results] = streaming_agg.finalize();
+  verify_against_groupby(out_keys, results, batches, KEY_COL, reqs);
+}
+
 TEST_F(StreamingGroupbyTest, EmptyBatch)
 {
   using K = int32_t;
@@ -849,6 +926,14 @@ TEST_F(StreamingGroupbyTest, ExceedsDistinctKeyCapacityThrows)
   cudf::test::fixed_width_column_wrapper<K> k2{4};
   cudf::test::fixed_width_column_wrapper<V> v2{50};
   EXPECT_THROW(streaming_agg.aggregate(cudf::table_view{{k2, v2}}), cudf::logic_error);
+
+  // The object is now invalidated: even an empty batch is rejected, while finalize() still
+  // recovers the groups inserted before the failure.
+  cudf::test::fixed_width_column_wrapper<K> k_empty{};
+  cudf::test::fixed_width_column_wrapper<V> v_empty{};
+  EXPECT_THROW(streaming_agg.aggregate(cudf::table_view{{k_empty, v_empty}}), cudf::logic_error);
+  auto [keys, results] = streaming_agg.finalize();
+  EXPECT_EQ(keys->num_rows(), 4);
 }
 
 // Test that sliced input columns with non-zero offsets work correctly.