Skip to content

Commit 40ac4ea

Browse files
zjw1111zjw1111
committed
fix(testing): address review issues
1 parent 445d4e6 commit 40ac4ea

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"
@@ -283,7 +284,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestSimple) {
283284
/*enable_adaptive_prefetch_strategy=*/false, executor_,
284285
/*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS,
285286
CacheConfig(), GetDefaultPool()));
286-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
287+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
288+
reader->GetPreviousBatchFirstRowNumber().value());
287289
ASSERT_OK_AND_ASSIGN(auto result_array,
288290
ReadResultCollector::CollectResult(
289291
reader.get(), /*max simulated data processing time*/ 100));
@@ -605,7 +607,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestReadWithLargeBatchSize) {
605607
prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_,
606608
/*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS,
607609
CacheConfig(), GetDefaultPool()));
608-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
610+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
611+
reader->GetPreviousBatchFirstRowNumber().value());
609612
ASSERT_OK_AND_ASSIGN(auto result_array,
610613
ReadResultCollector::CollectResult(
611614
reader.get(), /*max simulated data processing time*/ 100));
@@ -633,7 +636,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPartialReaderSuccessRead) {
633636
}
634637

635638
arrow::ArrayVector result_array_vector;
636-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
639+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
640+
reader->GetPreviousBatchFirstRowNumber().value());
637641
ASSERT_OK_AND_ASSIGN(auto batch_with_bitmap, reader->NextBatchWithBitmap());
638642
auto& [batch, bitmap] = batch_with_bitmap;
639643
ASSERT_EQ(batch.first->length, bitmap.Cardinality());
@@ -678,9 +682,11 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestAllReaderFailedWithIOError) {
678682
->SetNextBatchStatus(Status::IOError("mock error"));
679683
}
680684

681-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
685+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
686+
reader->GetPreviousBatchFirstRowNumber().value());
682687
auto batch_result = reader->NextBatchWithBitmap();
683-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
688+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
689+
reader->GetPreviousBatchFirstRowNumber().value());
684690
ASSERT_FALSE(batch_result.ok());
685691
ASSERT_TRUE(batch_result.status().IsIOError());
686692
ASSERT_FALSE(prefetch_reader->is_shutdown_);
@@ -689,7 +695,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestAllReaderFailedWithIOError) {
689695

690696
// call NextBatch again, will still return error status
691697
auto batch_result2 = reader->NextBatchWithBitmap();
692-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
698+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
699+
reader->GetPreviousBatchFirstRowNumber().value());
693700
ASSERT_FALSE(batch_result2.ok());
694701
ASSERT_TRUE(batch_result2.status().IsIOError());
695702
}
@@ -706,7 +713,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestPrefetchWithEmptyData) {
706713
prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_,
707714
/*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS,
708715
CacheConfig(), GetDefaultPool()));
709-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
716+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
717+
reader->GetPreviousBatchFirstRowNumber().value());
710718
ASSERT_OK_AND_ASSIGN(auto result_array,
711719
ReadResultCollector::CollectResult(
712720
reader.get(), /*max simulated data processing time*/ 100));
@@ -726,7 +734,8 @@ TEST_F(PrefetchFileBatchReaderImplTest, TestCallNextBatchAfterReadingEof) {
726734
prefetch_max_parallel_num * 2, /*enable_adaptive_prefetch_strategy=*/false, executor_,
727735
/*initialize_read_ranges=*/true, /*prefetch_cache_mode=*/PrefetchCacheMode::ALWAYS,
728736
CacheConfig(), GetDefaultPool()));
729-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
737+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
738+
reader->GetPreviousBatchFirstRowNumber().value());
730739
ASSERT_OK_AND_ASSIGN(auto result_array,
731740
ReadResultCollector::CollectResult(
732741
reader.get(), /*max simulated data processing time*/ 100));
@@ -832,7 +841,8 @@ TEST_P(PrefetchFileBatchReaderImplTest, TestPrefetchWithPredicatePushdownWithCom
832841
PreparePrefetchReader(file_format, schema.get(), predicate,
833842
/*selection_bitmap=*/std::nullopt,
834843
/*batch_size=*/10, /*prefetch_max_parallel_num=*/3, cache_mode);
835-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
844+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
845+
reader->GetPreviousBatchFirstRowNumber().value());
836846
ASSERT_OK_AND_ASSIGN(auto result_array,
837847
ReadResultCollector::CollectResult(
838848
reader.get(), /*max simulated data processing time*/ 100));
@@ -868,7 +878,8 @@ TEST_P(PrefetchFileBatchReaderImplTest,
868878
/*selection_bitmap=*/std::nullopt,
869879
/*batch_size=*/10, /*prefetch_max_parallel_num=*/3, cache_mode);
870880
ASSERT_OK(reader->RefreshReadRanges());
871-
ASSERT_EQ(reader->GetPreviousBatchFirstRowNumber().value(), -1);
881+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
882+
reader->GetPreviousBatchFirstRowNumber().value());
872883
ASSERT_OK_AND_ASSIGN(auto result_array,
873884
ReadResultCollector::CollectResult(
874885
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>
@@ -492,7 +493,8 @@ TEST_P(OrcFileBatchReaderTest, TestNextBatchSimple) {
492493
for (auto batch_size : {1, 2, 3, 5, 8, 10}) {
493494
auto orc_batch_reader =
494495
PrepareOrcFileBatchReader(file_name, &read_schema, batch_size, natural_read_size);
495-
ASSERT_EQ(orc_batch_reader->GetPreviousBatchFirstRowNumber().value(), -1);
496+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
497+
orc_batch_reader->GetPreviousBatchFirstRowNumber().value());
496498
ASSERT_OK_AND_ASSIGN(auto result_array, paimon::test::ReadResultCollector::CollectResult(
497499
orc_batch_reader.get()));
498500
ASSERT_EQ(orc_batch_reader->GetPreviousBatchFirstRowNumber().value(), 8);
@@ -766,7 +768,8 @@ TEST_F(OrcFileBatchReaderTest, TestReadNoField) {
766768
auto orc_batch_reader = PrepareOrcFileBatchReader(file_name, &read_schema, /*batch_size=*/3,
767769
/*natural_read_size=*/10);
768770
// read 3 rows
769-
ASSERT_EQ(orc_batch_reader->GetPreviousBatchFirstRowNumber().value(), -1);
771+
ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
772+
orc_batch_reader->GetPreviousBatchFirstRowNumber().value());
770773
ASSERT_OK_AND_ASSIGN(auto batch1, orc_batch_reader->NextBatch());
771774
ASSERT_EQ(orc_batch_reader->GetPreviousBatchFirstRowNumber().value(), 0);
772775
// 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
@@ -18,7 +18,6 @@
1818

1919
#include <map>
2020
#include <string>
21-
#include <utility>
2221

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

3938
Status MockFormatWriter::AddBatch(ArrowArray* batch) {
4039
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)