sort_merge_join.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <cudf/column/column.hpp>
#include <cudf/column/column_view.hpp>
#include <cudf/table/table_view.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/export.hpp>
#include <cudf/utilities/memory_resource.hpp>
 
#include <rmm/cuda_stream_view.hpp>
 
#include <thrust/iterator/counting_iterator.h>
 
#include <optional>
#include <variant>
 
namespace CUDF_EXPORT cudf {
 
class sort_merge_join {
 public:
  struct match_context {
    table_view _left_table;  
    std::unique_ptr<rmm::device_uvector<size_type>>
      _match_counts;  
  };
 
  struct partition_context {
    match_context
      left_table_context;      
    size_type left_start_idx;  
    size_type
      left_end_idx;  
  };
 
  sort_merge_join(table_view const& right,
                  sorted is_right_sorted,
                  null_equality compare_nulls  = null_equality::EQUAL,
                  rmm::cuda_stream_view stream = cudf::get_default_stream());
 
  std::pair<std::unique_ptr<rmm::device_uvector<size_type>>,
            std::unique_ptr<rmm::device_uvector<size_type>>>
  inner_join(table_view const& left,
             sorted is_left_sorted,
             rmm::cuda_stream_view stream      = cudf::get_default_stream(),
             rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
  match_context inner_join_match_context(
    table_view const& left,
    sorted is_left_sorted,
    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
    rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
  std::pair<std::unique_ptr<rmm::device_uvector<size_type>>,
            std::unique_ptr<rmm::device_uvector<size_type>>>
  partitioned_inner_join(
    partition_context const& context,
    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
    rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 private:
  struct preprocessed_table {
    table_view _table_view;  
 
    table_view
      _null_processed_table_view;  
 
    std::optional<rmm::device_buffer> _validity_mask =
      std::nullopt;  
    std::optional<size_type> _num_nulls =
      std::nullopt;  
    std::optional<std::unique_ptr<table>> _null_processed_table =
      std::nullopt;  
 
    std::optional<std::unique_ptr<column>> _null_processed_table_sorted_order =
      std::nullopt;  
 
    void populate_nonnull_filter(rmm::cuda_stream_view stream);
 
    void apply_nonnull_filter(rmm::cuda_stream_view stream);
 
    void preprocess_unprocessed_table(rmm::cuda_stream_view stream);
 
    void get_sorted_order(rmm::cuda_stream_view stream);
 
    rmm::device_uvector<size_type> map_table_to_unprocessed(rmm::cuda_stream_view stream);
  };
  preprocessed_table preprocessed_left;
  preprocessed_table preprocessed_right;
  null_equality compare_nulls;
 
  void postprocess_indices(device_span<size_type> smaller_indices,
                           device_span<size_type> larger_indices,
                           rmm::cuda_stream_view stream);
 
  template <typename MergeOperation>
  auto invoke_merge(table_view right_view, table_view left_view, MergeOperation&& op);
};
 
std::pair<std::unique_ptr<rmm::device_uvector<size_type>>,
          std::unique_ptr<rmm::device_uvector<size_type>>>
sort_merge_inner_join(cudf::table_view const& left_keys,
                      cudf::table_view const& right_keys,
                      null_equality compare_nulls       = null_equality::EQUAL,
                      rmm::cuda_stream_view stream      = cudf::get_default_stream(),
                      rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
std::pair<std::unique_ptr<rmm::device_uvector<size_type>>,
          std::unique_ptr<rmm::device_uvector<size_type>>>
merge_inner_join(cudf::table_view const& left_keys,
                 cudf::table_view const& right_keys,
                 null_equality compare_nulls       = null_equality::EQUAL,
                 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