add batch runner

Author: tjhu

format.sh

@@ -1,11 +1,11 @@
 # Format selected folders
 
 # Format folders recursively
-for folder in examples unittests include/input_reader src/input_reader; do
+for folder in examples unittests include/input_reader src/input_reader include/hashtables/batch_runner; do
   find $folder -regex '.*\.\(cpp\|hpp\|cc\|cxx\|c\|h\)' -exec clang-format -style=file -i {} \;
 done 
 
 # Format files
-for file in include/tests/HashjoinTest.hpp src/tests/hashjoin_test.cpp include/hashtables/batch_inserter.hpp; do
+for file in include/tests/HashjoinTest.hpp src/tests/hashjoin_test.cpp; do
   clang-format -style=file -i $file
 done 

include/hashtables/batch_runner/batch_finder.hpp

@@ -0,0 +1,99 @@
+#ifndef BATCH_RUNNER_BATCH_FINDER_HPP
+#define BATCH_RUNNER_BATCH_FINDER_HPP
+
+#include <functional>
+#include <span>
+
+#include "constants.hpp"
+#include "hashtables/base_kht.hpp"
+#include "types.hpp"
+
+namespace kmercounter {
+template <size_t N = HT_TESTS_BATCH_LENGTH>
+class HTBatchFinder {
+ public:
+  using FindCallback = std::function<void(const FindResult&)>;
+
+  HTBatchFinder() : HTBatchFinder(nullptr) {}
+  HTBatchFinder(BaseHashTable* ht) : HTBatchFinder(ht, nullptr) {}
+  HTBatchFinder(BaseHashTable* ht, FindCallback callback_fn)
+      : ht_(ht),
+        buffer_size_(0),
+        results_(0, result_buffer_),
+        callback_fn_(callback_fn) {}
+  ~HTBatchFinder() { flush(); }
+
+  /// Find a key. `id` is used to track the find operation.
+  /// Set `parition_id` to the actual partition if you have more than one
+  /// partition when using PartitionedHT.
+  void find(const uint64_t key, const uint64_t id,
+            const uint64_t partition_id = 0) {
+    // Append kv to `buffer_`
+    buffer_[buffer_size_].key = key;
+    buffer_[buffer_size_].id = id;
+    buffer_[buffer_size_].part_id = partition_id;
+    buffer_size_++;
+
+    // Flush if `buffer_` is full.
+    if (buffer_size_ >= N) {
+      flush_buffer();
+    }
+  }
+
+  /// Flush everything to the hashtable and flush the hashtable find queue.
+  void flush() {
+    if (buffer_size_ > 0) {
+      flush_buffer();
+    }
+    flush_ht();
+  }
+
+  // Returns the number of elements flushed.
+  size_t num_flushed() { return num_flushed_; }
+
+  // Set the callback function.
+  void set_callback(FindCallback callback_fn) { callback_fn_ = callback_fn; }
+
+ private:
+  // Flush the insertion buffer without checking `buffer_size_`.
+  void flush_buffer() {
+    ht_->find_batch(InsertFindArguments(buffer_, buffer_size_), results_);
+    num_flushed_ += buffer_size_;
+    buffer_size_ = 0;
+    process_results();
+  }
+
+  // Issue a flush to the hashtable.
+  void flush_ht() {
+    ht_->flush_find_queue(results_);
+    process_results();
+  }
+
+  /// Process each result, if there's any.
+  void process_results() {
+    for (const auto& result : std::span(results_.second, results_.first)) {
+      callback_fn_(result);
+    }
+    results_.first = 0;
+  }
+
+  // Target hashtable.
+  BaseHashTable* ht_;
+  // Buffer to hold the arguments for batch insertion.
+  __attribute__((aligned(64))) InsertFindArgument buffer_[N];
+  // Current size of the buffer.
+  size_t buffer_size_;
+  // Total number of elements flushed.
+  size_t num_flushed_;
+  // The buffer for storing the results.
+  __attribute__((alignas(64))) FindResult result_buffer_[N];
+  // The results of finds.
+  ValuePairs results_;
+  // A user provided function for processing a result
+  FindCallback callback_fn_;
+
+  // Sanity checks
+  static_assert(N > 0);
+};
+}  // namespace kmercounter
+#endif  // BATCH_RUNNER_BATCH_FINDER_HPP

include/hashtables/batch_inserter.hpp include/hashtables/batch_runner/batch_inserter.hpp

@@ -1,8 +1,8 @@
 #ifndef HASHTABLES_BATCH_INSERTER_HPP
 #define HASHTABLES_BATCH_INSERTER_HPP
 
-#include "base_kht.hpp"
 #include "constants.hpp"
+#include "hashtables/base_kht.hpp"
 #include "types.hpp"
 
 namespace kmercounter {
@@ -22,31 +22,31 @@ class HTBatchInserter {
 
     // Flush if `buffer_` is full.
     if (buffer_size_ >= N) {
-      flush_buffer_();
+      flush_buffer();
     }
   }
 
   // Flush everything to the hashtable and flush the hashtable insert queue.
   void flush() {
     if (buffer_size_ > 0) {
-      flush_buffer_();
+      flush_buffer();
     }
-    flush_ht_();
+    flush_ht();
   }
 
   // Returns the number of elements flushed.
   size_t num_flushed() { return num_flushed_; }
 
  private:
   // Flush the insertion buffer without checking `buffer_size_`.
-  void flush_buffer_() {
+  void flush_buffer() {
     ht_->insert_batch(InsertFindArguments(buffer_, buffer_size_));
     num_flushed_ += buffer_size_;
     buffer_size_ = 0;
   }
 
   // Issue a flush to the hashtable.
-  void flush_ht_() { ht_->flush_insert_queue(); }
+  void flush_ht() { ht_->flush_insert_queue(); }
 
   // Target hashtable.
   BaseHashTable* ht_;

include/hashtables/batch_runner/batch_runner.hpp

@@ -0,0 +1,49 @@
+#ifndef BATCH_RUNNER_BATCH_RUNNER_HPP
+#define BATCH_RUNNER_BATCH_RUNNER_HPP
+
+#include "batch_finder.hpp"
+#include "batch_inserter.hpp"
+#include "hashtables/base_kht.hpp"
+
+namespace kmercounter {
+/// A wrapper around `HTBatchInserter` and `HTBatchFinder`.
+template <size_t N = HT_TESTS_BATCH_LENGTH>
+class HTBatchRunner : public HTBatchInserter<N>, public HTBatchFinder<N> {
+ public:
+  using FindCallback = HTBatchFinder<N>::FindCallback;
+
+  HTBatchRunner() = default;
+  HTBatchRunner(BaseHashTable* ht) : HTBatchRunner(ht, nullptr) {}
+  HTBatchRunner(BaseHashTable* ht, FindCallback find_callback)
+      : HTBatchInserter<N>(ht), HTBatchFinder<N>(ht, find_callback) {}
+  ~HTBatchRunner() { flush(); }
+
+  /// Insert one kv pair.
+  void insert(const uint64_t key, const uint64_t value) {
+    HTBatchInserter<N>::insert(key, value);
+  }
+
+  /// Flush both insert and find queue.
+  void flush() {
+    flush_insert();
+    flush_find();
+  }
+
+  /// Flush insert queue.
+  void flush_insert() { HTBatchInserter<N>::flush(); }
+
+  /// Flush find queue.
+  void flush_find() { HTBatchFinder<N>::flush(); }
+
+  /// Returns the number of inserts flushed.
+  size_t num_insert_flushed() { return HTBatchInserter<N>::num_flushed(); }
+
+  /// Returns the number of inserts flushed.
+  size_t num_find_flushed() { return HTBatchFinder<N>::num_flushed(); }
+
+  // Sanity checks
+  static_assert(N > 0);
+};
+}  // namespace kmercounter
+
+#endif  // BATCH_RUNNER_BATCH_RUNNER_HPP