Skip to content

Commit ecfff72

Browse files
lxy-9602lxy-9602
authored
perf(blob): optimize array & map blob reading (#395)
1 parent 7e18577 commit ecfff72

5 files changed

Lines changed: 628 additions & 285 deletions

File tree

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

Lines changed: 73 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,9 @@
2323
#include "arrow/array/util.h"
2424
#include "arrow/c/abi.h"
2525
#include "arrow/c/bridge.h"
26-
#include "fmt/format.h"
2726
#include "paimon/common/metrics/metrics_impl.h"
2827
#include "paimon/common/reader/reader_utils.h"
28+
#include "paimon/common/utils/arrow/arrow_utils.h"
2929
#include "paimon/common/utils/arrow/mem_utils.h"
3030
#include "paimon/common/utils/arrow/status_utils.h"
3131
#include "paimon/common/utils/checked_cast.h"
@@ -59,25 +59,43 @@ Result<std::unique_ptr<DataEvolutionFileReader>> DataEvolutionFileReader::Create
5959
}
6060

6161
Result<BatchReader::ReadBatchWithBitmap> DataEvolutionFileReader::NextBatchWithBitmap() {
62+
int64_t array_length = read_batch_size_;
63+
bool has_active_reader = false;
64+
bool has_available_array = false;
65+
bool has_eof_reader = false;
66+
for (size_t i = 0; i < readers_.size(); i++) {
67+
if (!readers_[i]) {
68+
continue;
69+
}
70+
has_active_reader = true;
71+
PAIMON_ASSIGN_OR_RAISE(bool has_cached_array, EnsureCachedArray(i));
72+
if (!has_cached_array) {
73+
has_eof_reader = true;
74+
continue;
75+
}
76+
has_available_array = true;
77+
array_length = std::min(array_length, CalculateCachedArrayLength(i));
78+
}
79+
if (!has_active_reader) {
80+
return Status::Invalid("data evolution reader has no active inner reader");
81+
}
82+
if (has_eof_reader) {
83+
if (has_available_array) {
84+
return Status::Invalid("array for single reader length mismatch others");
85+
}
86+
return BatchReader::MakeEofBatchWithBitmap();
87+
}
88+
6289
std::vector<std::shared_ptr<arrow::StructArray>> array_for_each_reader;
6390
array_for_each_reader.reserve(readers_.size());
64-
int64_t array_length = -1;
6591
for (size_t i = 0; i < readers_.size(); i++) {
6692
if (!readers_[i]) {
6793
// no read field from readers_[i]
6894
array_for_each_reader.push_back(nullptr);
6995
continue;
7096
}
71-
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> array, NextBatchForSingleReader(i));
72-
if (array == nullptr) {
73-
// read eof
74-
return BatchReader::MakeEofBatchWithBitmap();
75-
}
76-
if (array_length == -1) {
77-
array_length = array->length();
78-
} else if (array_length != array->length()) {
79-
return Status::Invalid("array for single reader length mismatch others");
80-
}
97+
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> array,
98+
TakeCachedArray(i, array_length));
8199
auto struct_array = checked_pointer_cast<arrow::StructArray>(array);
82100
array_for_each_reader.push_back(struct_array);
83101
}
@@ -132,24 +150,19 @@ int64_t DataEvolutionFileReader::CalculateCachedArrayLength(size_t reader_idx) c
132150
return total_length;
133151
}
134152

135-
Result<std::shared_ptr<arrow::Array>> DataEvolutionFileReader::NextBatchForSingleReader(
136-
size_t reader_idx) {
137-
int64_t total_array_length = CalculateCachedArrayLength(reader_idx);
138-
if (total_array_length >= read_batch_size_) {
139-
assert(false);
140-
return Status::Invalid(fmt::format(
141-
"Unexpected: the length of cached array in last turn {} exceed read batch size {}",
142-
total_array_length, read_batch_size_));
143-
}
144-
// array left for last turn
145-
arrow::ArrayVector concat_array_vec = std::move(cached_array_vec_[reader_idx]);
146-
cached_array_vec_[reader_idx].clear();
147-
while (total_array_length < read_batch_size_) {
153+
Result<bool> DataEvolutionFileReader::EnsureCachedArray(size_t reader_idx) {
154+
if (!cached_array_vec_[reader_idx].empty()) {
155+
return true;
156+
}
157+
if (reader_eof_[reader_idx]) {
158+
return false;
159+
}
160+
while (cached_array_vec_[reader_idx].empty()) {
148161
PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap src_array_with_bitmap,
149162
readers_[reader_idx]->NextBatchWithBitmap());
150163
if (BatchReader::IsEofBatch(src_array_with_bitmap)) {
151-
// read finish
152-
break;
164+
reader_eof_[reader_idx] = true;
165+
return false;
153166
}
154167
auto& [read_batch, bitmap] = src_array_with_bitmap;
155168
auto& [c_array, c_schema] = read_batch;
@@ -158,36 +171,45 @@ Result<std::shared_ptr<arrow::Array>> DataEvolutionFileReader::NextBatchForSingl
158171
PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector selected_array_vec,
159172
ReaderUtils::GenerateFilteredArrayVector(src_array, bitmap));
160173
for (const auto& selected_array : selected_array_vec) {
161-
if (total_array_length + selected_array->length() > read_batch_size_) {
162-
// need truncate current array to align read_batch_size_
163-
int64_t truncated_length = read_batch_size_ - total_array_length;
164-
if (truncated_length == 0) {
165-
// total_array_length equals to read_batch_size_, all selected_array left will
166-
// be added to cached_array_vec_
167-
cached_array_vec_[reader_idx].push_back(selected_array);
168-
} else {
169-
concat_array_vec.push_back(selected_array->Slice(0, truncated_length));
170-
cached_array_vec_[reader_idx].push_back(
171-
selected_array->Slice(truncated_length));
172-
total_array_length += truncated_length;
173-
}
174-
} else {
175-
concat_array_vec.push_back(selected_array);
176-
total_array_length += selected_array->length();
174+
if (selected_array->length() > 0) {
175+
cached_array_vec_[reader_idx].push_back(selected_array);
177176
}
178177
}
179178
}
180-
if (concat_array_vec.empty()) {
181-
return std::shared_ptr<arrow::Array>();
179+
return true;
180+
}
181+
182+
Result<std::shared_ptr<arrow::Array>> DataEvolutionFileReader::TakeCachedArray(
183+
size_t reader_idx, int64_t array_length) {
184+
if (array_length <= 0) {
185+
return Status::Invalid("array length must be larger than zero");
186+
}
187+
if (CalculateCachedArrayLength(reader_idx) < array_length) {
188+
return Status::Invalid("requested array length exceeds cached array length");
182189
}
183-
if (concat_array_vec.size() == 1 && concat_array_vec[0]->offset() == 0) {
184-
// Avoid data copy when the array is already normalized.
185-
return concat_array_vec[0];
190+
arrow::ArrayVector selected_array_vec;
191+
int64_t remaining_length = array_length;
192+
auto& cached_array_vec = cached_array_vec_[reader_idx];
193+
while (remaining_length > 0) {
194+
const auto& cached_array = cached_array_vec.front();
195+
if (cached_array->length() <= remaining_length) {
196+
selected_array_vec.push_back(cached_array);
197+
remaining_length -= cached_array->length();
198+
cached_array_vec.erase(cached_array_vec.begin());
199+
} else {
200+
selected_array_vec.push_back(cached_array->Slice(0, remaining_length));
201+
cached_array_vec.front() = cached_array->Slice(remaining_length);
202+
remaining_length = 0;
203+
}
204+
}
205+
if (selected_array_vec.size() == 1) {
206+
return ArrowUtils::NormalizeArrayOffsets(selected_array_vec[0], arrow_pool_.get());
186207
}
187208
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> concat_array,
188-
arrow::Concatenate(concat_array_vec, arrow_pool_.get()));
189-
assert(concat_array->length() == total_array_length);
190-
assert(concat_array->length() <= read_batch_size_);
209+
arrow::Concatenate(selected_array_vec, arrow_pool_.get()));
210+
if (concat_array->length() != array_length) {
211+
return Status::Invalid("concatenated array length does not match requested array length");
212+
}
191213
return concat_array;
192214
}
193215

‎src/paimon/common/reader/data_evolution_file_reader.h‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,11 +85,14 @@ class DataEvolutionFileReader : public BatchReader {
8585
reader_offsets_(reader_offsets),
8686
field_offsets_(field_offsets),
8787
cached_array_vec_(readers_.size()),
88+
reader_eof_(readers_.size(), false),
8889
non_exist_array_vec_(read_schema->num_fields(), nullptr) {}
8990

9091
int64_t CalculateCachedArrayLength(size_t reader_idx) const;
9192

92-
Result<std::shared_ptr<arrow::Array>> NextBatchForSingleReader(size_t reader_idx);
93+
Result<bool> EnsureCachedArray(size_t reader_idx);
94+
95+
Result<std::shared_ptr<arrow::Array>> TakeCachedArray(size_t reader_idx, int64_t array_length);
9396

9497
Result<std::shared_ptr<arrow::Array>> GetOrCreateNonExistArray(int32_t field_idx,
9598
int64_t array_length);
@@ -102,6 +105,7 @@ class DataEvolutionFileReader : public BatchReader {
102105
std::vector<int32_t> reader_offsets_;
103106
std::vector<int32_t> field_offsets_;
104107
std::vector<arrow::ArrayVector> cached_array_vec_;
108+
std::vector<bool> reader_eof_;
105109
arrow::ArrayVector non_exist_array_vec_;
106110
};
107111
} // namespace paimon

0 commit comments

Comments
 (0)