zanmato1984 commented on code in PR #55170:
URL: https://github.com/apache/arrow/pull/55170#discussion_r4394039666
##########
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:
Added `ValidateFull()` before checking the group IDs.
##########
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:
Added `ValidateFull()` before comparing the lookup result with the expected
IDs.
##########
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:
Changed the scalar case to broadcast over three rows and check every length
counter, rather than only exercising a single row.
##########
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:
Changed the array case to three rows: a value larger than 2 GiB, an empty
value, and a null. It reuses the same value buffer and checks all three
counters, including accumulation on top of the scalar lengths.
##########
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:
With the default true filter, `LEFT_ANTI` has no right-side output, so the
large payload columns are projected away. Adding it to this test would
therefore not exercise the payload buffer exceeding 2 GiB. A meaningful
anti-join payload case would need a residual filter using those payload
columns. I kept this test focused on `INNER`; the separate large-key regression
still covers both `INNER` and `LEFT_ANTI`.
##########
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:
Added `ValidateFull()` on the complete decoded array before slicing it and
comparing the values.
##########
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:
Added `ValidateFull()`, length and null-count checks before accessing the
individual group IDs. The final lookup result is now validated as well.
##########
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:
Agreed, both types use `VarLengthKeyEncoder<LargeBinaryType>` here. I kept
only `large_binary()` for the payload regression, while retaining both
small-key encodings and serial/parallel execution.
##########
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:
It is the build-side encoder storing `r_payload1` and `r_payload2`, not the
join result. I clarified the comment to distinguish that buffer from the key
buffer.
##########
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:
`batch.length` is already `int64_t`; this caches the loop bound rather than
casting it. With Apple Clang on arm64, using the member directly caused the
prefix-sum loop to reload it on each iteration. A local A/B benchmark showed
about a 1% benefit from caching it for million-row fixed-width encoding
batches. I kept the local variable for that small benefit; it is not required
for correctness.
--
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]
