[DNR] Improving mixed join perf with sort-merge join

cpp/CMakeLists.txt

@@ -577,6 +577,11 @@ add_library(
   src/join/conditional_join.cu
   src/join/cross_join.cu
   src/join/distinct_hash_join.cu
+  src/join/filter_join_indices.cu
+  src/join/filter_join_indices_kernel_complex.cu
+  src/join/filter_join_indices_kernel_null_complex.cu
+  src/join/filter_join_indices_kernel_null_primitive.cu
+  src/join/filter_join_indices_kernel_primitive.cu
   src/join/filtered_join.cu
   src/join/hash_join.cu
   src/join/join.cu

cpp/benchmarks/CMakeLists.txt

@@ -148,7 +148,7 @@ ConfigureNVBench(
 # * join benchmark --------------------------------------------------------------------------------
 ConfigureNVBench(
   JOIN_NVBENCH join/left_join.cu join/conditional_join.cu join/join.cu join/mixed_join.cu
-  join/distinct_join.cu join/multiplicity_join.cu
+  join/distinct_join.cu join/filter_join_indices.cu join/multiplicity_join.cu
 )
 
 # ##################################################################################################

cpp/benchmarks/join/filter_join_indices.cu

@@ -0,0 +1,162 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <benchmarks/join/join_common.hpp>
+
+#include <cudf/join/hash_join.hpp>
+#include <cudf/join/join.hpp>
+
+auto const num_keys = 2;
+
+template <bool Nullable, cudf::null_equality NullEquality, data_type DataType>
+void nvbench_filter_join_indices_inner_join(nvbench::state& state,
+                                            nvbench::type_list<nvbench::enum_type<Nullable>,
+                                                               nvbench::enum_type<NullEquality>,
+                                                               nvbench::enum_type<DataType>>)
+{
+  auto join = [](cudf::table_view const& left_equality_input,
+                 cudf::table_view const& right_equality_input,
+                 cudf::table_view const& left_conditional_input,
+                 cudf::table_view const& right_conditional_input,
+                 cudf::ast::operation binary_pred,
+                 cudf::null_equality compare_nulls) {
+    auto hash_joiner                   = cudf::hash_join(right_equality_input, compare_nulls);
+    auto [left_indices, right_indices] = hash_joiner.inner_join(left_equality_input);
+
+    return cudf::filter_join_indices(left_conditional_input,
+                                     right_conditional_input,
+                                     cudf::device_span<cudf::size_type const>(*left_indices),
+                                     cudf::device_span<cudf::size_type const>(*right_indices),
+                                     binary_pred,
+                                     cudf::join_kind::INNER_JOIN);
+  };
+
+  auto dtypes = cycle_dtypes(get_type_or_group(static_cast<int32_t>(DataType)), num_keys);
+  BM_join<Nullable, join_t::MIXED, NullEquality>(state, dtypes, join);
+}
+
+template <bool Nullable, cudf::null_equality NullEquality, data_type DataType>
+void nvbench_filter_join_indices_inner_join_complex_ast(
+  nvbench::state& state,
+  nvbench::type_list<nvbench::enum_type<Nullable>,
+                     nvbench::enum_type<NullEquality>,
+                     nvbench::enum_type<DataType>>)
+{
+  auto const ast_levels = static_cast<cudf::size_type>(state.get_int64("ast_levels"));
+
+  auto join = [ast_levels](cudf::table_view const& left_equality_input,
+                           cudf::table_view const& right_equality_input,
+                           cudf::table_view const& left_conditional_input,
+                           cudf::table_view const& right_conditional_input,
+                           cudf::ast::operation binary_pred,
+                           cudf::null_equality compare_nulls) {
+    cudf::ast::tree tree;
+    create_complex_ast_expression(tree, ast_levels);
+
+    auto hash_joiner                   = cudf::hash_join(right_equality_input, compare_nulls);
+    auto [left_indices, right_indices] = hash_joiner.inner_join(left_equality_input);
+
+    return cudf::filter_join_indices(left_conditional_input,
+                                     right_conditional_input,
+                                     cudf::device_span<cudf::size_type const>(*left_indices),
+                                     cudf::device_span<cudf::size_type const>(*right_indices),
+                                     tree.back(),
+                                     cudf::join_kind::INNER_JOIN);
+  };
+
+  auto dtypes = cycle_dtypes(get_type_or_group(static_cast<int32_t>(DataType)), num_keys);
+  BM_join<Nullable, join_t::MIXED, NullEquality>(state, dtypes, join);
+}
+
+template <bool Nullable, cudf::null_equality NullEquality, data_type DataType>
+void nvbench_filter_join_indices_left_join(nvbench::state& state,
+                                           nvbench::type_list<nvbench::enum_type<Nullable>,
+                                                              nvbench::enum_type<NullEquality>,
+                                                              nvbench::enum_type<DataType>>)
+{
+  auto join = [](cudf::table_view const& left_equality_input,
+                 cudf::table_view const& right_equality_input,
+                 cudf::table_view const& left_conditional_input,
+                 cudf::table_view const& right_conditional_input,
+                 cudf::ast::operation binary_pred,
+                 cudf::null_equality compare_nulls) {
+    auto hash_joiner                   = cudf::hash_join(right_equality_input, compare_nulls);
+    auto [left_indices, right_indices] = hash_joiner.left_join(left_equality_input);
+
+    return cudf::filter_join_indices(left_conditional_input,
+                                     right_conditional_input,
+                                     cudf::device_span<cudf::size_type const>(*left_indices),
+                                     cudf::device_span<cudf::size_type const>(*right_indices),
+                                     binary_pred,
+                                     cudf::join_kind::LEFT_JOIN);
+  };
+
+  auto dtypes = cycle_dtypes(get_type_or_group(static_cast<int32_t>(DataType)), num_keys);
+  BM_join<Nullable, join_t::MIXED, NullEquality>(state, dtypes, join);
+}
+
+template <bool Nullable, cudf::null_equality NullEquality, data_type DataType>
+void nvbench_filter_join_indices_full_join(nvbench::state& state,
+                                           nvbench::type_list<nvbench::enum_type<Nullable>,
+                                                              nvbench::enum_type<NullEquality>,
+                                                              nvbench::enum_type<DataType>>)
+{
+  auto join = [](cudf::table_view const& left_equality_input,
+                 cudf::table_view const& right_equality_input,
+                 cudf::table_view const& left_conditional_input,
+                 cudf::table_view const& right_conditional_input,
+                 cudf::ast::operation binary_pred,
+                 cudf::null_equality compare_nulls) {
+    auto hash_joiner                   = cudf::hash_join(right_equality_input, compare_nulls);
+    auto [left_indices, right_indices] = hash_joiner.full_join(left_equality_input);
+
+    return cudf::filter_join_indices(left_conditional_input,
+                                     right_conditional_input,
+                                     cudf::device_span<cudf::size_type const>(*left_indices),
+                                     cudf::device_span<cudf::size_type const>(*right_indices),
+                                     binary_pred,
+                                     cudf::join_kind::FULL_JOIN);
+  };
+
+  auto dtypes = cycle_dtypes(get_type_or_group(static_cast<int32_t>(DataType)), num_keys);
+  BM_join<Nullable, join_t::MIXED, NullEquality>(state, dtypes, join);
+}
+
+NVBENCH_BENCH_TYPES(nvbench_filter_join_indices_inner_join,
+                    NVBENCH_TYPE_AXES(JOIN_NULLABLE_RANGE,
+                                      DEFAULT_JOIN_NULL_EQUALITY,
+                                      DEFAULT_JOIN_DATATYPES))
+  .set_name("filter_join_indices_inner_join")
+  .set_type_axes_names({"Nullable", "NullEquality", "DataType"})
+  .add_int64_axis("left_size", JOIN_SIZE_RANGE)
+  .add_int64_axis("right_size", JOIN_SIZE_RANGE);
+
+NVBENCH_BENCH_TYPES(nvbench_filter_join_indices_inner_join_complex_ast,
+                    NVBENCH_TYPE_AXES(JOIN_NULLABLE_RANGE,
+                                      DEFAULT_JOIN_NULL_EQUALITY,
+                                      DEFAULT_JOIN_DATATYPES))
+  .set_name("filter_join_indices_inner_join_complex_ast")
+  .set_type_axes_names({"Nullable", "NullEquality", "DataType"})
+  .add_int64_axis("left_size", JOIN_SIZE_RANGE)
+  .add_int64_axis("right_size", JOIN_SIZE_RANGE)
+  .add_int64_axis("ast_levels", {1, 5, 10});
+
+NVBENCH_BENCH_TYPES(nvbench_filter_join_indices_left_join,
+                    NVBENCH_TYPE_AXES(JOIN_NULLABLE_RANGE,
+                                      DEFAULT_JOIN_NULL_EQUALITY,
+                                      DEFAULT_JOIN_DATATYPES))
+  .set_name("filter_join_indices_left_join")
+  .set_type_axes_names({"Nullable", "NullEquality", "DataType"})
+  .add_int64_axis("left_size", JOIN_SIZE_RANGE)
+  .add_int64_axis("right_size", JOIN_SIZE_RANGE);
+
+NVBENCH_BENCH_TYPES(nvbench_filter_join_indices_full_join,
+                    NVBENCH_TYPE_AXES(JOIN_NULLABLE_RANGE,
+                                      DEFAULT_JOIN_NULL_EQUALITY,
+                                      DEFAULT_JOIN_DATATYPES))
+  .set_name("filter_join_indices_full_join")
+  .set_type_axes_names({"Nullable", "NullEquality", "DataType"})
+  .add_int64_axis("left_size", JOIN_SIZE_RANGE)
+  .add_int64_axis("right_size", JOIN_SIZE_RANGE);

cpp/benchmarks/join/join_common.hpp

@@ -13,13 +13,9 @@
 #include <benchmarks/join/nvbench_helpers.hpp>
 
 #include <cudf/ast/expressions.hpp>
-#include <cudf/column/column_factories.hpp>
-#include <cudf/filling.hpp>
-#include <cudf/scalar/scalar_factories.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/error.hpp>
-#include <cudf/utilities/memory_resource.hpp>
 
 #include <nvbench/nvbench.cuh>
 
@@ -42,6 +38,25 @@ using JOIN_NULL_EQUALITY =
 using DEFAULT_JOIN_DATATYPES     = nvbench::enum_type_list<data_type::INT32>;
 using DEFAULT_JOIN_NULL_EQUALITY = nvbench::enum_type_list<cudf::null_equality::UNEQUAL>;
 
+inline void create_complex_ast_expression(cudf::ast::tree& tree, cudf::size_type ast_levels)
+{
+  CUDF_EXPECTS(ast_levels > 0, "Number of AST levels must be greater than 0");
+
+  tree.push(cudf::ast::column_reference(0));
+  tree.push(cudf::ast::column_reference(0, cudf::ast::table_reference::RIGHT));
+
+  tree.push(cudf::ast::operation(cudf::ast::ast_operator::EQUAL, tree.at(0), tree.at(1)));
+
+  if (ast_levels == 1) { return; }
+
+  for (cudf::size_type i = 1; i < ast_levels; i++) {
+    tree.push(cudf::ast::operation(cudf::ast::ast_operator::EQUAL, tree.at(0), tree.at(1)));
+
+    tree.push(cudf::ast::operation(
+      cudf::ast::ast_operator::LOGICAL_AND, tree.at(tree.size() - 2), tree.back()));
+  }
+}
+
 template <bool Nullable,
           join_t join_type                  = join_t::HASH,
           cudf::null_equality compare_nulls = cudf::null_equality::UNEQUAL,

cpp/benchmarks/join/mixed_join.cu

@@ -9,30 +9,6 @@
 
 auto const num_keys = 2;
 
-void create_complex_ast_expression(cudf::ast::tree& tree, cudf::size_type ast_levels)
-{
-  CUDF_EXPECTS(ast_levels > 0, "Number of AST levels must be greater than 0");
-
-  // For mixed joins, the conditional tables only have 1 column each (column 0)
-  // So we'll create multiple comparisons of the same column to stress the AST evaluation
-  tree.push(cudf::ast::column_reference(0));
-  tree.push(cudf::ast::column_reference(0, cudf::ast::table_reference::RIGHT));
-
-  tree.push(cudf::ast::operation(cudf::ast::ast_operator::EQUAL, tree.at(0), tree.at(1)));
-
-  if (ast_levels == 1) { return; }
-
-  // For multiple levels, create additional comparisons of the same columns
-  // This will create expressions like: (col0_L == col0_R) && (col0_L == col0_R) && ...
-  // Total operators created: (2 * ast_levels - 1) = ast_levels EQUAL + (ast_levels-1) LOGICAL_AND
-  for (cudf::size_type i = 1; i < ast_levels; i++) {
-    tree.push(cudf::ast::operation(cudf::ast::ast_operator::EQUAL, tree.at(0), tree.at(1)));
-
-    tree.push(cudf::ast::operation(
-      cudf::ast::ast_operator::LOGICAL_AND, tree.at(tree.size() - 2), tree.back()));
-  }
-}
-
 template <bool Nullable, cudf::null_equality NullEquality, data_type DataType>
 void nvbench_mixed_inner_join(nvbench::state& state,
                               nvbench::type_list<nvbench::enum_type<Nullable>,

cpp/benchmarks/join/multiplicity_join.cu

@@ -52,20 +52,34 @@ void nvbench_hm_left_join(nvbench::state& state,
                                              nvbench::enum_type<DataType>,
                                              nvbench::enum_type<Algorithm>>)
 {
-  if constexpr (Algorithm == join_t::SORT_MERGE) {
-    state.skip("Left join using sort-merge algorithm not yet implemented");
+  if constexpr (not Nullable && NullEquality == cudf::null_equality::UNEQUAL) {
+    state.skip(
+      "Since the keys are not nullable, how null entries are to be compared by the join algorithm "
+      "is immaterial. Therefore, we skip running the benchmark when null equality is set to "
+      "UNEQUAL since the performance numbers will be the same as when null equality is set to "
+      "EQUAL.");
     return;
   }
   auto const multiplicity = static_cast<cudf::size_type>(state.get_int64("multiplicity"));
   auto const num_keys     = state.get_int64("num_keys");
   auto dtypes = cycle_dtypes(get_type_or_group(static_cast<int32_t>(DataType)), num_keys);
 
-  auto join = [](cudf::table_view const& left_input,
-                 cudf::table_view const& right_input,
-                 cudf::null_equality compare_nulls) {
+  auto hash_join = [](cudf::table_view const& left_input,
+                      cudf::table_view const& right_input,
+                      cudf::null_equality compare_nulls) {
     return cudf::left_join(left_input, right_input, compare_nulls);
   };
-  BM_join<Nullable, Algorithm, NullEquality>(state, dtypes, join, multiplicity);
+  auto sort_merge_join = [](cudf::table_view const& left_input,
+                            cudf::table_view const& right_input,
+                            cudf::null_equality compare_nulls) {
+    auto smj = cudf::sort_merge_join(right_input, cudf::sorted::NO, compare_nulls);
+    return smj.left_join(left_input, cudf::sorted::NO);
+  };
+  if constexpr (Algorithm == join_t::HASH) {
+    BM_join<Nullable, Algorithm, NullEquality>(state, dtypes, hash_join, multiplicity);
+  } else if constexpr (Algorithm == join_t::SORT_MERGE) {
+    BM_join<Nullable, Algorithm, NullEquality>(state, dtypes, sort_merge_join, multiplicity);
+  }
 }
 
 template <bool Nullable, cudf::null_equality NullEquality, data_type DataType, join_t Algorithm>