Trying to perform retrieve with input more than just the probe keys.

[kevkrist]

For database workloads, when probing, it is often necessary to carry along with the probe, for example, the rowids of the tuples that pass the probe, so that the attributes for tuples that are emitted by the join can then be later gathered. I'm trying to do something like that here:

  cuco::experimental::static_multimap<test_key_t, rowid_t> lineitem_table(
    std::ceil(num_join_4 / load_factor),
    cuco::empty_key{empty_key_sentinel},
    cuco::empty_value{empty_rowid_sentinel});

  // Build
  auto insert_iterator_5 = thrust::make_transform_iterator(
    thrust::make_zip_iterator(thrust::make_permutation_iterator(
                                l_suppkey.begin(),
                                thrust::make_permutation_iterator(l_rowid_3, derived_rowid_4)),
                              derived_rowid_4),
    MakeCucoPair<test_key_t, rowid_t>{});
  lineitem_table.insert(insert_iterator_5, insert_iterator_5 + num_join_4);

  // Probe
  auto probe_iterator_5 = thrust::make_zip_iterator(s_suppkey.begin(), s_nationkey.begin());
  lineitem_table.retrieve(probe_iterator_5,
                          probe_iterator_5 + num_supplier_mut,
                          thrust::make_zip_iterator(thrust::make_discard_iterator(), s_nationkey_5),
                          derived_rowid_5);

The problem, of course, is that dereferencing probe_iterator_5 returns a thrust tuple, whereas I really only want to probe the hash map with the first element of the tuple. Also, of course, a transform_iterator does not help here, because then the second element of the tuple is not exposed to the write to the output_probe iterator. Perhaps there is something I can do with rebinding the equality operator for the map? Much appreciation for your help.

EDIT: It looks like I can use for_each to accomplish my task after looking closer at the docs...

[PointKernel]

for_each will work but the performance will be sub-optimal.

Perhaps there is something I can do with rebinding the equality operator for the map? Much appreciation for your help.

Exactly. We expose retrieve API taking custom hasher and key equal in other containers like multiset, e.g.

cuCollections/include/cuco/static_multiset.cuh

Lines 746 to 757 in eb9319b

	template <class InputProbeIt,
	class ProbeEqual,
	class ProbeHash,
	class OutputProbeIt,
	class OutputMatchIt>
	std::pair<OutputProbeIt, OutputMatchIt> retrieve(InputProbeIt first,
	InputProbeIt last,
	ProbeEqual const& probe_equal,
	ProbeHash const& probe_hash,
	OutputProbeIt output_probe,
	OutputMatchIt output_match,
	cuda::stream_ref stream = {}) const;

And we definitely should expose this feature to multimap as well. Let me open a PR for this.

[kevkrist]

I know you already opened a PR for multimap, but I'm wondering whether the same functionality can be added to map and set, or whether there's a way to achieve the same task in a different way that I'm not seeing.

[PointKernel]

For set, find + thrust::copy_if can do the same work as retrieve and delivers better performance in general:

cuCollections/include/cuco/static_set.cuh

Lines 613 to 619 in a4fb985

	template <typename InputIt, typename ProbeEqual, typename ProbeHash, typename OutputIt>
	void find_async(InputIt first,
	InputIt last,
	ProbeEqual const& probe_equal,
	ProbeHash const& probe_hash,
	OutputIt output_begin,
	cuda::stream_ref stream = {}) const;

Let me open another PR adding the same find to map.

[kevkrist]

I see. So a separate compaction kernel generally provides better performance?

[PointKernel]

This conclusion is based on empirical experience in our workflow (see rapidsai/cudf#17278 (comment)).

retrieve is computationally expensive due to the use of atomic operations and high register pressure. Using find together with a separate compaction kernel can eliminate the costly atomic operation and reduce excessive register usage, leading to better performance for relatively large input data. However, performance can be influenced by various factors such as map occupancy, data size, and data type, etc. If you are aiming for the speed of light performance in a specific scenario, profiling the code and fine-tuning it accordingly may be necessary.

[kevkrist]

Roger that. We have found (unfortunately I can't point you to anything) in our own experiments that in-place compactions provide better performance when the join is highly selective (say, below 30%), because the atomics, register pressure, etc. are amortized by the reduced memory traffic, but of course this is a specific case. So, just putting the feature request out there for custom equals with retrieve one day. Thanks for all this!