GH-54885: [C++][Compute] Use 64-bit offsets for basic row encoding

cpp/src/arrow/acero/hash_join_node_test.cc

@@ -3312,6 +3312,110 @@ void AssertRowCountEq(Declaration source, int64_t expected) {
 
 }  // namespace
 
+TEST(HashJoin, LARGE_MEMORY_TEST(BuildSideOver2GBLargeBinaryKeys)) {
+  constexpr int64_t kValueLength = 64 * 1024;
+  constexpr int32_t kBatchLength = 1024;
+  constexpr int32_t kNumBatches = 33;
+  // large_utf8 selects the basic hash join. The build-side key encoder stores
+  // 33 * 1024 * 64 KiB of string bytes (plus row metadata), exceeding 2 GiB.
+  // Reuse a 64 MiB input array across batches rather than allocating 33 copies.
+  const auto type = large_utf8();
+  const auto value = std::make_shared<LargeStringScalar>(std::string(kValueLength, 'x'));
+  auto l_schema = schema({field("l_id", int32()), field("l_key", type)});
+  auto r_schema = schema({field("r_id", int32()), field("r_key", type)});
+  ASSERT_OK_AND_ASSIGN(auto l_keys, MakeArrayFromScalar(*value, /*length=*/2));
+  ASSERT_OK_AND_ASSIGN(auto r_keys, MakeArrayFromScalar(*value, kBatchLength));
+  Int32Builder ids_builder;
+  for (int32_t i = 0; i < kBatchLength; ++i) {
+    ASSERT_OK(ids_builder.Append(i));
+  }
+  ASSERT_OK_AND_ASSIGN(auto ids, ids_builder.Finish());
+  BatchesWithSchema left{{ExecBatch({ArrayFromJSON(int32(), "[0, -1]"), l_keys},
+                                    /*length=*/2)},
+                         l_schema};
+  BatchesWithSchema right{{}, r_schema};
+  right.batches.assign(kNumBatches, ExecBatch({ids, r_keys}, kBatchLength));
+
+  // Join on both id and string: id 0 matches once in each build batch, while
+  // id -1 never matches. INNER returns 33 rows and LEFT_ANTI returns only -1.
+  // Keep expected output small while comparing the complete large string values.
+  ASSERT_OK_AND_ASSIGN(auto expected_key, MakeArrayFromScalar(*value, /*length=*/1));
+  const auto zero = ArrayFromJSON(int32(), "[0]");
+  BatchesWithSchema expected{{},
+                             schema({field("l_id", int32()), field("l_key", type),
+                                     field("r_id", int32()), field("r_key", type)})};
+  expected.batches.assign(
+      kNumBatches, ExecBatch({zero, expected_key, zero, expected_key}, /*length=*/1));
+  BatchesWithSchema anti_expected{
+      {ExecBatch({ArrayFromJSON(int32(), "[-1]"), expected_key}, /*length=*/1)},
+      l_schema};
+  // Check both join types in serial and parallel execution.
+  for (bool parallel : {false, true}) {
+    CheckRunOutput(JoinType::INNER, left, right, {"l_id", "l_key"}, {"r_id", "r_key"},
+                   expected, parallel);
+    CheckRunOutput(JoinType::LEFT_ANTI, left, right, {"l_id", "l_key"}, {"r_id", "r_key"},
+                   anti_expected, parallel);
+  }
+}
+
+TEST(HashJoin, LARGE_MEMORY_TEST(BuildSideOver2GBLargeBinaryPayloads)) {
+  constexpr int64_t kValueLength = 64 * 1024;
+  constexpr int32_t kBatchLength = 1024;
+  constexpr int32_t kNumBatches = 17;
+  // Keep join keys small so that it is the build-side payload encoder, not the
+  // key encoder, that crosses 2 GiB. Both payload columns enter that row buffer.
+  Int32Builder ids_builder;
+  BinaryBuilder keys_builder;
+  for (int32_t i = 0; i < kBatchLength; ++i) {
+    ASSERT_OK(ids_builder.Append(i));
+    ASSERT_OK(keys_builder.Append(std::to_string(i)));
+  }
+  ASSERT_OK_AND_ASSIGN(auto ids, ids_builder.Finish());
+  ASSERT_OK_AND_ASSIGN(auto keys, keys_builder.Finish());
+  const auto zero_id = ArrayFromJSON(int32(), "[0]");
+  const auto zero_key = ArrayFromJSON(binary(), R"(["0"])");
+
+  for (const auto& type : {large_utf8(), large_binary()}) {
+    SCOPED_TRACE(type->ToString());
+    ASSERT_OK_AND_ASSIGN(auto value, MakeScalar(type, std::string(kValueLength, 'x')));
+    ASSERT_OK_AND_ASSIGN(auto payload, MakeArrayFromScalar(*value, kBatchLength));
+    auto l_schema = schema({field("l_id", int32()), field("l_key", binary())});
+    auto r_schema = schema({field("r_id", int32()), field("r_key", binary()),
+                            field("r_payload1", type), field("r_payload2", type)});
+    BatchesWithSchema left{{ExecBatch({ArrayFromJSON(int32(), "[0, -1]"),
+                                       ArrayFromJSON(binary(), R"(["0", "-1"])")},
+                                      /*length=*/2)},
+                           l_schema};
+    BatchesWithSchema right{{}, r_schema};
+    // 17 * 1024 * 2 * 64 KiB exceeds 2 GiB even before encoding metadata.
+    // Input columns and batches alias the same array; the payload encoder still
+    // writes every row and both columns into its persistent build-side buffer.
+    right.batches.assign(kNumBatches,
+                         ExecBatch({ids, keys, payload, payload}, kBatchLength));
+
+    // id/key 0 matches once per build batch and -1 never matches. Compare all
+    // six output columns, including both complete payload values, in 17 rows.
+    ASSERT_OK_AND_ASSIGN(auto expected_payload,
+                         MakeArrayFromScalar(*value, /*length=*/1));
+    BatchesWithSchema expected{
+        {},
+        schema({field("l_id", int32()), field("l_key", binary()), field("r_id", int32()),
+                field("r_key", binary()), field("r_payload1", type),
+                field("r_payload2", type)})};
+    expected.batches.assign(kNumBatches, ExecBatch({zero_id, zero_key, zero_id, zero_key,
+                                                    expected_payload, expected_payload},
+                                                   /*length=*/1));
+    // Two payload types, two small-key encodings and serial/parallel execution
+    // cover eight INNER join combinations without growing the expected output.
+    for (bool binary_keys : {false, true}) {
+      for (bool parallel : {false, true}) {
+        CheckRunOutput(JoinType::INNER, left, right, {binary_keys ? "l_key" : "l_id"},
+                       {binary_keys ? "r_key" : "r_id"}, expected, parallel);
+      }
+    }
+  }
+}
+
 // GH-43495: Test that both the key and the payload of the right side (the build side) are
 // fixed length and larger than 4GB, and the 64-bit offset in the hash table can handle it
 // correctly.

cpp/src/arrow/compute/row/grouper.cc

@@ -411,11 +411,11 @@ struct GrouperImpl : public Grouper {
   }
 
   template <typename VisitGroupFunc, typename VisitUnknownGroupFunc>
-  void VisitKeys(int64_t length, const int32_t* key_offsets, const uint8_t* key_data,
+  void VisitKeys(int64_t length, const int64_t* key_offsets, const uint8_t* key_data,
                  bool insert_new_keys, VisitGroupFunc&& visit_group,
                  VisitUnknownGroupFunc&& visit_unknown_group) {
     for (int64_t i = 0; i < length; ++i) {
-      const int32_t key_length = key_offsets[i + 1] - key_offsets[i];
+      const int64_t key_length = key_offsets[i + 1] - key_offsets[i];
       const uint8_t* key_ptr = key_data + key_offsets[i];
       std::string key(reinterpret_cast<const char*>(key_ptr), key_length);
 
@@ -426,7 +426,7 @@ struct GrouperImpl : public Grouper {
           // New key: update offsets and key_bytes
           ++num_groups_;
           if (key_length > 0) {
-            const auto next_key_offset = static_cast<int32_t>(key_bytes_.size());
+            const auto next_key_offset = static_cast<int64_t>(key_bytes_.size());
             key_bytes_.resize(next_key_offset + key_length);
             offsets_.push_back(next_key_offset + key_length);
             memcpy(key_bytes_.data() + next_key_offset, key_ptr, key_length);
@@ -453,12 +453,12 @@ struct GrouperImpl : public Grouper {
       auto batch_slice = batch.ToExecBatch().Slice(offset, length);
       return ConsumeImpl(ExecSpan(batch_slice), 0, -1, mode);
     }
-    std::vector<int32_t> offsets_batch(batch.length + 1);
+    std::vector<int64_t> offsets_batch(batch.length + 1);
     for (int i = 0; i < batch.num_values(); ++i) {
       encoders_[i]->AddLength(batch[i], batch.length, offsets_batch.data());
     }
 
-    int32_t total_length = 0;
+    int64_t total_length = 0;
     for (int64_t i = 0; i < batch.length; ++i) {
       auto total_length_before = total_length;
       total_length += offsets_batch[i];
@@ -546,7 +546,7 @@ struct GrouperImpl : public Grouper {
   ExecContext* ctx_;
   // TODO We could use std::string_view since the keys are copied in key_bytes_.
   std::unordered_map<std::string, uint32_t> map_;
-  std::vector<int32_t> offsets_ = {0};
+  std::vector<int64_t> offsets_ = {0};
   std::vector<uint8_t> key_bytes_;
   uint32_t num_groups_ = 0;
   std::vector<std::unique_ptr<internal::KeyEncoder>> encoders_;

cpp/src/arrow/compute/row/grouper_test.cc

@@ -21,13 +21,15 @@
 #include <gtest/gtest.h>
 
 #include "arrow/array.h"
+#include "arrow/array/builder_primitive.h"
 #include "arrow/array/util.h"
 #include "arrow/compute/api_vector.h"
 #include "arrow/compute/cast.h"
 #include "arrow/compute/exec.h"
 #include "arrow/compute/row/grouper.h"
 #include "arrow/compute/row/grouper_internal.h"
 #include "arrow/compute/test_util_internal.h"
+#include "arrow/scalar.h"
 #include "arrow/testing/gtest_util.h"
 #include "arrow/testing/matchers.h"
 #include "arrow/testing/random.h"
@@ -927,6 +929,97 @@ TEST(Grouper, StringKey) {
   }
 }
 
+TEST(Grouper, LARGE_MEMORY_TEST(BatchOver2GB)) {
+  constexpr int64_t kValueLength = 64 * 1024;
+  constexpr int64_t kBatchLength = 32768;
+  // large_utf8 selects GrouperImpl. Repeated scalar keys keep the input and
+  // unique-key storage small. 32,768 * 64 KiB is exactly 2 GiB of value bytes;
+  // adding the null byte and int64_t length per row takes the encoded batch
+  // beyond 2 GiB, even though there is only one distinct key.
+  const auto value = std::make_shared<LargeStringScalar>(std::string(kValueLength, 'x'));
+  const ExecBatch batch({value}, kBatchLength);
+  ASSERT_OK_AND_ASSIGN(auto grouper, Grouper::Make({large_utf8()}));
+  ASSERT_OK_AND_ASSIGN(auto expected, MakeArrayFromScalar(UInt32Scalar{0}, kBatchLength));
+
+  // All rows must receive group 0, despite the large temporary batch encoding.
+  ASSERT_OK_AND_ASSIGN(auto groups, grouper->Consume(ExecSpan(batch)));
+  AssertArraysEqual(*expected, *groups.make_array());
+  ASSERT_EQ(grouper->num_groups(), 1);
+
+  // Re-inserting the same key must not create more groups; Lookup must find it
+  // for every row, and GetUniques must decode just the one original value.
+  ASSERT_OK(grouper->Populate(ExecSpan(batch)));
+  ASSERT_OK_AND_ASSIGN(auto found, grouper->Lookup(ExecSpan(batch)));
+  AssertArraysEqual(*expected, *found.make_array());
+  ASSERT_EQ(grouper->num_groups(), 1);
+
+  ASSERT_OK_AND_ASSIGN(auto uniques, grouper->GetUniques());
+  ASSERT_OK_AND_ASSIGN(auto expected_unique, MakeArrayFromScalar(*value, /*length=*/1));
+  AssertArraysEqual(*expected_unique, *uniques[0].make_array());
+}
+
+TEST(Grouper, LARGE_MEMORY_TEST(UniqueKeysOver2GB)) {
+  constexpr int32_t kBatchLength = 1024;
+  constexpr int32_t kNumBatches = 32;
+  constexpr int32_t kNumGroups = kBatchLength * kNumBatches;
+  constexpr int64_t kValueLength = 64 * 1024;
+  // Each batch encodes about 64 MiB, but 32 * 1024 * 64 KiB is already 2 GiB
+  // of value bytes. The int32 key, null bytes and length prefixes take the
+  // persistent unique-key buffer beyond 2 GiB. Distinct integer ids prevent
+  // deduplication even though every row reuses the same large binary scalar.
+  const auto value = std::make_shared<LargeBinaryScalar>(std::string(kValueLength, 'x'));
+  ASSERT_OK_AND_ASSIGN(auto grouper, Grouper::Make({int32(), large_binary()}));
+
+  for (int32_t batch_index = 0; batch_index < kNumBatches; ++batch_index) {
+    SCOPED_TRACE(batch_index);
+    Int32Builder builder;
+    const int32_t first_id = batch_index * kBatchLength;
+    for (int32_t i = 0; i < kBatchLength; ++i) {
+      ASSERT_OK(builder.Append(first_id + i));
+    }
+    ASSERT_OK_AND_ASSIGN(auto ids, builder.Finish());
+    const ExecBatch batch({ids, value}, kBatchLength);
+    // Exercise both insertion APIs on the same persistent buffer: Populate
+    // first, then Consume, whose returned group ids must follow insertion order.
+    if (batch_index < kNumBatches / 2) {
+      ASSERT_OK(grouper->Populate(ExecSpan(batch)));
+    } else {
+      ASSERT_OK_AND_ASSIGN(auto groups, grouper->Consume(ExecSpan(batch)));
+      auto group_array = groups.make_array();
+      const auto& group_ids = checked_cast<const UInt32Array&>(*group_array);
+      for (int32_t i = 0; i < kBatchLength; ++i) {
+        ASSERT_EQ(group_ids.Value(i), first_id + i);
+      }
+    }
+  }
+  ASSERT_EQ(grouper->num_groups(), kNumGroups);
+
+  // Look up the last key (32767), the first key and a missing key. Lookup must
+  // return null for the miss without inserting a new group.
+  const ExecBatch query({ArrayFromJSON(int32(), "[32767, 0, -1]"), value}, /*length=*/3);
+  ASSERT_OK_AND_ASSIGN(auto found, grouper->Lookup(ExecSpan(query)));
+  AssertArraysEqual(*ArrayFromJSON(uint32(), "[32767, 0, null]"), *found.make_array());
+  ASSERT_EQ(grouper->num_groups(), kNumGroups);
+
+  // Decode all unique keys to validate the full offsets and value buffer, then
+  // check insertion order and sample large values at the beginning, middle and end.
+  ASSERT_OK_AND_ASSIGN(auto uniques, grouper->GetUniques());
+  ASSERT_EQ(uniques.length, kNumGroups);
+  auto unique_ids = uniques[0].make_array();
+  auto unique_values = uniques[1].make_array();
+  ASSERT_OK(unique_ids->ValidateFull());
+  ASSERT_OK(unique_values->ValidateFull());
+  const auto& ids = checked_cast<const Int32Array&>(*unique_ids);
+  const auto& values = checked_cast<const LargeBinaryArray&>(*unique_values);
+  ASSERT_EQ(values.value_data()->size(), kValueLength * kNumGroups);
+  for (int32_t i = 0; i < kNumGroups; ++i) {
+    ASSERT_EQ(ids.Value(i), i);
+  }
+  for (int32_t i : {0, kNumGroups / 2, kNumGroups - 1}) {
+    ASSERT_EQ(values.GetView(i), std::string_view(*value->value));
+  }
+}
+
 TEST(Grouper, StringViewKey) {
   // Mix inline (<=12 byte) and out-of-line view keys.
   for (auto ty : {utf8_view(), binary_view()}) {

cpp/src/arrow/compute/row/row_encoder_internal.cc

@@ -64,13 +64,13 @@ Status KeyEncoder::DecodeNulls(MemoryPool* pool, int32_t length, uint8_t** encod
 }
 
 void BooleanKeyEncoder::AddLength(const ExecValue&, int64_t batch_length,
-                                  int32_t* lengths) {
+                                  int64_t* lengths) {
   for (int64_t i = 0; i < batch_length; ++i) {
     lengths[i] += kByteWidth + kExtraByteForNull;
   }
 }
 
-void BooleanKeyEncoder::AddLengthNull(int32_t* length) {
+void BooleanKeyEncoder::AddLengthNull(int64_t* length) {
   *length += kByteWidth + kExtraByteForNull;
 }
 
@@ -134,13 +134,13 @@ Result<std::shared_ptr<ArrayData>> BooleanKeyEncoder::Decode(uint8_t** encoded_b
 }
 
 void FixedWidthKeyEncoder::AddLength(const ExecValue&, int64_t batch_length,
-                                     int32_t* lengths) {
+                                     int64_t* lengths) {
   for (int64_t i = 0; i < batch_length; ++i) {
     lengths[i] += byte_width_ + kExtraByteForNull;
   }
 }
 
-void FixedWidthKeyEncoder::AddLengthNull(int32_t* length) {
+void FixedWidthKeyEncoder::AddLengthNull(int64_t* length) {
   *length += byte_width_ + kExtraByteForNull;
 }
 
@@ -197,7 +197,8 @@ Result<std::shared_ptr<ArrayData>> FixedWidthKeyEncoder::Decode(uint8_t** encode
   int32_t null_count;
   RETURN_NOT_OK(DecodeNulls(pool, length, encoded_bytes, &null_buf, &null_count));
 
-  ARROW_ASSIGN_OR_RAISE(auto key_buf, AllocateBuffer(length * byte_width_, pool));
+  ARROW_ASSIGN_OR_RAISE(auto key_buf,
+                        AllocateBuffer(int64_t{length} * byte_width_, pool));
 
   uint8_t* raw_output = key_buf->mutable_data();
   for (int32_t i = 0; i < length; ++i) {
@@ -339,7 +340,7 @@ void RowEncoder::Init(const std::vector<TypeHolder>& column_types, ExecContext*
     ARROW_DCHECK(false);
   }
 
-  int32_t total_length = 0;
+  int64_t total_length = 0;
   for (size_t i = 0; i < column_types.size(); ++i) {
     encoders_[i]->AddLengthNull(&total_length);
   }
@@ -370,7 +371,7 @@ Status RowEncoder::EncodeAndAppend(const ExecSpan& batch) {
     encoders_[i]->AddLength(batch[i], batch.length, offsets_.data() + length_before + 1);
   }
 
-  int32_t total_length = offsets_[length_before];
+  int64_t total_length = offsets_[length_before];
   for (int64_t i = 0; i < batch.length; ++i) {
     total_length += offsets_[length_before + 1 + i];
     offsets_[length_before + 1 + i] = total_length;