Add a direct_inner_join API for pre-hashed distinct UINT32 keys (rapidsai#23147)

PointKernel · abigalekim · commit bb3346ed5828 · 2026-07-20T21:28:35.000Z
Closes rapidsai#23146 This PR adds a `direct_inner_join` free function to libcudf, the first step of the perfect hash join effort in rapidsai#23126. The keys are a single `UINT32` column per side, produced by a prior perfect hashing pass such as `cudf::key_remapping`, dictionary encoding, or dense integer primary keys. The right keys act as a perfect hash of the right rows: a lookup table of caller-specified `capacity` entries maps each key value to its row index and each left key probes that table directly, so the join performs no hashing or key comparison at all. - The caller controls the memory footprint via the explicit `capacity` argument. All key values must be in `[0, capacity)` and the right keys must be distinct; behavior is undefined otherwise. - Inner join only, as a free function: no table reuse across probes is needed, so there is no join object. - The build scatters right row indices into the lookup table with `cub::DeviceTransform::Fill` + `cub::DeviceFor::Bulk`; the probe is a single `cub`-based `copy_if` pass emitting the matched index pairs. - A new `JOIN_NVBENCH` benchmark compares `inner_join`, `distinct_hash_join`, and `direct_inner_join` on identical conforming input; results in the comment below. Authors: - Yunsong Wang (https://github.com/PointKernel) Approvers: - Tianyu Liu (https://github.com/kingcrimsontianyu) - Shruti Shivakumar (https://github.com/shrshi) - Muhammad Haseeb (https://github.com/mhaseeb123) URL: rapidsai#23147
diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
@@ -841,6 +841,7 @@ add_library(
   src/jit/util.cpp
   src/join/conditional_join.cu
   src/join/cross_join.cu
+  src/join/direct_join.cu
   src/join/distinct_hash_join.cu
   src/join/filter_join_indices/filter_join_indices.cu
   src/join/filter_join_indices/filter_join_indices_jit.cu
diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt
@@ -150,6 +150,7 @@ ConfigureNVBench(
 ConfigureNVBench(
   JOIN_NVBENCH
   join/conditional_join.cpp
+  join/direct_join.cu
   join/distinct_join.cpp
   join/filter_join_indices.cpp
   join/filter_join_indices_jit.cu
diff --git a/cpp/benchmarks/join/direct_join.cu b/cpp/benchmarks/join/direct_join.cu
@@ -0,0 +1,90 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "join_common.hpp"
+
+#include <cudf/column/column_factories.hpp>
+#include <cudf/join/direct_join.hpp>
+#include <cudf/join/distinct_hash_join.hpp>
+#include <cudf/join/join.hpp>
+
+#include <thrust/execution_policy.h>
+#include <thrust/functional.h>
+#include <thrust/random.h>
+#include <thrust/sequence.h>
+#include <thrust/shuffle.h>
+#include <thrust/tabulate.h>
+
+// Apples-to-apples comparison of inner join implementations on input that satisfies
+// `direct_inner_join`'s preconditions: a single UINT32 key column per side, distinct right keys,
+// and all key values in [0, capacity) with capacity = right_size. The right keys are the shuffled
+// dense values [0, right_size), the key_remapping/dense-primary-key case, so every left key
+// matches and the input is identical for all algorithms.
+void nvbench_direct_inner_join(nvbench::state& state)
+{
+  if (should_skip_large_sizes(state)) { return; }
+
+  auto const right_size = static_cast<cudf::size_type>(state.get_int64("right_size"));
+  auto const left_size  = static_cast<cudf::size_type>(state.get_int64("left_size"));
+  auto const algorithm  = state.get_string("algorithm");
+  auto const capacity   = static_cast<std::size_t>(right_size);
+
+  // Dense distinct right keys: a shuffled sequence of [0, capacity)
+  auto right = cudf::make_numeric_column(
+    cudf::data_type{cudf::type_id::UINT32}, right_size, cudf::mask_state::UNALLOCATED);
+  thrust::sequence(thrust::device,
+                   right->mutable_view().begin<std::uint32_t>(),
+                   right->mutable_view().end<std::uint32_t>());
+  thrust::shuffle(thrust::device,
+                  right->mutable_view().begin<std::uint32_t>(),
+                  right->mutable_view().end<std::uint32_t>(),
+                  thrust::default_random_engine{12345});
+
+  // Left keys cycle through [0, capacity), then shuffled
+  auto left = cudf::make_numeric_column(
+    cudf::data_type{cudf::type_id::UINT32}, left_size, cudf::mask_state::UNALLOCATED);
+  thrust::tabulate(thrust::device,
+                   left->mutable_view().begin<std::uint32_t>(),
+                   left->mutable_view().end<std::uint32_t>(),
+                   thrust::placeholders::_1 % static_cast<std::uint32_t>(right_size));
+  thrust::shuffle(thrust::device,
+                  left->mutable_view().begin<std::uint32_t>(),
+                  left->mutable_view().end<std::uint32_t>(),
+                  thrust::default_random_engine{67890});
+
+  auto const left_view  = left->view();
+  auto const right_view = right->view();
+  auto const left_keys  = cudf::table_view{{left_view}};
+  auto const right_keys = cudf::table_view{{right_view}};
+
+  auto const input_bytes = estimate_size(left_keys) + estimate_size(right_keys);
+  state.set_cuda_stream(nvbench::make_cuda_stream_view(cudf::get_default_stream().value()));
+  state.add_element_count(input_bytes, "input_bytes");
+  state.add_global_memory_reads<nvbench::int8_t>(input_bytes);
+
+  if (algorithm == "hash") {
+    state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) {
+      auto result = cudf::inner_join(left_keys, right_keys, cudf::null_equality::UNEQUAL);
+    });
+  } else if (algorithm == "distinct_hash") {
+    state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) {
+      auto hj_obj = cudf::distinct_hash_join{right_keys, cudf::null_equality::UNEQUAL, 0.5};
+      auto result = hj_obj.inner_join(left_keys);
+    });
+  } else if (algorithm == "direct") {
+    state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) {
+      auto result = cudf::direct_inner_join(left_view, right_view, capacity);
+    });
+  } else {
+    state.skip("unknown algorithm");
+  }
+}
+
+NVBENCH_BENCH(nvbench_direct_inner_join)
+  .set_name("direct_inner_join")
+  .add_string_axis("algorithm", {"hash", "distinct_hash", "direct"})
+  .add_int64_axis("left_size", JOIN_SIZE_RANGE)
+  .add_int64_axis("right_size", JOIN_SIZE_RANGE)
+  .add_int64_axis("skip_large_sizes", {1});
diff --git a/cpp/include/cudf/join/direct_join.hpp b/cpp/include/cudf/join/direct_join.hpp
@@ -0,0 +1,66 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cudf/column/column_view.hpp>
+#include <cudf/types.hpp>
+#include <cudf/utilities/default_stream.hpp>
+#include <cudf/utilities/export.hpp>
+#include <cudf/utilities/memory_resource.hpp>
+
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/device_uvector.hpp>
+
+#include <memory>
+#include <utility>
+
+namespace CUDF_EXPORT cudf {
+
+/**
+ * @addtogroup column_join
+ * @{
+ * @file
+ * @brief Direct join APIs for pre-hashed integer keys
+ */
+
+/**
+ * @brief Returns the row indices that can be used to construct the result of performing an inner
+ * join between two key columns whose values directly determine the matched row index
+ *
+ * The right keys are treated as a perfect hash of the right rows: a lookup table of `capacity`
+ * entries maps each key value to its row index, and each left key probes that table directly. No
+ * hashing or key comparison is performed. Left keys that do not occur in the right keys produce no
+ * output pair.
+ *
+ * @note Behavior is undefined if any key value is not less than `capacity`, or if the right keys
+ * contain duplicates.
+ *
+ * @throw cudf::data_type_error if the key columns are not of type UINT32
+ * @throw std::invalid_argument if the key columns contain nulls
+ * @throw std::invalid_argument if `capacity` is less than the number of right keys
+ *
+ * @param left_keys The left key column containing pre-hashed keys in `[0, capacity)`, from which
+ * the keys are probed
+ * @param right_keys The right key column containing distinct pre-hashed keys in `[0, capacity)`
+ * @param capacity The number of entries in the lookup table
+ * @param stream CUDA stream used for device memory operations and kernel launches
+ * @param mr Device memory resource used to allocate the returned indices' device memory
+ *
+ * @return A pair of vectors [`left_indices`, `right_indices`] that can be used to construct the
+ * result of performing an inner join between two tables with `left_keys` and `right_keys` as the
+ * join keys
+ */
+[[nodiscard]] std::pair<std::unique_ptr<rmm::device_uvector<size_type>>,
+                        std::unique_ptr<rmm::device_uvector<size_type>>>
+direct_inner_join(column_view const& left_keys,
+                  column_view const& right_keys,
+                  std::size_t capacity,
+                  rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+                  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
+
+/** @} */  // end of group
+
+}  // namespace CUDF_EXPORT cudf
diff --git a/cpp/src/join/direct_join.cu b/cpp/src/join/direct_join.cu
@@ -0,0 +1,140 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <cudf/column/column_view.hpp>
+#include <cudf/detail/algorithms/copy_if.cuh>
+#include <cudf/detail/nvtx/ranges.hpp>
+#include <cudf/join/direct_join.hpp>
+#include <cudf/join/join.hpp>
+#include <cudf/types.hpp>
+#include <cudf/utilities/error.hpp>
+#include <cudf/utilities/memory_resource.hpp>
+
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/device_uvector.hpp>
+
+#include <cub/device/device_for.cuh>
+#include <cub/device/device_transform.cuh>
+#include <cuda/iterator>
+#include <cuda/std/iterator>
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+
+namespace cudf {
+namespace detail {
+namespace {
+
+// Scatters each right row index to the lookup slot addressed by its key value
+struct scatter_right_index {
+  size_type* lookup;
+  std::uint32_t const* right_keys;
+
+  __device__ void operator()(size_type right_idx) const
+  {
+    lookup[right_keys[right_idx]] = right_idx;
+  }
+};
+
+// Writes the (left, right) index pair of the `out_idx`-th match, given a matched left row index
+struct emit_match_pair {
+  size_type* left_out;
+  size_type* right_out;
+  size_type const* lookup;
+  std::uint32_t const* left_keys;
+
+  __device__ void operator()(size_type out_idx, size_type left_idx) const
+  {
+    left_out[out_idx]  = left_idx;
+    right_out[out_idx] = lookup[left_keys[left_idx]];
+  }
+};
+
+// Returns true if the left row's key hits a right row in the lookup table
+struct is_match {
+  size_type const* lookup;
+  std::uint32_t const* left_keys;
+
+  __device__ bool operator()(size_type left_idx) const
+  {
+    return lookup[left_keys[left_idx]] != JoinNoMatch;
+  }
+};
+
+}  // namespace
+
+std::pair<std::unique_ptr<rmm::device_uvector<size_type>>,
+          std::unique_ptr<rmm::device_uvector<size_type>>>
+direct_inner_join(column_view const& left_keys,
+                  column_view const& right_keys,
+                  std::size_t capacity,
+                  rmm::cuda_stream_view stream,
+                  rmm::device_async_resource_ref mr)
+{
+  CUDF_EXPECTS(
+    left_keys.type().id() == type_id::UINT32 and right_keys.type().id() == type_id::UINT32,
+    "direct_inner_join keys must be of type UINT32",
+    cudf::data_type_error);
+  CUDF_EXPECTS(not left_keys.has_nulls() and not right_keys.has_nulls(),
+               "direct_inner_join keys must not contain nulls",
+               std::invalid_argument);
+  CUDF_EXPECTS(static_cast<std::size_t>(right_keys.size()) <= capacity,
+               "capacity must be at least the number of right keys",
+               std::invalid_argument);
+
+  if (left_keys.is_empty() or right_keys.is_empty()) {
+    return std::pair(std::make_unique<rmm::device_uvector<size_type>>(0, stream, mr),
+                     std::make_unique<rmm::device_uvector<size_type>>(0, stream, mr));
+  }
+
+  // Build: scatter each right row index to the slot addressed by its key value
+  auto lookup =
+    rmm::device_uvector<size_type>(capacity, stream, cudf::get_current_device_resource_ref());
+  CUDF_CUDA_TRY(
+    cub::DeviceTransform::Fill(lookup.begin(), lookup.size(), JoinNoMatch, stream.value()));
+  CUDF_CUDA_TRY(
+    cub::DeviceFor::Bulk(right_keys.size(),
+                         scatter_right_index{lookup.data(), right_keys.begin<std::uint32_t>()},
+                         stream.value()));
+
+  // Probe: a single pass emitting the (left index, matched right index) pairs
+  auto left_indices =
+    std::make_unique<rmm::device_uvector<size_type>>(left_keys.size(), stream, mr);
+  auto right_indices =
+    std::make_unique<rmm::device_uvector<size_type>>(left_keys.size(), stream, mr);
+
+  auto const d_left_keys = left_keys.begin<std::uint32_t>();
+  auto const out_iter    = cuda::tabulate_output_iterator{
+    emit_match_pair{left_indices->data(), right_indices->data(), lookup.data(), d_left_keys}};
+
+  auto const out_end = cudf::detail::copy_if(cuda::counting_iterator<size_type>{0},
+                                             cuda::counting_iterator<size_type>{left_keys.size()},
+                                             out_iter,
+                                             is_match{lookup.data(), d_left_keys},
+                                             stream);
+
+  auto const num_matches = cuda::std::distance(out_iter, out_end);
+  left_indices->resize(num_matches, stream);
+  right_indices->resize(num_matches, stream);
+
+  return std::pair(std::move(left_indices), std::move(right_indices));
+}
+
+}  // namespace detail
+
+std::pair<std::unique_ptr<rmm::device_uvector<size_type>>,
+          std::unique_ptr<rmm::device_uvector<size_type>>>
+direct_inner_join(column_view const& left_keys,
+                  column_view const& right_keys,
+                  std::size_t capacity,
+                  rmm::cuda_stream_view stream,
+                  rmm::device_async_resource_ref mr)
+{
+  CUDF_FUNC_RANGE();
+  return detail::direct_inner_join(left_keys, right_keys, capacity, stream, mr);
+}
+
+}  // namespace cudf
diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
@@ -176,6 +176,7 @@ ConfigureTest(
   join/cross_join_tests.cpp
   join/semi_anti_join_tests.cpp
   join/mixed_join_tests.cu
+  join/direct_join_tests.cpp
   join/distinct_join_tests.cpp
   join/key_remapping_tests.cpp
   GPUS 1
diff --git a/cpp/tests/join/direct_join_tests.cpp b/cpp/tests/join/direct_join_tests.cpp
diff --git a/cpp/tests/streams/join_test.cpp b/cpp/tests/streams/join_test.cpp
