GH-45506: [C++][Acero] More overflow-safe Swiss table #45515

zanmato1984 · 2025-02-12T12:41:31Z

Rationale for this change

See #45506.

What changes are included in this PR?

Abstract current overflow-prone block data access into functions that do proper type promotion to avoid overflow. Also remove the old block base address accessor.
Unify the data types used for various concepts as they naturally are (i.e., w/o explicit promotion): uint32_t for block_id, int for num_xxx_bits/bytes, uint32_t for group_id, int for local_slot_id and uint32_t for global_slot_id.
Abstract several constants and utility functions for readability and maintainability.

Are these changes tested?

Existing tests should suffice.

It is really hard (gosh I did try) to create a concrete test case that fails w/o this change and passes w/ this change.

Are there any user-facing changes?

None.

GitHub Issue: [C++][Acero] Swiss table still has risks of overflow #45506

zanmato1984 · 2025-02-12T12:42:16Z

Most of this change are cleanup and refinement. @pitrou mind to take a look? Thanks.

zanmato1984 · 2025-02-12T12:48:28Z

cpp/src/arrow/acero/swiss_join.cc

@@ -643,37 +643,36 @@ void SwissTableMerge::MergePartition(SwissTable* target, const SwissTable* sourc
  //
  int source_group_id_bits =
      SwissTable::num_groupid_bits_from_log_blocks(source->log_blocks());
-  uint64_t source_group_id_mask = ~0ULL >> (64 - source_group_id_bits);


After cleaning up these unnecessary 64-bit in this file, we can further cleanup some temp states as mentioned in #45336 (comment) .

zanmato1984 · 2025-02-12T12:56:57Z

cpp/src/arrow/compute/key_map_internal.h

@@ -81,18 +81,29 @@ class ARROW_EXPORT SwissTable {

  void num_inserted(uint32_t i) { num_inserted_ = i; }

-  uint8_t* blocks() const { return blocks_->mutable_data(); }


This is the source of all evil. Let's get rid of it!

pitrou · 2025-02-12T15:49:53Z

cpp/src/arrow/compute/key_map_internal.cc

+      uint32_t group_id = *reinterpret_cast<const uint32_t*>(
+          block_data(block_id, num_block_bytes) + local_slots[id] * num_groupid_bytes +
+          bytes_status_in_block_);
+      group_id &= group_id_mask;


So we always issue a 32-bit load but then we optionally mask if the actual group id width is smaller? Don't we risk reading past block_data bounds here?

(also, should we use an unaligned load? see the SafeLoad and SafeLoadAs utility functions)

So we always issue a 32-bit load but then we optionally mask if the actual group id width is smaller? Don't we risk reading past block_data bounds here?

There will always be padding_ (64) extra bytes at the buffer end.

(also, should we use an unaligned load? see the SafeLoad and SafeLoadAs utility functions)

It seems so indeed, though I didn't change how the original code does it.

I'll update later.

(also, should we use an unaligned load? see the SafeLoad and SafeLoadAs utility functions)

It seems so indeed, though I didn't change how the original code does it.

I'll update later.

OK, turns out I was wrong. The original code uses aligned read and my change made it unaligned. I'll need to update it with more care. Thank you for pointing this out.

Changed back to the original aligned read with minor refinement.

Besides, I've also did some more cleanup during fixing the alignment issue. See my latest commit. Thanks.

pitrou · 2025-02-12T15:56:15Z

cpp/src/arrow/compute/key_map_internal.h

+    return static_cast<uint32_t>((1ULL << num_groupid_bits) - 1);
+  }
+
+  static constexpr int bytes_status_in_block_ = 8;


Usually, compile-time constants should follow the naming convention kBytesStatusInBlock. Perhaps we can have a renaming pass in this PR or another one?

Yes it is supposed to be. However I was also following the naming convention of several existing compile time constants in this class. I would like to to change them all in another PR to keep this one solely focused on the purpose the overflow prevention.

pitrou · 2025-02-12T15:57:22Z

cpp/src/arrow/compute/key_map_internal_avx2.cc

+    uint32_t mask = num_groupid_bytes == 1   ? 0xFF
+                    : num_groupid_bytes == 2 ? 0xFFFF
+                                             : 0xFFFFFFFF;


Is there a reason for expanding the possible values instead of simply using the usual bitshift formula?

Not particularly. This is just moving the original code.

pitrou · 2025-02-12T16:00:18Z

cpp/src/arrow/compute/key_map_internal_avx2.cc

-      __m128i group_id_lo = _mm256_i64gather_epi32(elements, pos_lo, 1);
-      __m128i group_id_hi = _mm256_i64gather_epi32(elements, pos_hi, 1);
+      local_slot_lo =
+          _mm256_mul_epu32(local_slot_lo, _mm256_set1_epi32(num_groupid_bytes));


This is using a 32-bit multiply even though local_slot_lo is supposed to be a vector of 64-bit ints? This might be correct because most bytes are zero, but I would at least expect an explanatory comment :)

Ah, my bad, _mm256_mul_epu32 is actually a 64-bit multiply.

cpp/src/arrow/compute/key_map_internal.h

pitrou · 2025-02-13T15:03:51Z

cpp/src/arrow/compute/key_map_internal.h

-                                   uint64_t group_id_mask) const;
+  static uint32_t extract_group_id(const uint8_t* block_ptr, int local_slot,
+                                   int num_group_id_bits) {
+    // Extract group id using aligned 32-bit read.


Why not use SafeLoad as in other places already?
(also, since this is non-trivial, factoring out the loading of a group id could go into a dedicated inline function)

The original code uses aligned read + masking so I'm following it, possibly for performance sake I guess?

If SafeLoad is preferred (i.e., it doesn't hurt performance), then yes it is possible to factor this piece of code out.

For the record, there are three places doing group id extraction:

Here, extracting single group id, publicly used by swiss join: currently using aligned read + masking;

extract_group_ids, extracting a vector of group ids, internally used: using aligned read w/o masking (the number of bits is constant-ized as template parameter);

grow_double, extracting single group id inside a big loop, inlined: using unaligned read + masking.

I think we should at least keep 2) as is because it makes perfect sense. 1) and 3) can be unified, either aligned or unaligned.

What do you think?

I think we should at least keep 2) as is because it makes perfect sense. 1) and 3) can be unified, either aligned or unaligned.

Agreed. Feel free to choose whatever approach you prefer!

Addressed. Thank you for the suggestion.

This PR is good for review again.

zanmato1984 requested a review from westonpace as a code owner February 12, 2025 12:41

github-actions bot added the Component: C++ label Feb 12, 2025

github-actions bot added the awaiting review Awaiting review label Feb 12, 2025

zanmato1984 commented Feb 12, 2025

View reviewed changes

github-actions bot added awaiting committer review Awaiting committer review and removed awaiting review Awaiting review labels Feb 12, 2025

More overflow-safe swiss table.

1dd5c08

zanmato1984 force-pushed the more-overflow-safe-swiss-table branch from af1c470 to 1dd5c08 Compare February 12, 2025 12:53

zanmato1984 commented Feb 12, 2025

View reviewed changes

zanmato1984 added 2 commits February 12, 2025 20:59

Remove already implied casting

a1f9758

Revert some unintended reordering

f5db159

pitrou reviewed Feb 12, 2025

View reviewed changes

zanmato1984 commented Feb 13, 2025

View reviewed changes

cpp/src/arrow/compute/key_map_internal.h Show resolved Hide resolved

zanmato1984 added 2 commits February 13, 2025 19:58

Use aligned read to extract group id

7af1d3c

Some more cleanup found in the last commit

aeb8b9d

zanmato1984 force-pushed the more-overflow-safe-swiss-table branch from 1f9dde4 to aeb8b9d Compare February 13, 2025 14:40

pitrou reviewed Feb 13, 2025

View reviewed changes

Unify extracting group id code by using aligned 32-bit read

95c6990

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

GH-45506: [C++][Acero] More overflow-safe Swiss table #45515

GH-45506: [C++][Acero] More overflow-safe Swiss table #45515

zanmato1984 commented Feb 12, 2025 •

edited

Loading

zanmato1984 commented Feb 12, 2025

zanmato1984 Feb 12, 2025

zanmato1984 Feb 12, 2025

pitrou Feb 12, 2025 •

edited

Loading

zanmato1984 Feb 12, 2025

zanmato1984 Feb 13, 2025

zanmato1984 Feb 13, 2025

pitrou Feb 12, 2025

zanmato1984 Feb 12, 2025

pitrou Feb 12, 2025

zanmato1984 Feb 12, 2025

pitrou Feb 12, 2025

pitrou Feb 12, 2025

pitrou Feb 13, 2025

zanmato1984 Feb 13, 2025

zanmato1984 Feb 13, 2025

pitrou Feb 13, 2025

zanmato1984 Feb 13, 2025

		@@ -81,18 +81,29 @@ class ARROW_EXPORT SwissTable {

		void num_inserted(uint32_t i) { num_inserted_ = i; }

		uint8_t* blocks() const { return blocks_->mutable_data(); }

GH-45506: [C++][Acero] More overflow-safe Swiss table #45515

Are you sure you want to change the base?

GH-45506: [C++][Acero] More overflow-safe Swiss table #45515

Conversation

zanmato1984 commented Feb 12, 2025 • edited Loading

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

zanmato1984 commented Feb 12, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pitrou Feb 12, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zanmato1984 commented Feb 12, 2025 •

edited

Loading

pitrou Feb 12, 2025 •

edited

Loading