Skip to content

Commit 9793306

Browse files
zjw1111zjw1111
committed
fix(testing): address review issues
1 parent 737ddc7 commit 9793306

12 files changed

Lines changed: 72 additions & 37 deletions

‎src/paimon/common/reader/prefetch_file_batch_reader_impl_test.cpp‎

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include "paimon/common/reader/prefetch_file_batch_reader_impl.h"
1818

1919
#include <atomic>
20+
#include <limits>
2021
#include <set>
2122

2223
#include "arrow/compute/api.h"
@@ -281,7 +282,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestSimple) {
281282
/*enable_adaptive_prefetch_strategy=*/false, executor_,
282283
/*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS,
283284
CacheConfig(), GetDefaultPool()));
284-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
285+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
286+
reader->GetPreviousBatchFirstRowNumber().value());
285287
ASSERT_OK_AND_ASSIGN(auto result_array,
286288
ReadResultCollector::CollectResult(
287289
reader.get(), /*max simulated data processing time*/ 100));
@@ -603,7 +605,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLargeBatchSize) {
603605
prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_,
604606
/*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS,
605607
CacheConfig(), GetDefaultPool()));
606-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
608+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
609+
reader->GetPreviousBatchFirstRowNumber().value());
607610
ASSERT_OK_AND_ASSIGN(auto result_array,
608611
ReadResultCollector::CollectResult(
609612
reader.get(), /*max simulated data processing time*/ 100));
@@ -631,7 +634,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPartialReaderSuccessRead) {
631634
}
632635

633636
arrow::ArrayVector result_array_vector;
634-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
637+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
638+
reader->GetPreviousBatchFirstRowNumber().value());
635639
ASSERT_OK_AND_ASSIGN(auto batch_with_bitmap, reader->NextBatchWithBitmap());
636640
auto& [batch, bitmap] = batch_with_bitmap;
637641
ASSERT_EQ(batch.first->length, bitmap.Cardinality());
@@ -676,9 +680,11 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestAllReaderFailedWithIOError) {
676680
->SetNextBatchStatus(Status::IOError("mock error"));
677681
}
678682

679-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
683+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
684+
reader->GetPreviousBatchFirstRowNumber().value());
680685
auto batch_result = reader->NextBatchWithBitmap();
681-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
686+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
687+
reader->GetPreviousBatchFirstRowNumber().value());
682688
ASSERT_FALSE(batch_result.ok());
683689
ASSERT_TRUE(batch_result.status().IsIOError());
684690
ASSERT_FALSE(prefetch_reader->is_shutdown_);
@@ -687,7 +693,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestAllReaderFailedWithIOError) {
687693

688694
// call NextBatch again, will still return error status
689695
auto batch_result2 = reader->NextBatchWithBitmap();
690-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
696+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
697+
reader->GetPreviousBatchFirstRowNumber().value());
691698
ASSERT_FALSE(batch_result2.ok());
692699
ASSERT_TRUE(batch_result2.status().IsIOError());
693700
}
@@ -704,7 +711,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithEmptyData) {
704711
prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_,
705712
/*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS,
706713
CacheConfig(), GetDefaultPool()));
707-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
714+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
715+
reader->GetPreviousBatchFirstRowNumber().value());
708716
ASSERT_OK_AND_ASSIGN(auto result_array,
709717
ReadResultCollector::CollectResult(
710718
reader.get(), /*max simulated data processing time*/ 100));
@@ -724,7 +732,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCallNextBatchAfterReadingEof) {
724732
prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_,
725733
/*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS,
726734
CacheConfig(), GetDefaultPool()));
727-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
735+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
736+
reader->GetPreviousBatchFirstRowNumber().value());
728737
ASSERT_OK_AND_ASSIGN(auto result_array,
729738
ReadResultCollector::CollectResult(
730739
reader.get(), /*max simulated data processing time*/ 100));
@@ -830,7 +839,8 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestPrefetchWithPredicatePushdownWithCom
830839
PreparePrefetchReader(file_format, schema.get(), predicate,
831840
/*selection_bitmap=*/std::nullopt,
832841
/*batch_size=*/10, /*prefetch_max_parallel_num=*/3, cache_mode);
833-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
842+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
843+
reader->GetPreviousBatchFirstRowNumber().value());
834844
ASSERT_OK_AND_ASSIGN(auto result_array,
835845
ReadResultCollector::CollectResult(
836846
reader.get(), /*max simulated data processing time*/ 100));
@@ -866,7 +876,8 @@ TEST_P(PrefetchFileBatchReaderImplTest,
866876
/*selection_bitmap=*/std::nullopt,
867877
/*batch_size=*/10, /*prefetch_max_parallel_num=*/3, cache_mode);
868878
ASSERT_OK(reader->RefreshReadRanges());
869-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
879+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
880+
reader->GetPreviousBatchFirstRowNumber().value());
870881
ASSERT_OK_AND_ASSIGN(auto result_array,
871882
ReadResultCollector::CollectResult(
872883
reader.get(), /*max simulated data processing time*/ 100));

‎src/paimon/common/utils/arrow/arrow_utils.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "arrow/array/array_base.h"
2020
#include "arrow/array/array_nested.h"
2121
#include "arrow/util/compression.h"
22+
#include "fmt/format.h"
2223
#include "paimon/common/utils/arrow/status_utils.h"
2324
#include "paimon/common/utils/string_utils.h"
2425

‎src/paimon/common/utils/arrow/arrow_utils.h‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020

2121
#include "arrow/api.h"
2222
#include "arrow/util/type_fwd.h"
23-
#include "fmt/format.h"
2423
#include "paimon/result.h"
2524

2625
namespace paimon {

‎src/paimon/common/utils/arrow/arrow_utils_test.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include "gtest/gtest.h"
2222
#include "paimon/common/types/data_field.h"
2323
#include "paimon/testing/utils/testharness.h"
24+
2425
namespace paimon::test {
2526

2627
TEST(ArrowUtilsTest, TestCreateProjection) {

‎src/paimon/common/utils/arrow/mem_utils.cpp‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
#include "arrow/memory_pool.h"
2424
#include "arrow/status.h"
25+
#include "fmt/format.h"
2526
#include "paimon/memory/memory_pool.h"
2627

2728
namespace paimon {
@@ -33,13 +34,22 @@ class ArrowMemPoolAdaptor : public arrow::MemoryPool {
3334

3435
arrow::Status Allocate(int64_t size, int64_t alignment, uint8_t** out) override {
3536
*out = reinterpret_cast<uint8_t*>(pool_.Malloc(size, alignment));
37+
if (size > 0 && *out == nullptr) {
38+
return arrow::Status::OutOfMemory(fmt::format("failed to allocate {} bytes", size));
39+
}
3640
stats_.DidAllocateBytes(size);
3741
return arrow::Status::OK();
3842
}
3943

4044
arrow::Status Reallocate(int64_t old_size, int64_t new_size, int64_t alignment,
4145
uint8_t** ptr) override {
42-
*ptr = reinterpret_cast<uint8_t*>(pool_.Realloc(*ptr, old_size, new_size, alignment));
46+
auto* new_ptr =
47+
reinterpret_cast<uint8_t*>(pool_.Realloc(*ptr, old_size, new_size, alignment));
48+
if (new_size > 0 && new_ptr == nullptr) {
49+
return arrow::Status::OutOfMemory(
50+
fmt::format("failed to reallocate memory from {} to {} bytes", old_size, new_size));
51+
}
52+
*ptr = new_ptr;
4353
stats_.DidReallocateBytes(old_size, new_size);
4454
return arrow::Status::OK();
4555
}

‎src/paimon/common/utils/arrow/mem_utils.h‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,6 @@
2323
#include "paimon/visibility.h"
2424

2525
namespace paimon {
26-
class MemoryPool;
27-
28-
PAIMON_EXPORT std::unique_ptr<arrow::MemoryPool> GetArrowPool(MemoryPool& pool);
2926

3027
PAIMON_EXPORT std::unique_ptr<arrow::MemoryPool> GetArrowPool(
3128
const std::shared_ptr<MemoryPool>& pool);

‎src/paimon/format/orc/orc_file_batch_reader_test.cpp‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
#include "paimon/format/orc/orc_file_batch_reader.h"
1818

19+
#include <limits>
1920
#include <list>
2021
#include <map>
2122
#include <memory>
@@ -400,7 +401,8 @@ TEST_P(OrcFileBatchReaderTest, TestNextBatchSimple) {
400401
for (auto batch_size : {1, 2, 3, 5, 8, 10}) {
401402
auto orc_batch_reader =
402403
PrepareOrcFileBatchReader(file_name, &read_schema, batch_size, natural_read_size);
403-
ASSERT_EQ(orc_batch_reader->GetPreviousBatchFirstRowNumber().value(), -1);
404+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
405+
orc_batch_reader->GetPreviousBatchFirstRowNumber().value());
404406
ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult(
405407
orc_batch_reader.get()));
406408
ASSERT_EQ(orc_batch_reader->GetPreviousBatchFirstRowNumber().value(), 8);
@@ -674,7 +676,8 @@ TEST_F(OrcFileBatchReaderTest, TestReadNoField) {
674676
auto orc_batch_reader = PrepareOrcFileBatchReader(file_name, &read_schema, /*batch_size=*/3,
675677
/*natural_read_size=*/10);
676678
// read 3 rows
677-
ASSERT_EQ(orc_batch_reader->GetPreviousBatchFirstRowNumber().value(), -1);
679+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
680+
orc_batch_reader->GetPreviousBatchFirstRowNumber().value());
678681
ASSERT_OK_AND_ASSIGN(auto batch1, orc_batch_reader->NextBatch());
679682
ASSERT_EQ(orc_batch_reader->GetPreviousBatchFirstRowNumber().value(), 0);
680683
// read 3 rows

‎src/paimon/testing/mock/mock_file_batch_reader.h‎

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,11 @@
1717
#pragma once
1818

1919
#include <algorithm>
20+
#include <cassert>
21+
#include <cstdint>
22+
#include <limits>
2023
#include <memory>
24+
#include <random>
2125
#include <utility>
2226
#include <vector>
2327

@@ -26,8 +30,8 @@
2630
#include "paimon/common/metrics/metrics_impl.h"
2731
#include "paimon/common/reader/reader_utils.h"
2832
#include "paimon/common/utils/arrow/status_utils.h"
29-
#include "paimon/common/utils/date_time_utils.h"
3033
#include "paimon/reader/prefetch_file_batch_reader.h"
34+
3135
namespace paimon::test {
3236

3337
class MockFileBatchReader : public PrefetchFileBatchReader {
@@ -39,13 +43,14 @@ class MockFileBatchReader : public PrefetchFileBatchReader {
3943
file_schema_(file_schema),
4044
read_schema_(arrow::schema(file_schema->fields())),
4145
batch_size_(read_batch_size) {
46+
assert(read_batch_size > 0);
4247
// add all valid bitmap
43-
int32_t data_length = data_ ? data_->length() : 0;
48+
int64_t data_length = data_ ? data_->length() : 0;
49+
assert(data_length >= 0);
50+
assert(data_length <= static_cast<int64_t>(std::numeric_limits<int32_t>::max()));
4451
bitmap_ = RoaringBitmap32();
45-
bitmap_.AddRange(0, data_length);
46-
read_end_pos_ = data_length;
47-
int64_t seed = DateTimeUtils::GetCurrentUTCTimeUs();
48-
std::srand(seed);
52+
bitmap_.AddRange(0, static_cast<int32_t>(data_length));
53+
read_end_pos_ = static_cast<int32_t>(data_length);
4954
}
5055

5156
MockFileBatchReader(const std::shared_ptr<arrow::Array>& data,
@@ -93,7 +98,8 @@ class MockFileBatchReader : public PrefetchFileBatchReader {
9398
}
9499

95100
Status SeekToRow(uint64_t row_number) override {
96-
current_pos_ = row_number;
101+
assert(row_number <= static_cast<uint64_t>(std::numeric_limits<int32_t>::max()));
102+
current_pos_ = static_cast<int32_t>(row_number);
97103
return Status::OK();
98104
}
99105

@@ -117,7 +123,8 @@ class MockFileBatchReader : public PrefetchFileBatchReader {
117123
}
118124
int32_t actual_batch_size = batch_size_;
119125
if (enable_randomize_batch_size_) {
120-
actual_batch_size = std::rand() % batch_size_ + 1;
126+
std::uniform_int_distribution<int32_t> distribution(1, batch_size_);
127+
actual_batch_size = distribution(random_engine_);
121128
}
122129
int32_t batch_end_pos = std::min(read_end_pos_, current_pos_ + actual_batch_size);
123130
auto slice = data_->Slice(current_pos_, batch_end_pos - current_pos_);
@@ -150,14 +157,14 @@ class MockFileBatchReader : public PrefetchFileBatchReader {
150157
}
151158

152159
Result<uint64_t> GetPreviousBatchFirstRowNumber() const override {
153-
return previous_batch_first_row_num_;
160+
return ToReaderRowNumber(previous_batch_first_row_num_);
154161
}
155162

156163
Result<uint64_t> GetNumberOfRows() const override {
157-
return data_ ? data_->length() : 0;
164+
return ToReaderRowNumber(read_end_pos_);
158165
}
159166
uint64_t GetNextRowToRead() const override {
160-
return current_pos_;
167+
return ToReaderRowNumber(current_pos_);
161168
}
162169
void Close() override {}
163170

@@ -170,6 +177,13 @@ class MockFileBatchReader : public PrefetchFileBatchReader {
170177
}
171178

172179
private:
180+
static uint64_t ToReaderRowNumber(int32_t row_number) {
181+
if (row_number < 0) {
182+
return std::numeric_limits<uint64_t>::max();
183+
}
184+
return static_cast<uint64_t>(row_number);
185+
}
186+
173187
std::shared_ptr<arrow::Array> data_;
174188
std::shared_ptr<arrow::DataType> file_schema_;
175189
std::shared_ptr<arrow::Schema> read_schema_;
@@ -181,6 +195,7 @@ class MockFileBatchReader : public PrefetchFileBatchReader {
181195
Status next_batch_status_;
182196
bool enable_randomize_batch_size_ = true;
183197
std::vector<std::pair<uint64_t, uint64_t>> read_ranges_;
198+
std::mt19937 random_engine_{std::random_device{}()};
184199
};
185200

186201
} // namespace paimon::test

‎src/paimon/testing/mock/mock_format_writer.cpp‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
#include "paimon/testing/mock/mock_format_writer.h"
1818

1919
#include <string>
20-
#include <utility>
2120

2221
#include "arrow/c/helpers.h"
2322
#include "paimon/common/utils/date_time_utils.h"
@@ -33,7 +32,7 @@ class MemoryPool;
3332
namespace paimon::test {
3433
MockFormatWriter::MockFormatWriter(const std::shared_ptr<OutputStream>& out,
3534
const std::shared_ptr<MemoryPool>& pool)
36-
: FormatWriter(), out_(std::move(out)), pool_(pool) {}
35+
: FormatWriter(), out_(out), pool_(pool) {}
3736

3837
Status MockFormatWriter::AddBatch(ArrowArray* batch) {
3938
ArrowArrayRelease(batch);

‎src/paimon/testing/mock/mock_format_writer_builder.h‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,6 @@ class MockFormatWriterBuilder : public WriterBuilder {
4242
const std::string& compression) override;
4343

4444
private:
45-
Status Prepare();
46-
4745
std::shared_ptr<MemoryPool> memory_pool_;
4846
};
4947

0 commit comments

Comments
 (0)