diff --git a/cpp/src/arrow/acero/hash_join_node_test.cc b/cpp/src/arrow/acero/hash_join_node_test.cc index 73355af8c4b4..43e3c1863708 100644 --- a/cpp/src/arrow/acero/hash_join_node_test.cc +++ b/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(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. diff --git a/cpp/src/arrow/compute/row/grouper.cc b/cpp/src/arrow/compute/row/grouper.cc index d3eafef2cfa7..c4ba75c27cca 100644 --- a/cpp/src/arrow/compute/row/grouper.cc +++ b/cpp/src/arrow/compute/row/grouper.cc @@ -411,11 +411,11 @@ struct GrouperImpl : public Grouper { } template - 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(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(key_bytes_.size()); + const auto next_key_offset = static_cast(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 offsets_batch(batch.length + 1); + std::vector 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 map_; - std::vector offsets_ = {0}; + std::vector offsets_ = {0}; std::vector key_bytes_; uint32_t num_groups_ = 0; std::vector> encoders_; diff --git a/cpp/src/arrow/compute/row/grouper_test.cc b/cpp/src/arrow/compute/row/grouper_test.cc index edbf07e78588..6be438a11b72 100644 --- a/cpp/src/arrow/compute/row/grouper_test.cc +++ b/cpp/src/arrow/compute/row/grouper_test.cc @@ -21,6 +21,7 @@ #include #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" @@ -28,6 +29,7 @@ #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(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(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(*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(*unique_ids); + const auto& values = checked_cast(*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()}) { diff --git a/cpp/src/arrow/compute/row/row_encoder_internal.cc b/cpp/src/arrow/compute/row/row_encoder_internal.cc index ccf6642b950f..3c67fe759d79 100644 --- a/cpp/src/arrow/compute/row/row_encoder_internal.cc +++ b/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> 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> 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& 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; diff --git a/cpp/src/arrow/compute/row/row_encoder_internal.h b/cpp/src/arrow/compute/row/row_encoder_internal.h index 66595f3ef3e1..c3aa38831a77 100644 --- a/cpp/src/arrow/compute/row/row_encoder_internal.h +++ b/cpp/src/arrow/compute/row/row_encoder_internal.h @@ -48,11 +48,11 @@ struct ARROW_COMPUTE_EXPORT KeyEncoder { // sizeof(Offset) + buffer_size. // If Encoder is for null type, the length would add 0. virtual void AddLength(const ExecValue& value, int64_t batch_length, - int32_t* lengths) = 0; + int64_t* lengths) = 0; // Increment the length by the length of an encoded null value. // It's a special case for AddLength like `AddLength(Null-Scalar, 1, lengths)`. - virtual void AddLengthNull(int32_t* length) = 0; + virtual void AddLengthNull(int64_t* length) = 0; // Encode the column into the encoded_bytes, which is an array of pointers to each row // buffer. @@ -89,9 +89,9 @@ struct ARROW_COMPUTE_EXPORT KeyEncoder { struct ARROW_COMPUTE_EXPORT BooleanKeyEncoder : KeyEncoder { static constexpr int kByteWidth = 1; - void AddLength(const ExecValue& data, int64_t batch_length, int32_t* lengths) override; + void AddLength(const ExecValue& data, int64_t batch_length, int64_t* lengths) override; - void AddLengthNull(int32_t* length) override; + void AddLengthNull(int64_t* length) override; Status Encode(const ExecValue& data, int64_t batch_length, uint8_t** encoded_bytes) override; @@ -107,9 +107,9 @@ struct ARROW_COMPUTE_EXPORT FixedWidthKeyEncoder : KeyEncoder { : type_(std::move(type)), byte_width_(checked_cast(*type_).bit_width() / 8) {} - void AddLength(const ExecValue& data, int64_t batch_length, int32_t* lengths) override; + void AddLength(const ExecValue& data, int64_t batch_length, int64_t* lengths) override; - void AddLengthNull(int32_t* length) override; + void AddLengthNull(int64_t* length) override; Status Encode(const ExecValue& data, int64_t batch_length, uint8_t** encoded_bytes) override; @@ -148,7 +148,7 @@ struct BaseVarLengthKeyEncoder : KeyEncoder { explicit BaseVarLengthKeyEncoder(std::shared_ptr type) : type_(std::move(type)) {} - void AddLength(const ExecValue& data, int64_t batch_length, int32_t* lengths) override { + void AddLength(const ExecValue& data, int64_t batch_length, int64_t* lengths) override { if (data.is_array()) { int64_t i = 0; ARROW_DCHECK_EQ(data.array.length, batch_length); @@ -156,14 +156,14 @@ struct BaseVarLengthKeyEncoder : KeyEncoder { data.array, [&](std::string_view bytes) { lengths[i++] += - kExtraByteForNull + sizeof(Offset) + static_cast(bytes.size()); + kExtraByteForNull + sizeof(Offset) + static_cast(bytes.size()); }, [&] { lengths[i++] += kExtraByteForNull + sizeof(Offset); }); } else { const Scalar& scalar = *data.scalar; - const int32_t buffer_size = + const int64_t buffer_size = scalar.is_valid - ? static_cast(UnboxScalar::Unbox(scalar).size()) + ? static_cast(UnboxScalar::Unbox(scalar).size()) : 0; for (int64_t i = 0; i < batch_length; i++) { lengths[i] += kExtraByteForNull + sizeof(Offset) + buffer_size; @@ -171,7 +171,7 @@ struct BaseVarLengthKeyEncoder : KeyEncoder { } } - void AddLengthNull(int32_t* length) override { + void AddLengthNull(int64_t* length) override { *length += kExtraByteForNull + sizeof(Offset); } @@ -281,9 +281,9 @@ struct ARROW_COMPUTE_EXPORT BinaryViewKeyEncoder }; struct ARROW_COMPUTE_EXPORT NullKeyEncoder : KeyEncoder { - void AddLength(const ExecValue&, int64_t batch_length, int32_t* lengths) override {} + void AddLength(const ExecValue&, int64_t batch_length, int64_t* lengths) override {} - void AddLengthNull(int32_t* length) override {} + void AddLengthNull(int64_t* length) override {} Status Encode(const ExecValue& data, int64_t batch_length, uint8_t** encoded_bytes) override { @@ -377,7 +377,7 @@ class ARROW_COMPUTE_EXPORT RowEncoder { return std::string(reinterpret_cast(encoded_nulls_.data()), encoded_nulls_.size()); } - int32_t row_length = offsets_[i + 1] - offsets_[i]; + int64_t row_length = offsets_[i + 1] - offsets_[i]; return std::string(reinterpret_cast(bytes_.data() + offsets_[i]), row_length); } @@ -393,8 +393,9 @@ class ARROW_COMPUTE_EXPORT RowEncoder { // within the bytes_ vector. This allows for quick access to individual rows. // // The size would be num_rows + 1 if not empty, the last element is the total - // length of the bytes_ vector. - std::vector offsets_; + // length of the bytes_ vector. Use 64-bit offsets for buffers over 2GB, + // independently of the 32-bit row ids. + std::vector offsets_; // The encoded bytes of all non "kRowIdForNulls" rows. std::vector bytes_; // A pre-encoded constant row with all its columns being null. Useful when diff --git a/cpp/src/arrow/compute/row/row_encoder_internal_test.cc b/cpp/src/arrow/compute/row/row_encoder_internal_test.cc index 22a72117e9ff..3786aa9c4483 100644 --- a/cpp/src/arrow/compute/row/row_encoder_internal_test.cc +++ b/cpp/src/arrow/compute/row/row_encoder_internal_test.cc @@ -22,6 +22,7 @@ #include "arrow/array.h" #include "arrow/array/validate.h" +#include "arrow/buffer.h" #include "arrow/scalar.h" #include "arrow/testing/gtest_util.h" #include "arrow/type.h" @@ -35,7 +36,7 @@ TEST(TestKeyEncoder, BooleanScalar) { BooleanKeyEncoder key_encoder; SCOPED_TRACE("scalar " + scalar.ToString()); constexpr int64_t kBatchLength = 10; - std::array lengths{}; + std::array lengths{}; key_encoder.AddLength(ExecValue{&scalar}, kBatchLength, lengths.data()); // Check that the lengths are all 2. constexpr int32_t kPayloadWidth = @@ -72,10 +73,10 @@ TEST(TestKeyEncoder, BooleanScalar) { Result> RoundTripThroughKeyEncoder(KeyEncoder* encoder, const ExecValue& value, int32_t length) { - std::vector lengths(length, 0); + std::vector lengths(length, 0); encoder->AddLength(value, length, lengths.data()); - std::vector offsets(length + 1, 0); + std::vector offsets(length + 1, 0); for (int32_t i = 0; i < length; ++i) { offsets[i + 1] = offsets[i] + lengths[i]; } @@ -135,4 +136,94 @@ TEST(TestKeyEncoder, BinaryViewScalar) { } } +TEST(TestKeyEncoder, LARGE_MEMORY_TEST(ValueLengthOver2GB)) { + constexpr int64_t kValueLength = (int64_t{1} << 31) + 13; + constexpr int64_t kEncodedValueOverhead = sizeof(uint8_t) + sizeof(int64_t); + // A single value exceeds the signed 32-bit length limit. Only its length is + // inspected, so the buffer need not be initialized or copied by this test. + ASSERT_OK_AND_ASSIGN(std::shared_ptr bytes, AllocateBuffer(kValueLength)); + LargeBinaryScalar scalar(bytes); + VarLengthKeyEncoder encoder(large_binary()); + int64_t length = 0; + encoder.AddLength(ExecValue(&scalar), /*batch_length=*/1, &length); + // The encoded value includes a null byte and an int64_t length prefix. + ASSERT_EQ(length, kValueLength + kEncodedValueOverhead); + + // Length calculation must not narrow either scalar or array values, nor the + // combined length of multiple columns in a single encoded row. + // Reuse the same buffer through the array path and add to the existing length + // rather than resetting it, as when encoding a second column in the same row. + auto offsets = Buffer::FromVector({0, kValueLength}); + auto array = ArrayData::Make(large_binary(), /*length=*/1, {nullptr, offsets, bytes}, + /*null_count=*/0); + encoder.AddLength(ExecValue(*array), /*batch_length=*/1, &length); + ASSERT_EQ(length, 2 * (kValueLength + kEncodedValueOverhead)); +} + +TEST(TestRowEncoder, LARGE_MEMORY_TEST(AppendOver2GB)) { + constexpr int64_t kValueLength = 64 * 1024; + constexpr int32_t kBatchLength = 16385; + constexpr int32_t kNumBatches = 2; + constexpr int32_t kNumRows = kBatchLength * kNumBatches; + constexpr int32_t kNumDecodedValues = 4; + ExecContext ctx; + RowEncoder encoder; + encoder.Init({large_utf8()}, &ctx); + // Each batch encodes 16,385 * (64 KiB + 9) bytes, below the signed 32-bit + // offset limit. The second append takes the persistent row buffer over 2 GiB. + // Scalar input avoids allocating another expanded copy of those strings. + const auto scalar = std::make_shared(std::string(kValueLength, 'x')); + const ExecBatch batch({scalar}, kBatchLength); + ASSERT_OK(encoder.EncodeAndAppend(ExecSpan(batch))); + ASSERT_OK(encoder.EncodeAndAppend(ExecSpan(batch))); + ASSERT_EQ(encoder.num_rows(), kNumRows); + + // Read rows on both sides of the signed 32-bit offset boundary, including + // the special all-null row, without decoding the entire large buffer. + const int32_t row_ids[] = {0, kBatchLength - 1, kNumRows - 2, kNumRows - 1, + RowEncoder::kRowIdForNulls()}; + ASSERT_OK_AND_ASSIGN(auto decoded, encoder.Decode(kNumDecodedValues + 1, row_ids)); + ASSERT_OK_AND_ASSIGN(auto expected, MakeArrayFromScalar(*scalar, kNumDecodedValues)); + ASSERT_OK_AND_ASSIGN(auto nulls, MakeArrayOfNull(large_utf8(), /*length=*/1)); + AssertArraysEqual(*expected, + *decoded[0].make_array()->Slice(/*offset=*/0, kNumDecodedValues)); + AssertArraysEqual(*nulls, *decoded[0].make_array()->Slice(kNumDecodedValues)); + ASSERT_EQ(encoder.encoded_row(0), encoder.encoded_row(kNumRows - 1)); + + // Clear must also reset the large offsets so that the encoder can be reused. + encoder.Clear(); + ASSERT_EQ(encoder.num_rows(), 0); + ASSERT_OK(encoder.EncodeAndAppend(ExecSpan(ExecBatch({scalar}, /*length=*/1)))); + ASSERT_EQ(encoder.num_rows(), 1); +} + +TEST(TestRowEncoder, LARGE_MEMORY_TEST(FixedWidthDecodeAt2GB)) { + constexpr int32_t kValueLength = 64 * 1024; + constexpr int32_t kNumDecodedRows = 32 * 1024; + constexpr int64_t kDecodedBytes = int64_t{kNumDecodedRows} * kValueLength; + const std::string value(kValueLength, 'x'); + const auto scalar = std::make_shared(value); + ExecContext ctx; + RowEncoder encoder; + encoder.Init({fixed_size_binary(kValueLength)}, &ctx); + ASSERT_OK(encoder.EncodeAndAppend(ExecSpan(ExecBatch({scalar}, /*length=*/1)))); + ASSERT_EQ(encoder.num_rows(), 1); + + // Repeating a single row keeps the encoded input small. Only the decoded + // output reaches 32,768 * 64 KiB = 2 GiB, exceeding the signed 32-bit size + // limit independently of the offsets used to store encoded rows. + const std::vector row_ids(kNumDecodedRows, 0); + ASSERT_OK_AND_ASSIGN(auto decoded, encoder.Decode(kNumDecodedRows, row_ids.data())); + const auto array = + std::static_pointer_cast(decoded[0].make_array()); + ASSERT_OK(array->ValidateFull()); + ASSERT_EQ(array->length(), kNumDecodedRows); + ASSERT_EQ(array->null_count(), 0); + ASSERT_EQ(array->data()->buffers[1]->size(), kDecodedBytes); + // Check values throughout the output without allocating a second 2 GiB array. + for (const int32_t row_id : {0, kNumDecodedRows / 2, kNumDecodedRows - 1}) { + ASSERT_EQ(array->GetView(row_id), value); + } +} + } // namespace arrow::compute::internal