diff --git a/cpp/src/arrow/acero/swiss_join.cc b/cpp/src/arrow/acero/swiss_join.cc index 3f11b89af39..6fc99409091 100644 --- a/cpp/src/arrow/acero/swiss_join.cc +++ b/cpp/src/arrow/acero/swiss_join.cc @@ -439,6 +439,7 @@ void RowArray::DebugPrintToFile(const char* filename, bool print_sorted) const { Status RowArrayMerge::PrepareForMerge(RowArray* target, const std::vector& sources, + bool is_key_data, std::vector* first_target_row_id, MemoryPool* pool) { ARROW_DCHECK(!sources.empty()); @@ -473,7 +474,7 @@ Status RowArrayMerge::PrepareForMerge(RowArray* target, (*first_target_row_id)[sources.size()] = num_rows; } - if (num_bytes > std::numeric_limits::max()) { + if (is_key_data && num_bytes > std::numeric_limits::max()) { return Status::Invalid( "There are more than 2^32 bytes of key data. Acero cannot " "process a join of this magnitude"); @@ -1330,7 +1331,8 @@ Status SwissTableForJoinBuild::PreparePrtnMerge() { for (int i = 0; i < num_prtns_; ++i) { partition_keys[i] = prtn_states_[i].keys.keys(); } - RETURN_NOT_OK(RowArrayMerge::PrepareForMerge(target_->map_.keys(), partition_keys, + + RETURN_NOT_OK(RowArrayMerge::PrepareForMerge(target_->map_.keys(), partition_keys, true, &partition_keys_first_row_id_, pool_)); // 2. SwissTable: @@ -1353,7 +1355,7 @@ Status SwissTableForJoinBuild::PreparePrtnMerge() { partition_payloads[i] = &prtn_states_[i].payloads; } RETURN_NOT_OK(RowArrayMerge::PrepareForMerge(&target_->payloads_, partition_payloads, - &partition_payloads_first_row_id_, + false, &partition_payloads_first_row_id_, pool_)); } diff --git a/cpp/src/arrow/acero/swiss_join_internal.h b/cpp/src/arrow/acero/swiss_join_internal.h index 88b80f06f57..50efaba9d9c 100644 --- a/cpp/src/arrow/acero/swiss_join_internal.h +++ b/cpp/src/arrow/acero/swiss_join_internal.h @@ -160,6 +160,7 @@ class RowArrayMerge { // caller can pass in nullptr to indicate that it is not needed. // static Status PrepareForMerge(RowArray* target, const std::vector& sources, + bool is_key_data, std::vector* first_target_row_id, MemoryPool* pool);