pitrou commented on code in PR #55170:
URL: https://github.com/apache/arrow/pull/55170#discussion_r4376114950
##########
cpp/src/arrow/compute/row/grouper_test.cc:
##########
@@ -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)));
Review Comment:
Same here: validate `found`
##########
cpp/src/arrow/compute/row/row_encoder_internal.cc:
##########
@@ -356,34 +357,33 @@ void RowEncoder::Clear() {
}
Status RowEncoder::EncodeAndAppend(const ExecSpan& batch) {
+ const int64_t batch_length = batch.length;
Review Comment:
Is there a point in this? Casting to `int64_t` perhaps?
##########
cpp/src/arrow/compute/row/grouper_test.cc:
##########
@@ -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)));
Review Comment:
Also validate `groups`?
##########
cpp/src/arrow/compute/row/row_encoder_internal_test.cc:
##########
@@ -135,4 +136,122 @@ TEST(TestKeyEncoder, BinaryViewScalar) {
}
}
+TEST(TestRowEncoder, AppendSlicesAndReuse) {
+ for (const auto& type :
+ {binary(), utf8(), large_binary(), large_utf8(), binary_view(),
utf8_view()}) {
+ SCOPED_TRACE(type->ToString());
+ auto array = ArrayFromJSON(
+ type, R"(["unused", "first", null, "", "a long out-of-line value",
"last"])");
+ auto expected = array->Slice(/*offset=*/1);
+ ExecContext ctx;
+ RowEncoder encoder;
+ encoder.Init({type}, &ctx);
+ // Repeat after Clear to exercise appends into retained, previously used
capacity.
+ for (int round = 0; round < 2; ++round) {
+ for (const auto& [offset, length] :
+ {std::pair{1, 3}, std::pair{4, 0}, std::pair{4, 2}}) {
+ const ExecBatch batch({array->Slice(offset, length)}, length);
+ ASSERT_OK(encoder.EncodeAndAppend(ExecSpan(batch)));
+ }
+ ASSERT_EQ(encoder.num_rows(), expected->length());
+ const int32_t row_ids[] = {0, 1, 2, 3, 4};
+ ASSERT_OK_AND_ASSIGN(auto decoded, encoder.Decode(expected->length(),
row_ids));
+ ASSERT_OK(decoded[0].make_array()->ValidateFull());
+ AssertArraysEqual(*expected, *decoded[0].make_array(), /*verbose=*/true);
+ encoder.Clear();
+ ASSERT_EQ(encoder.num_rows(), 0);
+ }
+ }
+}
+
+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<Buffer> bytes,
AllocateBuffer(kValueLength));
+ LargeBinaryScalar scalar(bytes);
+ VarLengthKeyEncoder<LargeBinaryType> 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<int64_t>({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<LargeStringScalar>(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));
Review Comment:
Validate `decoded`?
##########
cpp/src/arrow/compute/row/grouper_test.cc:
##########
@@ -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();
Review Comment:
Same: validate `groups`
##########
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
Review Comment:
"Payload encoder" means the join result, or something else?
##########
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()}) {
Review Comment:
Do we actually need to test both types?
##########
cpp/src/arrow/compute/row/row_encoder_internal_test.cc:
##########
@@ -135,4 +136,122 @@ TEST(TestKeyEncoder, BinaryViewScalar) {
}
}
+TEST(TestRowEncoder, AppendSlicesAndReuse) {
+ for (const auto& type :
+ {binary(), utf8(), large_binary(), large_utf8(), binary_view(),
utf8_view()}) {
+ SCOPED_TRACE(type->ToString());
+ auto array = ArrayFromJSON(
+ type, R"(["unused", "first", null, "", "a long out-of-line value",
"last"])");
+ auto expected = array->Slice(/*offset=*/1);
+ ExecContext ctx;
+ RowEncoder encoder;
+ encoder.Init({type}, &ctx);
+ // Repeat after Clear to exercise appends into retained, previously used
capacity.
+ for (int round = 0; round < 2; ++round) {
+ for (const auto& [offset, length] :
+ {std::pair{1, 3}, std::pair{4, 0}, std::pair{4, 2}}) {
+ const ExecBatch batch({array->Slice(offset, length)}, length);
+ ASSERT_OK(encoder.EncodeAndAppend(ExecSpan(batch)));
+ }
+ ASSERT_EQ(encoder.num_rows(), expected->length());
+ const int32_t row_ids[] = {0, 1, 2, 3, 4};
+ ASSERT_OK_AND_ASSIGN(auto decoded, encoder.Decode(expected->length(),
row_ids));
+ ASSERT_OK(decoded[0].make_array()->ValidateFull());
+ AssertArraysEqual(*expected, *decoded[0].make_array(), /*verbose=*/true);
+ encoder.Clear();
+ ASSERT_EQ(encoder.num_rows(), 0);
+ }
+ }
+}
+
+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<Buffer> bytes,
AllocateBuffer(kValueLength));
+ LargeBinaryScalar scalar(bytes);
+ VarLengthKeyEncoder<LargeBinaryType> 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<int64_t>({0, kValueLength});
+ auto array = ArrayData::Make(large_binary(), /*length=*/1, {nullptr,
offsets, bytes},
+ /*null_count=*/0);
+ encoder.AddLength(ExecValue(*array), /*batch_length=*/1, &length);
Review Comment:
Same question here.
##########
cpp/src/arrow/compute/row/row_encoder_internal_test.cc:
##########
@@ -135,4 +136,122 @@ TEST(TestKeyEncoder, BinaryViewScalar) {
}
}
+TEST(TestRowEncoder, AppendSlicesAndReuse) {
+ for (const auto& type :
+ {binary(), utf8(), large_binary(), large_utf8(), binary_view(),
utf8_view()}) {
+ SCOPED_TRACE(type->ToString());
+ auto array = ArrayFromJSON(
+ type, R"(["unused", "first", null, "", "a long out-of-line value",
"last"])");
+ auto expected = array->Slice(/*offset=*/1);
+ ExecContext ctx;
+ RowEncoder encoder;
+ encoder.Init({type}, &ctx);
+ // Repeat after Clear to exercise appends into retained, previously used
capacity.
+ for (int round = 0; round < 2; ++round) {
+ for (const auto& [offset, length] :
+ {std::pair{1, 3}, std::pair{4, 0}, std::pair{4, 2}}) {
+ const ExecBatch batch({array->Slice(offset, length)}, length);
+ ASSERT_OK(encoder.EncodeAndAppend(ExecSpan(batch)));
+ }
+ ASSERT_EQ(encoder.num_rows(), expected->length());
+ const int32_t row_ids[] = {0, 1, 2, 3, 4};
+ ASSERT_OK_AND_ASSIGN(auto decoded, encoder.Decode(expected->length(),
row_ids));
+ ASSERT_OK(decoded[0].make_array()->ValidateFull());
+ AssertArraysEqual(*expected, *decoded[0].make_array(), /*verbose=*/true);
+ encoder.Clear();
+ ASSERT_EQ(encoder.num_rows(), 0);
+ }
+ }
+}
+
+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<Buffer> bytes,
AllocateBuffer(kValueLength));
+ LargeBinaryScalar scalar(bytes);
+ VarLengthKeyEncoder<LargeBinaryType> encoder(large_binary());
+ int64_t length = 0;
+ encoder.AddLength(ExecValue(&scalar), /*batch_length=*/1, &length);
Review Comment:
Why not use a non-trivial batch length?
##########
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"},
Review Comment:
Not testing `LEFT_ANTI` as well?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
