From 1bb0ae6ce9475b1902b9c66e92c172d4862ded7e Mon Sep 17 00:00:00 2001
From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com>
Date: Thu, 21 Nov 2024 11:34:22 -0800
Subject: [PATCH] Allow `cuco::arrow_filter_policy` to accept a custom
 implementation of `xxhash_64` (#642)

---
 include/cuco/bloom_filter_policies.cuh                   | 8 ++++++--
 include/cuco/detail/bloom_filter/arrow_filter_policy.cuh | 7 ++++---
 tests/bloom_filter/unique_sequence_test.cu               | 4 ++--
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/include/cuco/bloom_filter_policies.cuh b/include/cuco/bloom_filter_policies.cuh
index cf9ddb371..a433e7cca 100644
--- a/include/cuco/bloom_filter_policies.cuh
+++ b/include/cuco/bloom_filter_policies.cuh
@@ -18,6 +18,7 @@
 
 #include <cuco/detail/bloom_filter/arrow_filter_policy.cuh>
 #include <cuco/detail/bloom_filter/default_filter_policy_impl.cuh>
+#include <cuco/hash_functions.cuh>
 
 #include <cstdint>
 
@@ -28,9 +29,12 @@ namespace cuco {
  * fingerprint.
  *
  * @tparam Key The type of the values to generate a fingerprint for.
+ * @tparam XXHash64 Custom (64 bit) XXHash hasher to generate a key's fingerprint.
+ * By default, cuco::xxhash_64 hasher will be used.
+ *
  */
-template <class Key>
-using arrow_filter_policy = detail::arrow_filter_policy<Key>;
+template <class Key, class XXHash64 = cuco::xxhash_64<Key>>
+using arrow_filter_policy = detail::arrow_filter_policy<Key, XXHash64>;
 
 /**
  * @brief The default policy that defines how a Blocked Bloom Filter generates and stores a key's
diff --git a/include/cuco/detail/bloom_filter/arrow_filter_policy.cuh b/include/cuco/detail/bloom_filter/arrow_filter_policy.cuh
index 23b95793e..e1f1b641a 100644
--- a/include/cuco/detail/bloom_filter/arrow_filter_policy.cuh
+++ b/include/cuco/detail/bloom_filter/arrow_filter_policy.cuh
@@ -78,12 +78,13 @@ namespace cuco::detail {
  * @endcode
  *
  * @tparam Key The type of the values to generate a fingerprint for.
+ * @tparam XXHash64 64-bit XXHash hasher implementation for fingerprint generation.
  */
-template <class Key>
+template <class Key, class XXHash64>
 class arrow_filter_policy {
  public:
-  using hasher    = cuco::xxhash_64<Key>;  ///< xxhash_64 hasher for Arrow bloom filter policy
-  using word_type = std::uint32_t;         ///< uint32_t for Arrow bloom filter policy
+  using hasher             = XXHash64;       ///< 64-bit XXHash hasher for Arrow bloom filter policy
+  using word_type          = std::uint32_t;  ///< uint32_t for Arrow bloom filter policy
   using hash_argument_type = typename hasher::argument_type;  ///< Hash function input type
   using hash_result_type   = decltype(std::declval<hasher>()(
     std::declval<hash_argument_type>()));  ///< hash function output type
diff --git a/tests/bloom_filter/unique_sequence_test.cu b/tests/bloom_filter/unique_sequence_test.cu
index c5970b17c..b23564270 100644
--- a/tests/bloom_filter/unique_sequence_test.cu
+++ b/tests/bloom_filter/unique_sequence_test.cu
@@ -35,8 +35,8 @@ void test_unique_sequence(Filter& filter, size_type num_keys)
 {
   using Key = typename Filter::key_type;
 
+  // Generate keys
   thrust::device_vector<Key> keys(num_keys);
-
   thrust::sequence(thrust::device, keys.begin(), keys.end());
 
   thrust::device_vector<bool> contained(num_keys, false);
@@ -119,4 +119,4 @@ TEMPLATE_TEST_CASE_SIG("bloom_filter arrow policy tests",
   auto filter = filter_type{1000};
 
   test_unique_sequence(filter, num_keys);
-}
+}
\ No newline at end of file