This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-cpp.git
The following commit(s) were added to refs/heads/main by this push:
new ecfff724 perf(blob): optimize array & map blob reading (#395)
ecfff724 is described below
commit ecfff724448cd88badbe06b8125378b17c9d511d
Author: lxy <[email protected]>
AuthorDate: Tue Sep 29 13:49:48 2026 +0800
perf(blob): optimize array & map blob reading (#395)
---
.../common/reader/data_evolution_file_reader.cpp | 124 ++---
.../common/reader/data_evolution_file_reader.h | 6 +-
.../reader/data_evolution_file_reader_test.cpp | 273 +++++------
src/paimon/format/blob/blob_file_batch_reader.cpp | 499 ++++++++++++++++++---
src/paimon/format/blob/blob_file_batch_reader.h | 11 +-
5 files changed, 628 insertions(+), 285 deletions(-)
diff --git a/src/paimon/common/reader/data_evolution_file_reader.cpp
b/src/paimon/common/reader/data_evolution_file_reader.cpp
index 347fe439..5cace418 100644
--- a/src/paimon/common/reader/data_evolution_file_reader.cpp
+++ b/src/paimon/common/reader/data_evolution_file_reader.cpp
@@ -23,9 +23,9 @@
#include "arrow/array/util.h"
#include "arrow/c/abi.h"
#include "arrow/c/bridge.h"
-#include "fmt/format.h"
#include "paimon/common/metrics/metrics_impl.h"
#include "paimon/common/reader/reader_utils.h"
+#include "paimon/common/utils/arrow/arrow_utils.h"
#include "paimon/common/utils/arrow/mem_utils.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/common/utils/checked_cast.h"
@@ -59,25 +59,43 @@ Result<std::unique_ptr<DataEvolutionFileReader>>
DataEvolutionFileReader::Create
}
Result<BatchReader::ReadBatchWithBitmap>
DataEvolutionFileReader::NextBatchWithBitmap() {
+ int64_t array_length = read_batch_size_;
+ bool has_active_reader = false;
+ bool has_available_array = false;
+ bool has_eof_reader = false;
+ for (size_t i = 0; i < readers_.size(); i++) {
+ if (!readers_[i]) {
+ continue;
+ }
+ has_active_reader = true;
+ PAIMON_ASSIGN_OR_RAISE(bool has_cached_array, EnsureCachedArray(i));
+ if (!has_cached_array) {
+ has_eof_reader = true;
+ continue;
+ }
+ has_available_array = true;
+ array_length = std::min(array_length, CalculateCachedArrayLength(i));
+ }
+ if (!has_active_reader) {
+ return Status::Invalid("data evolution reader has no active inner
reader");
+ }
+ if (has_eof_reader) {
+ if (has_available_array) {
+ return Status::Invalid("array for single reader length mismatch
others");
+ }
+ return BatchReader::MakeEofBatchWithBitmap();
+ }
+
std::vector<std::shared_ptr<arrow::StructArray>> array_for_each_reader;
array_for_each_reader.reserve(readers_.size());
- int64_t array_length = -1;
for (size_t i = 0; i < readers_.size(); i++) {
if (!readers_[i]) {
// no read field from readers_[i]
array_for_each_reader.push_back(nullptr);
continue;
}
- PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> array,
NextBatchForSingleReader(i));
- if (array == nullptr) {
- // read eof
- return BatchReader::MakeEofBatchWithBitmap();
- }
- if (array_length == -1) {
- array_length = array->length();
- } else if (array_length != array->length()) {
- return Status::Invalid("array for single reader length mismatch
others");
- }
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> array,
+ TakeCachedArray(i, array_length));
auto struct_array = checked_pointer_cast<arrow::StructArray>(array);
array_for_each_reader.push_back(struct_array);
}
@@ -132,24 +150,19 @@ int64_t
DataEvolutionFileReader::CalculateCachedArrayLength(size_t reader_idx) c
return total_length;
}
-Result<std::shared_ptr<arrow::Array>>
DataEvolutionFileReader::NextBatchForSingleReader(
- size_t reader_idx) {
- int64_t total_array_length = CalculateCachedArrayLength(reader_idx);
- if (total_array_length >= read_batch_size_) {
- assert(false);
- return Status::Invalid(fmt::format(
- "Unexpected: the length of cached array in last turn {} exceed
read batch size {}",
- total_array_length, read_batch_size_));
- }
- // array left for last turn
- arrow::ArrayVector concat_array_vec =
std::move(cached_array_vec_[reader_idx]);
- cached_array_vec_[reader_idx].clear();
- while (total_array_length < read_batch_size_) {
+Result<bool> DataEvolutionFileReader::EnsureCachedArray(size_t reader_idx) {
+ if (!cached_array_vec_[reader_idx].empty()) {
+ return true;
+ }
+ if (reader_eof_[reader_idx]) {
+ return false;
+ }
+ while (cached_array_vec_[reader_idx].empty()) {
PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap src_array_with_bitmap,
readers_[reader_idx]->NextBatchWithBitmap());
if (BatchReader::IsEofBatch(src_array_with_bitmap)) {
- // read finish
- break;
+ reader_eof_[reader_idx] = true;
+ return false;
}
auto& [read_batch, bitmap] = src_array_with_bitmap;
auto& [c_array, c_schema] = read_batch;
@@ -158,36 +171,45 @@ Result<std::shared_ptr<arrow::Array>>
DataEvolutionFileReader::NextBatchForSingl
PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector selected_array_vec,
ReaderUtils::GenerateFilteredArrayVector(src_array, bitmap));
for (const auto& selected_array : selected_array_vec) {
- if (total_array_length + selected_array->length() >
read_batch_size_) {
- // need truncate current array to align read_batch_size_
- int64_t truncated_length = read_batch_size_ -
total_array_length;
- if (truncated_length == 0) {
- // total_array_length equals to read_batch_size_, all
selected_array left will
- // be added to cached_array_vec_
- cached_array_vec_[reader_idx].push_back(selected_array);
- } else {
- concat_array_vec.push_back(selected_array->Slice(0,
truncated_length));
- cached_array_vec_[reader_idx].push_back(
- selected_array->Slice(truncated_length));
- total_array_length += truncated_length;
- }
- } else {
- concat_array_vec.push_back(selected_array);
- total_array_length += selected_array->length();
+ if (selected_array->length() > 0) {
+ cached_array_vec_[reader_idx].push_back(selected_array);
}
}
}
- if (concat_array_vec.empty()) {
- return std::shared_ptr<arrow::Array>();
+ return true;
+}
+
+Result<std::shared_ptr<arrow::Array>> DataEvolutionFileReader::TakeCachedArray(
+ size_t reader_idx, int64_t array_length) {
+ if (array_length <= 0) {
+ return Status::Invalid("array length must be larger than zero");
+ }
+ if (CalculateCachedArrayLength(reader_idx) < array_length) {
+ return Status::Invalid("requested array length exceeds cached array
length");
}
- if (concat_array_vec.size() == 1 && concat_array_vec[0]->offset() == 0) {
- // Avoid data copy when the array is already normalized.
- return concat_array_vec[0];
+ arrow::ArrayVector selected_array_vec;
+ int64_t remaining_length = array_length;
+ auto& cached_array_vec = cached_array_vec_[reader_idx];
+ while (remaining_length > 0) {
+ const auto& cached_array = cached_array_vec.front();
+ if (cached_array->length() <= remaining_length) {
+ selected_array_vec.push_back(cached_array);
+ remaining_length -= cached_array->length();
+ cached_array_vec.erase(cached_array_vec.begin());
+ } else {
+ selected_array_vec.push_back(cached_array->Slice(0,
remaining_length));
+ cached_array_vec.front() = cached_array->Slice(remaining_length);
+ remaining_length = 0;
+ }
+ }
+ if (selected_array_vec.size() == 1) {
+ return ArrowUtils::NormalizeArrayOffsets(selected_array_vec[0],
arrow_pool_.get());
}
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array>
concat_array,
- arrow::Concatenate(concat_array_vec,
arrow_pool_.get()));
- assert(concat_array->length() == total_array_length);
- assert(concat_array->length() <= read_batch_size_);
+ arrow::Concatenate(selected_array_vec,
arrow_pool_.get()));
+ if (concat_array->length() != array_length) {
+ return Status::Invalid("concatenated array length does not match
requested array length");
+ }
return concat_array;
}
diff --git a/src/paimon/common/reader/data_evolution_file_reader.h
b/src/paimon/common/reader/data_evolution_file_reader.h
index 31743d1e..7f6961d8 100644
--- a/src/paimon/common/reader/data_evolution_file_reader.h
+++ b/src/paimon/common/reader/data_evolution_file_reader.h
@@ -85,11 +85,14 @@ class DataEvolutionFileReader : public BatchReader {
reader_offsets_(reader_offsets),
field_offsets_(field_offsets),
cached_array_vec_(readers_.size()),
+ reader_eof_(readers_.size(), false),
non_exist_array_vec_(read_schema->num_fields(), nullptr) {}
int64_t CalculateCachedArrayLength(size_t reader_idx) const;
- Result<std::shared_ptr<arrow::Array>> NextBatchForSingleReader(size_t
reader_idx);
+ Result<bool> EnsureCachedArray(size_t reader_idx);
+
+ Result<std::shared_ptr<arrow::Array>> TakeCachedArray(size_t reader_idx,
int64_t array_length);
Result<std::shared_ptr<arrow::Array>> GetOrCreateNonExistArray(int32_t
field_idx,
int64_t
array_length);
@@ -102,6 +105,7 @@ class DataEvolutionFileReader : public BatchReader {
std::vector<int32_t> reader_offsets_;
std::vector<int32_t> field_offsets_;
std::vector<arrow::ArrayVector> cached_array_vec_;
+ std::vector<bool> reader_eof_;
arrow::ArrayVector non_exist_array_vec_;
};
} // namespace paimon
diff --git a/src/paimon/common/reader/data_evolution_file_reader_test.cpp
b/src/paimon/common/reader/data_evolution_file_reader_test.cpp
index 35cdda40..3a98d42a 100644
--- a/src/paimon/common/reader/data_evolution_file_reader_test.cpp
+++ b/src/paimon/common/reader/data_evolution_file_reader_test.cpp
@@ -93,50 +93,20 @@ class DataEvolutionFileReaderTest : public ::testing::Test,
}
}
- void CheckNextBatchForSingleReader(int32_t inner_batch_size, int32_t
read_batch_size,
- const std::shared_ptr<arrow::Array>&
src_array,
- const std::optional<RoaringBitmap32>&
selection_bitmap,
- const std::shared_ptr<arrow::Array>&
expected_array) const {
- std::unique_ptr<MockFileBatchReader> file_batch_reader;
- if (selection_bitmap) {
- file_batch_reader = std::make_unique<MockFileBatchReader>(
- src_array, src_array->type(), selection_bitmap.value(),
inner_batch_size);
- } else {
- file_batch_reader =
std::make_unique<MockFileBatchReader>(src_array, src_array->type(),
-
inner_batch_size);
+ Result<std::shared_ptr<arrow::Array>>
ReadNextOutput(DataEvolutionFileReader* reader) const {
+ PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap
batch_with_bitmap,
+ reader->NextBatchWithBitmap());
+ if (BatchReader::IsEofBatch(batch_with_bitmap)) {
+ return std::shared_ptr<arrow::Array>();
}
- auto enable_randomize_batch_size = GetParam();
-
file_batch_reader->EnableRandomizeBatchSize(enable_randomize_batch_size);
- std::vector<std::unique_ptr<BatchReader>> readers;
- readers.push_back(std::move(file_batch_reader));
- DataEvolutionFileReader fake_data_evolution_reader(
- std::move(readers), /*read_schema=*/arrow::schema({}),
read_batch_size,
- /*reader_offsets=*/{}, /*field_offsets=*/{}, GetArrowPool(pool_));
- arrow::ArrayVector result_array_vec;
- while (true) {
- ASSERT_OK_AND_ASSIGN(auto result_array,
-
fake_data_evolution_reader.NextBatchForSingleReader(0));
- if (result_array == nullptr) {
- break;
- }
- ASSERT_EQ(result_array->offset(), 0);
- result_array_vec.push_back(result_array);
- }
- ASSERT_EQ(result_array_vec.size(),
- std::ceil(static_cast<double>(expected_array->length()) /
read_batch_size));
- // except for last batch, the length each array is expected to be
aligned to read_batch_size
- for (size_t i = 0; i < result_array_vec.size() - 1; i++) {
- ASSERT_EQ(result_array_vec[i]->length(), read_batch_size);
- }
- if (expected_array->length() % read_batch_size == 0) {
- ASSERT_EQ(result_array_vec.back()->length(), read_batch_size);
- } else {
- ASSERT_EQ(result_array_vec.back()->length(),
- expected_array->length() % read_batch_size);
+ auto& [batch, bitmap] = batch_with_bitmap;
+ auto& [c_array, c_schema] = batch;
+ if (bitmap.Cardinality() != c_array->length) {
+ return Status::Invalid("data evolution output bitmap should select
every row");
}
- auto result_chunk_array =
std::make_shared<arrow::ChunkedArray>(result_array_vec);
- auto expected_chunk_array =
std::make_shared<arrow::ChunkedArray>(expected_array);
- ASSERT_TRUE(result_chunk_array->Equals(expected_chunk_array));
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> array,
+ arrow::ImportArray(c_array.get(),
c_schema.get()));
+ return array;
}
private:
@@ -201,130 +171,103 @@ TEST_F(DataEvolutionFileReaderTest, TestInvalid) {
}
}
-TEST_P(DataEvolutionFileReaderTest, TestNextBatchForSingleReader) {
- auto prepare_array = [](int64_t array_length) ->
std::shared_ptr<arrow::Array> {
- auto array_builder = std::make_shared<arrow::Int32Builder>();
- for (int32_t i = 0; i < array_length; ++i) {
- EXPECT_TRUE(array_builder->Append(i).ok());
- }
- std::shared_ptr<arrow::Array> array;
- EXPECT_TRUE(array_builder->Finish(&array).ok());
- return array;
- };
- auto prepare_array_with_bitmap =
- [](const RoaringBitmap32& bitmap) -> std::shared_ptr<arrow::Array> {
- auto array_builder = std::make_shared<arrow::Int32Builder>();
- for (auto iter = bitmap.Begin(); iter != bitmap.End(); ++iter) {
- EXPECT_TRUE(array_builder->Append(*iter).ok());
- }
- std::shared_ptr<arrow::Array> array;
- EXPECT_TRUE(array_builder->Finish(&array).ok());
- return array;
- };
- {
- // src array length = 10, read batch size = 10
- auto src_array = prepare_array(10);
- for (int32_t inner_batch_size : arrow::internal::Iota(1, 10)) {
- CheckNextBatchForSingleReader(inner_batch_size,
/*read_batch_size=*/10, src_array,
- /*selection_bitmap=*/std::nullopt,
- /*expected_array=*/src_array);
- }
- }
- {
- // src array length = 10, read batch size = 6
- auto src_array = prepare_array(10);
- for (int32_t inner_batch_size : arrow::internal::Iota(1, 6)) {
- CheckNextBatchForSingleReader(inner_batch_size,
/*read_batch_size=*/6, src_array,
- /*selection_bitmap=*/std::nullopt,
- /*expected_array=*/src_array);
- }
- }
- {
- // src array length = 10, read batch size = 15
- auto src_array = prepare_array(10);
- for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
- CheckNextBatchForSingleReader(inner_batch_size,
/*read_batch_size=*/15, src_array,
- /*selection_bitmap=*/std::nullopt,
- /*expected_array=*/src_array);
- }
- }
- {
- // test bulk data, src array length = 10000, read batch size = 1024
- auto src_array = prepare_array(10000);
- for (int32_t inner_batch_size : {1, 2, 8, 16, 20, 100, 1024}) {
- CheckNextBatchForSingleReader(inner_batch_size,
/*read_batch_size=*/1024, src_array,
- /*selection_bitmap=*/std::nullopt,
- /*expected_array=*/src_array);
- }
- }
- {
- // src array length = 10, selection bitmap = {1, 3, 5}
- auto src_array = prepare_array(10);
- RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({1, 3, 5});
- auto expected_array = prepare_array_with_bitmap(selected_bitmap);
- for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
- CheckNextBatchForSingleReader(inner_batch_size,
/*read_batch_size=*/15, src_array,
- selected_bitmap, expected_array);
- }
- }
- {
- // src array length = 10, selection all
- auto src_array = prepare_array(10);
- RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({0, 1, 2, 3,
4, 5, 6, 7, 8, 9});
- for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
- CheckNextBatchForSingleReader(inner_batch_size,
/*read_batch_size=*/15, src_array,
- selected_bitmap,
/*expected_array=*/src_array);
- }
- }
- {
- // src array length = 10, selection first
- auto src_array = prepare_array(10);
- RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({0});
- auto expected_array = prepare_array_with_bitmap(selected_bitmap);
- for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
- CheckNextBatchForSingleReader(inner_batch_size,
/*read_batch_size=*/15, src_array,
- selected_bitmap, expected_array);
- }
- }
- {
- // src array length = 10, selection first
- auto src_array = prepare_array(10);
- RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({9});
- auto expected_array = prepare_array_with_bitmap(selected_bitmap);
- for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
- CheckNextBatchForSingleReader(inner_batch_size,
/*read_batch_size=*/15, src_array,
- selected_bitmap, expected_array);
- }
- }
- {
- // src array length = 10, selection consecutive positions
- auto src_array = prepare_array(10);
- RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({2, 3, 4, 5});
- auto expected_array = prepare_array_with_bitmap(selected_bitmap);
- for (int32_t inner_batch_size : arrow::internal::Iota(1, 15)) {
- CheckNextBatchForSingleReader(inner_batch_size,
/*read_batch_size=*/15, src_array,
- selected_bitmap, expected_array);
- }
- }
- {
- auto src_array = prepare_array(10);
- RoaringBitmap32 selected_bitmap = RoaringBitmap32::From({0, 1, 2, 3,
4, 5, 7, 8, 9});
- auto expected_array = prepare_array_with_bitmap(selected_bitmap);
- // inner batch: [0, 1, 2, 3] | [4, 5] [7, 8] | [9]
- CheckNextBatchForSingleReader(/*inner_batch_size=*/4,
/*read_batch_size=*/5, src_array,
- selected_bitmap, expected_array);
- }
- {
- // test bulk data, src array length = 10000, read batch size = 1024
- auto src_array = prepare_array(10000);
- RoaringBitmap32 selected_bitmap =
- RoaringBitmap32::From({0, 10, 1000, 2333, 4566, 7838, 8787, 9999});
- auto expected_array = prepare_array_with_bitmap(selected_bitmap);
- for (int32_t inner_batch_size : {1, 2, 8, 16, 20, 100, 1024}) {
- CheckNextBatchForSingleReader(inner_batch_size,
/*read_batch_size=*/1024, src_array,
- selected_bitmap, expected_array);
- }
+TEST_F(DataEvolutionFileReaderTest, TestNoActiveReader) {
+ auto read_schema = arrow::schema({arrow::field("missing",
arrow::int32())});
+ std::vector<std::unique_ptr<BatchReader>> readers;
+ readers.push_back(nullptr);
+ ASSERT_OK_AND_ASSIGN(auto reader,
+ DataEvolutionFileReader::Create(
+ std::move(readers), read_schema,
/*read_batch_size=*/10,
+ /*reader_offsets=*/{-1}, /*field_offsets=*/{-1},
GetArrowPool(pool_)));
+ ASSERT_NOK_WITH_MSG(reader->NextBatchWithBitmap(),
+ "data evolution reader has no active inner reader");
+}
+
+TEST_F(DataEvolutionFileReaderTest, TestDifferentInnerBatchSizes) {
+ auto f0 = arrow::field("f0", arrow::int32());
+ auto f1 = arrow::field("f1", arrow::int32());
+ auto array0 =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({f0}), R"([
+ [0], [1], [2], [3], [4], [5]
+ ])")
+ .ValueOrDie();
+ auto array1 =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({f1}), R"([
+ [10], [11], [12], [13], [14], [15]
+ ])")
+ .ValueOrDie();
+
+ std::vector<std::unique_ptr<BatchReader>> readers;
+ auto reader0 = std::make_unique<MockFileBatchReader>(array0,
array0->type(),
+
/*read_batch_size=*/5);
+ reader0->EnableRandomizeBatchSize(false);
+ readers.push_back(std::move(reader0));
+ auto reader1 = std::make_unique<MockFileBatchReader>(array1,
array1->type(),
+
/*read_batch_size=*/2);
+ reader1->EnableRandomizeBatchSize(false);
+ readers.push_back(std::move(reader1));
+ ASSERT_OK_AND_ASSIGN(auto reader, DataEvolutionFileReader::Create(
+ std::move(readers),
arrow::schema({f0, f1}),
+ /*read_batch_size=*/10,
/*reader_offsets=*/{0, 1},
+ /*field_offsets=*/{0, 0},
GetArrowPool(pool_)));
+
+ arrow::ArrayVector batches;
+ for (int64_t expected_length : {2, 2, 1, 1}) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> batch,
ReadNextOutput(reader.get()));
+ ASSERT_NE(batch, nullptr);
+ ASSERT_EQ(batch->length(), expected_length);
+ ASSERT_EQ(batch->offset(), 0);
+ auto struct_array =
std::dynamic_pointer_cast<arrow::StructArray>(batch);
+ ASSERT_TRUE(struct_array);
+ ASSERT_EQ(struct_array->field(0)->offset(), 0);
+ ASSERT_EQ(struct_array->field(1)->offset(), 0);
+ batches.push_back(std::move(batch));
}
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> eof,
ReadNextOutput(reader.get()));
+ ASSERT_EQ(eof, nullptr);
+ ASSERT_OK_AND_ASSIGN(eof, ReadNextOutput(reader.get()));
+ ASSERT_EQ(eof, nullptr);
+
+ auto expected =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({f0, f1}), R"([
+ [0, 10], [1, 11], [2, 12], [3, 13], [4, 14], [5, 15]
+ ])")
+ .ValueOrDie();
+ ASSERT_TRUE(std::make_shared<arrow::ChunkedArray>(batches)->Equals(
+ std::make_shared<arrow::ChunkedArray>(expected)));
+}
+
+TEST_F(DataEvolutionFileReaderTest, TestBitmapSkipsWholeBatches) {
+ auto f0 = arrow::field("f0", arrow::int32());
+ auto f1 = arrow::field("f1", arrow::int32());
+ auto array0 =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({f0}), R"([
+ [0], [1], [2], [3], [4], [5], [6], [7]
+ ])")
+ .ValueOrDie();
+ auto array1 =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({f1}), R"([
+ [100], [101], [102], [103], [104], [105], [106], [107]
+ ])")
+ .ValueOrDie();
+ RoaringBitmap32 bitmap = RoaringBitmap32::From({4, 7});
+
+ std::vector<std::unique_ptr<BatchReader>> readers;
+ auto reader0 = std::make_unique<MockFileBatchReader>(array0,
array0->type(), bitmap,
+
/*read_batch_size=*/2);
+ reader0->EnableRandomizeBatchSize(false);
+ readers.push_back(std::move(reader0));
+ auto reader1 = std::make_unique<MockFileBatchReader>(array1,
array1->type(), bitmap,
+
/*read_batch_size=*/3);
+ reader1->EnableRandomizeBatchSize(false);
+ readers.push_back(std::move(reader1));
+ ASSERT_OK_AND_ASSIGN(auto reader, DataEvolutionFileReader::Create(
+ std::move(readers),
arrow::schema({f0, f1}),
+ /*read_batch_size=*/10,
/*reader_offsets=*/{0, 1},
+ /*field_offsets=*/{0, 0},
GetArrowPool(pool_)));
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::ChunkedArray> actual,
+
ReadResultCollector::CollectResult(std::move(reader)));
+
+ auto expected =
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({f0, f1}), R"([
+ [4, 104], [7, 107]
+ ])")
+ .ValueOrDie();
+
ASSERT_TRUE(actual->Equals(std::make_shared<arrow::ChunkedArray>(expected)));
}
TEST_P(DataEvolutionFileReaderTest, TestSimple) {
diff --git a/src/paimon/format/blob/blob_file_batch_reader.cpp
b/src/paimon/format/blob/blob_file_batch_reader.cpp
index 531b32a0..0fba02e1 100644
--- a/src/paimon/format/blob/blob_file_batch_reader.cpp
+++ b/src/paimon/format/blob/blob_file_batch_reader.cpp
@@ -22,8 +22,9 @@
#include <array>
#include <cstring>
#include <numeric>
-#include <set>
#include <string>
+#include <string_view>
+#include <unordered_set>
#include "arrow/api.h"
#include "arrow/array/builder_dict.h"
@@ -90,6 +91,141 @@ Result<int32_t> GetMapBlobFixedKeyLength(const
std::shared_ptr<arrow::DataType>&
}
}
+Status AddBuilderCapacity(int64_t increment, int64_t* capacity, const char*
capacity_name) {
+ if (increment < 0 || increment > std::numeric_limits<int64_t>::max() -
*capacity) {
+ return Status::CapacityError(fmt::format("{} exceeds int64 capacity",
capacity_name));
+ }
+ *capacity += increment;
+ return Status::OK();
+}
+
+Status AddBuilderCapacityProduct(int64_t count, int64_t element_size, int64_t*
capacity,
+ const char* capacity_name) {
+ if (count < 0 || element_size < 0 ||
+ (count > 0 && element_size > (std::numeric_limits<int64_t>::max() -
*capacity) / count)) {
+ return Status::CapacityError(fmt::format("{} exceeds int64 capacity",
capacity_name));
+ }
+ *capacity += count * element_size;
+ return Status::OK();
+}
+
+Result<int64_t> GetSerializedDescriptorSize(const std::string& file_path,
+ const std::shared_ptr<MemoryPool>&
pool) {
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<Blob> blob,
+ Blob::FromPath(file_path, /*offset=*/0,
/*length=*/0));
+ PAIMON_UNIQUE_PTR<Bytes> descriptor = blob->ToDescriptor(pool);
+ if (!InRange<int64_t>(descriptor->size())) {
+ return Status::CapacityError("serialized BLOB descriptor exceeds int64
capacity");
+ }
+ return static_cast<int64_t>(descriptor->size());
+}
+
+template <typename T>
+Result<std::shared_ptr<arrow::Buffer>> AllocateTypedBuffer(
+ int64_t element_count, const std::shared_ptr<arrow::MemoryPool>& pool) {
+ if (element_count < 0 ||
+ element_count > std::numeric_limits<int64_t>::max() /
static_cast<int64_t>(sizeof(T))) {
+ return Status::CapacityError("Arrow buffer exceeds int64 capacity");
+ }
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ std::shared_ptr<arrow::Buffer> buffer,
+ arrow::AllocateBuffer(element_count * static_cast<int64_t>(sizeof(T)),
pool.get()));
+ return buffer;
+}
+
+Result<std::shared_ptr<arrow::Buffer>> AllocateValidityBitmap(
+ int64_t value_count, int64_t null_count, const
std::shared_ptr<arrow::MemoryPool>& pool) {
+ if (null_count == 0) {
+ return std::shared_ptr<arrow::Buffer>();
+ }
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Buffer>
null_bitmap,
+ arrow::AllocateBitmap(value_count,
pool.get()));
+ memset(null_bitmap->mutable_data(), 0xFF, null_bitmap->size());
+ return null_bitmap;
+}
+
+class PreallocatedLargeBinaryBuilder {
+ public:
+ static Result<std::unique_ptr<PreallocatedLargeBinaryBuilder>> Create(
+ int64_t value_count, int64_t data_length, int64_t null_count,
+ const std::shared_ptr<arrow::MemoryPool>& pool) {
+ if (value_count < 0 || value_count ==
std::numeric_limits<int64_t>::max() ||
+ data_length < 0 || null_count < 0 || null_count > value_count) {
+ return Status::Invalid("invalid preallocated large binary array
capacity");
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> offsets,
+ AllocateTypedBuffer<int64_t>(value_count + 1,
pool));
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Buffer> data,
+ arrow::AllocateBuffer(data_length,
pool.get()));
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> null_bitmap,
+ AllocateValidityBitmap(value_count, null_count,
pool));
+ auto builder =
+ std::unique_ptr<PreallocatedLargeBinaryBuilder>(new
PreallocatedLargeBinaryBuilder(
+ value_count, data_length, null_count, std::move(offsets),
std::move(data),
+ std::move(null_bitmap)));
+ return builder;
+ }
+
+ Result<uint8_t*> Append(int64_t length) {
+ if (length < 0 || value_index_ >= value_count_ || length >
data_length_ - data_offset_) {
+ return Status::Invalid("large binary value exceeds preallocated
capacity");
+ }
+ uint8_t* target = length == 0 ? nullptr : data_->mutable_data() +
data_offset_;
+ data_offset_ += length;
+ offsets_[++value_index_] = data_offset_;
+ return target;
+ }
+
+ Status AppendNull() {
+ if (value_index_ >= value_count_ || null_bitmap_ == nullptr) {
+ return Status::Invalid("large binary null exceeds preallocated
capacity");
+ }
+ arrow::bit_util::ClearBit(null_bitmap_->mutable_data(), value_index_);
+ offsets_[++value_index_] = data_offset_;
+ ++appended_null_count_;
+ return Status::OK();
+ }
+
+ int64_t length() const {
+ return value_index_;
+ }
+
+ Result<std::shared_ptr<arrow::LargeBinaryArray>> Finish() {
+ if (value_index_ != value_count_ || data_offset_ != data_length_ ||
+ appended_null_count_ != null_count_) {
+ return Status::Invalid("large binary output does not match
preallocated capacity");
+ }
+ return std::make_shared<arrow::LargeBinaryArray>(value_count_,
offsets_buffer_, data_,
+ null_bitmap_,
null_count_);
+ }
+
+ private:
+ PreallocatedLargeBinaryBuilder(int64_t value_count, int64_t data_length,
int64_t null_count,
+ std::shared_ptr<arrow::Buffer> offsets,
+ std::shared_ptr<arrow::Buffer> data,
+ std::shared_ptr<arrow::Buffer> null_bitmap)
+ : value_count_(value_count),
+ data_length_(data_length),
+ null_count_(null_count),
+ offsets_buffer_(std::move(offsets)),
+ data_(std::move(data)),
+ null_bitmap_(std::move(null_bitmap)),
+
offsets_(reinterpret_cast<int64_t*>(offsets_buffer_->mutable_data())) {
+ offsets_[0] = 0;
+ }
+
+ const int64_t value_count_;
+ const int64_t data_length_;
+ const int64_t null_count_;
+ std::shared_ptr<arrow::Buffer> offsets_buffer_;
+ std::shared_ptr<arrow::Buffer> data_;
+ std::shared_ptr<arrow::Buffer> null_bitmap_;
+ int64_t* offsets_;
+ int64_t value_index_ = 0;
+ int64_t data_offset_ = 0;
+ int64_t appended_null_count_ = 0;
+};
+
Status AppendMapBlobKey(const std::shared_ptr<arrow::DataType>& key_type,
const uint8_t* data,
int32_t length, arrow::ArrayBuilder* builder) {
switch (key_type->id()) {
@@ -251,6 +387,14 @@ BlobFileBatchReader::BlobFileBatchReader(
std::iota(target_blob_row_indexes_.begin(),
target_blob_row_indexes_.end(), 0);
}
+Result<int64_t> BlobFileBatchReader::GetOrCreateSerializedDescriptorSize()
const {
+ if (serialized_descriptor_size_ < 0) {
+ PAIMON_ASSIGN_OR_RAISE(serialized_descriptor_size_,
+ GetSerializedDescriptorSize(file_path_, pool_));
+ }
+ return serialized_descriptor_size_;
+}
+
Status BlobFileBatchReader::SetReadSchema(::ArrowSchema* read_schema,
const std::shared_ptr<Predicate>&
predicate,
const
std::optional<RoaringBitmap32>& selection_bitmap) {
@@ -469,31 +613,19 @@ Result<BlobFileBatchReader::ArrayBlobPayload>
BlobFileBatchReader::ReadArrayBlob
return ArrayBlobPayload{std::move(element_lengths), data_offset};
}
-Status BlobFileBatchReader::AppendArrayBlobValues(const ArrayBlobPayload&
payload,
- arrow::LargeBinaryBuilder*
blob_builder) const {
+Status BlobFileBatchReader::AppendArrayBlobDescriptors(
+ const ArrayBlobPayload& payload, arrow::LargeBinaryBuilder* blob_builder)
const {
int64_t element_offset = payload.data_offset;
for (int64_t element_length : payload.element_lengths) {
if (element_length == BlobDefs::kNullBinLength) {
PAIMON_RETURN_NOT_OK_FROM_ARROW(blob_builder->AppendNull());
continue;
}
- if (blob_as_descriptor_) {
- PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<Blob> blob,
- Blob::FromPath(file_path_, element_offset,
element_length));
- PAIMON_UNIQUE_PTR<Bytes> descriptor = blob->ToDescriptor(pool_);
- PAIMON_RETURN_NOT_OK_FROM_ARROW(
- blob_builder->Append(descriptor->data(), descriptor->size()));
- } else {
- PAIMON_UNIQUE_PTR<Bytes> element_bytes =
- Bytes::AllocateBytes(static_cast<size_t>(element_length),
pool_.get());
- if (element_length > 0) {
- PAIMON_RETURN_NOT_OK(
- ReadBlobContentAt(element_offset, element_length,
-
reinterpret_cast<uint8_t*>(element_bytes->data())));
- }
- PAIMON_RETURN_NOT_OK_FROM_ARROW(
- blob_builder->Append(element_bytes->data(), element_length));
- }
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<Blob> blob,
+ Blob::FromPath(file_path_, element_offset,
element_length));
+ PAIMON_UNIQUE_PTR<Bytes> descriptor = blob->ToDescriptor(pool_);
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(
+ blob_builder->Append(descriptor->data(), descriptor->size()));
element_offset += element_length;
}
return Status::OK();
@@ -509,12 +641,118 @@ Result<std::shared_ptr<arrow::Array>>
BlobFileBatchReader::BuildArrayBlobArray(
return Status::Invalid("ARRAY<BLOB> element type must be large
binary");
}
+ std::vector<ArrayBlobPayload> payloads;
+ payloads.reserve(rows_to_read);
+ int64_t value_count = 0;
+ int64_t value_data_length = 0;
+ int64_t descriptor_count = 0;
+ int64_t list_null_count = 0;
+ int64_t value_null_count = 0;
+ for (int32_t k = 0; k < rows_to_read; ++k) {
+ const size_t row_index = current_pos_ + k;
+ if (IsTargetPlaceholder(row_index)) {
+ PAIMON_RETURN_NOT_OK(AddBuilderCapacity(1, &value_count,
"ARRAY<BLOB> element count"));
+
PAIMON_RETURN_NOT_OK(AddBuilderCapacity(BlobDefs::kPlaceholderSentinelLength,
+ &value_data_length,
+ "ARRAY<BLOB> element data
size"));
+ continue;
+ }
+ if (IsTargetNull(row_index)) {
+ ++list_null_count;
+ continue;
+ }
+
+ PAIMON_ASSIGN_OR_RAISE(ArrayBlobPayload payload,
ReadArrayBlobPayload(row_index));
+ PAIMON_RETURN_NOT_OK(
+
AddBuilderCapacity(static_cast<int64_t>(payload.element_lengths.size()),
&value_count,
+ "ARRAY<BLOB> element count"));
+ for (int64_t element_length : payload.element_lengths) {
+ if (element_length == BlobDefs::kNullBinLength) {
+ ++value_null_count;
+ continue;
+ }
+ if (blob_as_descriptor_) {
+ PAIMON_RETURN_NOT_OK(
+ AddBuilderCapacity(1, &descriptor_count, "ARRAY<BLOB>
descriptor count"));
+ } else {
+ PAIMON_RETURN_NOT_OK(AddBuilderCapacity(element_length,
&value_data_length,
+ "ARRAY<BLOB> element
data size"));
+ }
+ }
+ payloads.emplace_back(std::move(payload));
+ }
+ if (descriptor_count > 0) {
+ PAIMON_ASSIGN_OR_RAISE(int64_t descriptor_size,
GetOrCreateSerializedDescriptorSize());
+ PAIMON_RETURN_NOT_OK(AddBuilderCapacityProduct(descriptor_count,
descriptor_size,
+ &value_data_length,
+ "ARRAY<BLOB> descriptor
data size"));
+ }
+
+ if (value_count > std::numeric_limits<int32_t>::max()) {
+ return Status::CapacityError("ARRAY<BLOB> element count exceeds list
offset capacity");
+ }
+ if (!blob_as_descriptor_) {
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> list_offsets,
+ AllocateTypedBuffer<int32_t>(rows_to_read + 1,
arrow_pool_));
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> list_null_bitmap,
+ AllocateValidityBitmap(rows_to_read,
list_null_count, arrow_pool_));
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<PreallocatedLargeBinaryBuilder>
value_builder,
+ PreallocatedLargeBinaryBuilder::Create(
+ value_count, value_data_length,
value_null_count, arrow_pool_));
+
+ auto* raw_list_offsets =
reinterpret_cast<int32_t*>(list_offsets->mutable_data());
+ raw_list_offsets[0] = 0;
+ size_t payload_index = 0;
+ for (int32_t k = 0; k < rows_to_read; ++k) {
+ const size_t row_index = current_pos_ + k;
+ if (IsTargetNull(row_index)) {
+ arrow::bit_util::ClearBit(list_null_bitmap->mutable_data(), k);
+ } else if (IsTargetPlaceholder(row_index)) {
+ PAIMON_ASSIGN_OR_RAISE(auto* target,
+
value_builder->Append(BlobDefs::kPlaceholderSentinelLength));
+ memcpy(target, BlobDefs::kPlaceholderSentinel,
+ BlobDefs::kPlaceholderSentinelLength);
+ } else {
+ const ArrayBlobPayload& payload = payloads[payload_index++];
+ uint8_t* row_target = nullptr;
+ int64_t row_data_length = 0;
+ for (int64_t element_length : payload.element_lengths) {
+ if (element_length == BlobDefs::kNullBinLength) {
+ PAIMON_RETURN_NOT_OK(value_builder->AppendNull());
+ continue;
+ }
+ PAIMON_ASSIGN_OR_RAISE(auto* target,
value_builder->Append(element_length));
+ if (row_target == nullptr && element_length > 0) {
+ row_target = target;
+ }
+ row_data_length += element_length;
+ }
+ if (row_data_length > 0) {
+ PAIMON_RETURN_NOT_OK(
+ ReadBlobContentAt(payload.data_offset,
row_data_length, row_target));
+ }
+ }
+ raw_list_offsets[k + 1] =
static_cast<int32_t>(value_builder->length());
+ }
+
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::LargeBinaryArray> values,
+ value_builder->Finish());
+ auto list_array = std::make_shared<arrow::ListArray>(
+ list_type, rows_to_read, list_offsets, values, list_null_bitmap,
list_null_count);
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::StructArray>
struct_array,
+
arrow::StructArray::Make({list_array}, {list_field}));
+ return struct_array;
+ }
+
std::unique_ptr<arrow::ArrayBuilder> array_builder;
PAIMON_RETURN_NOT_OK_FROM_ARROW(
arrow::MakeBuilder(arrow_pool_.get(), list_type, &array_builder));
auto* list_builder =
checked_cast<arrow::ListBuilder*>(array_builder.get());
auto* blob_builder =
checked_cast<arrow::LargeBinaryBuilder*>(list_builder->value_builder());
PAIMON_RETURN_NOT_OK_FROM_ARROW(list_builder->Reserve(rows_to_read));
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(blob_builder->Reserve(value_count));
+
PAIMON_RETURN_NOT_OK_FROM_ARROW(blob_builder->ReserveData(value_data_length));
+ size_t payload_index = 0;
for (int32_t k = 0; k < rows_to_read; ++k) {
const size_t row_index = current_pos_ + k;
const bool is_null = IsTargetNull(row_index);
@@ -523,8 +761,8 @@ Result<std::shared_ptr<arrow::Array>>
BlobFileBatchReader::BuildArrayBlobArray(
PAIMON_RETURN_NOT_OK_FROM_ARROW(blob_builder->Append(
BlobDefs::kPlaceholderSentinel,
BlobDefs::kPlaceholderSentinelLength));
} else if (!is_null) {
- PAIMON_ASSIGN_OR_RAISE(ArrayBlobPayload payload,
ReadArrayBlobPayload(row_index));
- PAIMON_RETURN_NOT_OK(AppendArrayBlobValues(payload, blob_builder));
+
PAIMON_RETURN_NOT_OK(AppendArrayBlobDescriptors(payloads[payload_index],
blob_builder));
+ ++payload_index;
}
}
@@ -658,56 +896,48 @@ Result<BlobFileBatchReader::MapBlobPayload>
BlobFileBatchReader::ReadMapBlobPayl
Status BlobFileBatchReader::AppendMapBlobKeys(const MapBlobPayload& payload,
const
std::shared_ptr<arrow::DataType>& key_type,
arrow::ArrayBuilder*
key_builder) const {
- int64_t key_offset = payload.data_offset;
- std::set<std::string> serialized_keys;
+ if (!InRange<size_t>(payload.key_data_length)) {
+ return Status::CapacityError("MAP<..., BLOB> key data exceeds size_t
capacity");
+ }
+ PAIMON_UNIQUE_PTR<Bytes> key_bytes =
+ Bytes::AllocateBytes(static_cast<size_t>(payload.key_data_length),
pool_.get());
+ if (payload.key_data_length > 0) {
+ PAIMON_RETURN_NOT_OK(ReadBlobContentAt(payload.data_offset,
payload.key_data_length,
+
reinterpret_cast<uint8_t*>(key_bytes->data())));
+ }
+ int64_t key_offset = 0;
+ std::unordered_set<std::string_view> serialized_keys;
+ serialized_keys.reserve(payload.key_lengths.size());
for (int64_t key_length_64 : payload.key_lengths) {
const auto key_length = static_cast<int32_t>(key_length_64);
- PAIMON_UNIQUE_PTR<Bytes> key_bytes =
- Bytes::AllocateBytes(static_cast<size_t>(key_length), pool_.get());
- if (key_length > 0) {
- PAIMON_RETURN_NOT_OK(ReadBlobContentAt(key_offset, key_length,
-
reinterpret_cast<uint8_t*>(key_bytes->data())));
- }
- std::string serialized_key;
- if (key_length > 0) {
- serialized_key.assign(key_bytes->data(), key_length);
- }
- if (!serialized_keys.emplace(std::move(serialized_key)).second) {
+ const char* key_data = key_length == 0 ? nullptr : key_bytes->data() +
key_offset;
+ const std::string_view serialized_key =
+ key_length == 0 ? std::string_view() : std::string_view(key_data,
key_length);
+ if (!serialized_keys.emplace(serialized_key).second) {
return Status::Invalid("invalid MAP<..., BLOB> payload: duplicate
key");
}
const uint8_t empty_key = 0;
- const uint8_t* key_data =
- key_length == 0 ? &empty_key : reinterpret_cast<const
uint8_t*>(key_bytes->data());
- PAIMON_RETURN_NOT_OK(AppendMapBlobKey(key_type, key_data, key_length,
key_builder));
+ const uint8_t* append_data =
+ key_length == 0 ? &empty_key : reinterpret_cast<const
uint8_t*>(key_data);
+ PAIMON_RETURN_NOT_OK(AppendMapBlobKey(key_type, append_data,
key_length, key_builder));
key_offset += key_length;
}
return Status::OK();
}
-Status BlobFileBatchReader::AppendMapBlobValues(const MapBlobPayload& payload,
- arrow::LargeBinaryBuilder*
blob_builder) const {
+Status BlobFileBatchReader::AppendMapBlobDescriptors(
+ const MapBlobPayload& payload, arrow::LargeBinaryBuilder* blob_builder)
const {
int64_t value_offset = payload.data_offset + payload.key_data_length;
for (int64_t value_length : payload.value_lengths) {
if (value_length == BlobDefs::kNullBinLength) {
PAIMON_RETURN_NOT_OK_FROM_ARROW(blob_builder->AppendNull());
continue;
}
- if (blob_as_descriptor_) {
- PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<Blob> blob,
- Blob::FromPath(file_path_, value_offset,
value_length));
- PAIMON_UNIQUE_PTR<Bytes> descriptor = blob->ToDescriptor(pool_);
- PAIMON_RETURN_NOT_OK_FROM_ARROW(
- blob_builder->Append(descriptor->data(), descriptor->size()));
- } else {
- PAIMON_UNIQUE_PTR<Bytes> value_bytes =
- Bytes::AllocateBytes(static_cast<size_t>(value_length),
pool_.get());
- if (value_length > 0) {
- PAIMON_RETURN_NOT_OK(ReadBlobContentAt(
- value_offset, value_length,
reinterpret_cast<uint8_t*>(value_bytes->data())));
- }
- PAIMON_RETURN_NOT_OK_FROM_ARROW(
- blob_builder->Append(value_bytes->data(), value_length));
- }
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<Blob> blob,
+ Blob::FromPath(file_path_, value_offset,
value_length));
+ PAIMON_UNIQUE_PTR<Bytes> descriptor = blob->ToDescriptor(pool_);
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(
+ blob_builder->Append(descriptor->data(), descriptor->size()));
value_offset += value_length;
}
return Status::OK();
@@ -724,19 +954,140 @@ Result<std::shared_ptr<arrow::Array>>
BlobFileBatchReader::BuildMapBlobArray(
}
PAIMON_ASSIGN_OR_RAISE(int32_t fixed_key_length,
GetMapBlobFixedKeyLength(key_type));
- PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::unique_ptr<arrow::ArrayBuilder>
key_builder_unique,
+ std::vector<MapBlobPayload> payloads;
+ payloads.reserve(rows_to_read);
+ int64_t entry_count = 0;
+ int64_t key_data_length = 0;
+ int64_t value_data_length = 0;
+ int64_t descriptor_count = 0;
+ int64_t map_null_count = 0;
+ int64_t value_null_count = 0;
+ for (int32_t k = 0; k < rows_to_read; ++k) {
+ const size_t row_index = current_pos_ + k;
+ if (IsTargetPlaceholder(row_index)) {
+ PAIMON_RETURN_NOT_OK(AddBuilderCapacity(2, &entry_count, "MAP<...,
BLOB> entry count"));
+ value_null_count += 2;
+ continue;
+ }
+ if (IsTargetNull(row_index)) {
+ ++map_null_count;
+ continue;
+ }
+
+ PAIMON_ASSIGN_OR_RAISE(MapBlobPayload payload,
+ ReadMapBlobPayload(row_index,
fixed_key_length));
+
PAIMON_RETURN_NOT_OK(AddBuilderCapacity(static_cast<int64_t>(payload.value_lengths.size()),
+ &entry_count, "MAP<..., BLOB>
entry count"));
+ PAIMON_RETURN_NOT_OK(AddBuilderCapacity(payload.key_data_length,
&key_data_length,
+ "MAP<..., BLOB> key data
size"));
+ for (int64_t value_length : payload.value_lengths) {
+ if (value_length == BlobDefs::kNullBinLength) {
+ ++value_null_count;
+ continue;
+ }
+ if (blob_as_descriptor_) {
+ PAIMON_RETURN_NOT_OK(
+ AddBuilderCapacity(1, &descriptor_count, "MAP<..., BLOB>
descriptor count"));
+ } else {
+ PAIMON_RETURN_NOT_OK(AddBuilderCapacity(value_length,
&value_data_length,
+ "MAP<..., BLOB> value
data size"));
+ }
+ }
+ payloads.emplace_back(std::move(payload));
+ }
+ if (descriptor_count > 0) {
+ PAIMON_ASSIGN_OR_RAISE(int64_t descriptor_size,
GetOrCreateSerializedDescriptorSize());
+ PAIMON_RETURN_NOT_OK(AddBuilderCapacityProduct(descriptor_count,
descriptor_size,
+ &value_data_length,
+ "MAP<..., BLOB>
descriptor data size"));
+ }
+
+ if (entry_count > std::numeric_limits<int32_t>::max()) {
+ return Status::CapacityError("MAP<..., BLOB> entry count exceeds map
offset capacity");
+ }
+
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::ArrayBuilder>
key_builder,
arrow::MakeBuilder(key_type,
arrow_pool_.get()));
- PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::unique_ptr<arrow::ArrayBuilder>
item_builder_unique,
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(key_builder->Reserve(entry_count));
+ if (key_type->id() == arrow::Type::STRING) {
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(
+
checked_cast<arrow::StringBuilder*>(key_builder.get())->ReserveData(key_data_length));
+ } else if (key_type->id() == arrow::Type::BINARY) {
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(
+
checked_cast<arrow::BinaryBuilder*>(key_builder.get())->ReserveData(key_data_length));
+ }
+
+ if (!blob_as_descriptor_) {
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> map_offsets,
+ AllocateTypedBuffer<int32_t>(rows_to_read + 1,
arrow_pool_));
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Buffer> map_null_bitmap,
+ AllocateValidityBitmap(rows_to_read,
map_null_count, arrow_pool_));
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<PreallocatedLargeBinaryBuilder>
value_builder,
+ PreallocatedLargeBinaryBuilder::Create(
+ entry_count, value_data_length,
value_null_count, arrow_pool_));
+
+ auto* raw_map_offsets =
reinterpret_cast<int32_t*>(map_offsets->mutable_data());
+ raw_map_offsets[0] = 0;
+ size_t payload_index = 0;
+ for (int32_t k = 0; k < rows_to_read; ++k) {
+ const size_t row_index = current_pos_ + k;
+ if (IsTargetNull(row_index)) {
+ arrow::bit_util::ClearBit(map_null_bitmap->mutable_data(), k);
+ } else if (IsTargetPlaceholder(row_index)) {
+
PAIMON_RETURN_NOT_OK_FROM_ARROW(key_builder->AppendEmptyValues(2));
+ PAIMON_RETURN_NOT_OK(value_builder->AppendNull());
+ PAIMON_RETURN_NOT_OK(value_builder->AppendNull());
+ } else {
+ const MapBlobPayload& payload = payloads[payload_index++];
+ PAIMON_RETURN_NOT_OK(AppendMapBlobKeys(payload, key_type,
key_builder.get()));
+ uint8_t* row_target = nullptr;
+ int64_t row_data_length = 0;
+ for (int64_t value_length : payload.value_lengths) {
+ if (value_length == BlobDefs::kNullBinLength) {
+ PAIMON_RETURN_NOT_OK(value_builder->AppendNull());
+ continue;
+ }
+ PAIMON_ASSIGN_OR_RAISE(auto* target,
value_builder->Append(value_length));
+ if (row_target == nullptr && value_length > 0) {
+ row_target = target;
+ }
+ row_data_length += value_length;
+ }
+ if (row_data_length > 0) {
+ PAIMON_RETURN_NOT_OK(
+ ReadBlobContentAt(payload.data_offset +
payload.key_data_length,
+ row_data_length, row_target));
+ }
+ }
+ raw_map_offsets[k + 1] =
static_cast<int32_t>(value_builder->length());
+ }
+
+ std::shared_ptr<arrow::Array> keys;
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(key_builder->Finish(&keys));
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::LargeBinaryArray> values,
+ value_builder->Finish());
+ auto map_array = std::make_shared<arrow::MapArray>(
+ map_type, rows_to_read, map_offsets, keys, values,
map_null_bitmap, map_null_count);
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::StructArray>
struct_array,
+
arrow::StructArray::Make({map_array}, {map_field}));
+ return struct_array;
+ }
+
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::ArrayBuilder>
item_builder,
arrow::MakeBuilder(map_type->item_type(), arrow_pool_.get()));
- std::shared_ptr<arrow::ArrayBuilder>
key_builder(std::move(key_builder_unique));
- std::shared_ptr<arrow::ArrayBuilder>
item_builder(std::move(item_builder_unique));
if (!item_builder || !item_builder->type() ||
item_builder->type()->id() != arrow::Type::LARGE_BINARY) {
return Status::Invalid("cast MAP<..., BLOB> item builder to large
binary builder failed");
}
auto* blob_builder =
checked_cast<arrow::LargeBinaryBuilder*>(item_builder.get());
arrow::MapBuilder map_builder(arrow_pool_.get(), key_builder,
item_builder, map_type);
+ auto* entry_builder =
checked_cast<arrow::StructBuilder*>(map_builder.value_builder());
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(map_builder.Reserve(rows_to_read));
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(entry_builder->Reserve(entry_count));
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(blob_builder->Reserve(entry_count));
+
PAIMON_RETURN_NOT_OK_FROM_ARROW(blob_builder->ReserveData(value_data_length));
+ size_t payload_index = 0;
for (int32_t k = 0; k < rows_to_read; ++k) {
const size_t row_index = current_pos_ + k;
if (IsTargetNull(row_index)) {
@@ -752,11 +1103,11 @@ Result<std::shared_ptr<arrow::Array>>
BlobFileBatchReader::BuildMapBlobArray(
continue;
}
- PAIMON_ASSIGN_OR_RAISE(MapBlobPayload payload,
- ReadMapBlobPayload(row_index,
fixed_key_length));
PAIMON_RETURN_NOT_OK_FROM_ARROW(map_builder.Append());
- PAIMON_RETURN_NOT_OK(AppendMapBlobKeys(payload, key_type,
key_builder.get()));
- PAIMON_RETURN_NOT_OK(AppendMapBlobValues(payload, blob_builder));
+ PAIMON_RETURN_NOT_OK(
+ AppendMapBlobKeys(payloads[payload_index], key_type,
key_builder.get()));
+ PAIMON_RETURN_NOT_OK(AppendMapBlobDescriptors(payloads[payload_index],
blob_builder));
+ ++payload_index;
}
std::shared_ptr<arrow::MapArray> built_map_array;
@@ -796,6 +1147,26 @@ Result<std::shared_ptr<arrow::Array>>
BlobFileBatchReader::BuildTargetArray(
return Status::Invalid("cast to large binary builder failed");
}
auto* field_builder =
checked_cast<arrow::LargeBinaryBuilder*>(field_builder_base);
+ int64_t descriptor_count = 0;
+ int64_t value_data_length = 0;
+ for (int32_t k = 0; k < rows_to_read; ++k) {
+ const size_t i = current_pos_ + k;
+ if (IsTargetPlaceholder(i)) {
+
PAIMON_RETURN_NOT_OK(AddBuilderCapacity(BlobDefs::kPlaceholderSentinelLength,
+ &value_data_length,
+ "BLOB descriptor data
size"));
+ } else if (!IsTargetNull(i)) {
+ PAIMON_RETURN_NOT_OK(AddBuilderCapacity(1, &descriptor_count,
"BLOB descriptor count"));
+ }
+ }
+ if (descriptor_count > 0) {
+ PAIMON_ASSIGN_OR_RAISE(int64_t descriptor_size,
GetOrCreateSerializedDescriptorSize());
+ PAIMON_RETURN_NOT_OK(AddBuilderCapacityProduct(
+ descriptor_count, descriptor_size, &value_data_length, "BLOB
descriptor data size"));
+ }
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(builder->Reserve(rows_to_read));
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(field_builder->Reserve(rows_to_read));
+
PAIMON_RETURN_NOT_OK_FROM_ARROW(field_builder->ReserveData(value_data_length));
for (int32_t k = 0; k < rows_to_read; ++k) {
const size_t i = current_pos_ + k;
PAIMON_RETURN_NOT_OK_FROM_ARROW(builder->Append());
diff --git a/src/paimon/format/blob/blob_file_batch_reader.h
b/src/paimon/format/blob/blob_file_batch_reader.h
index e49a2180..372c742d 100644
--- a/src/paimon/format/blob/blob_file_batch_reader.h
+++ b/src/paimon/format/blob/blob_file_batch_reader.h
@@ -178,6 +178,8 @@ class BlobFileBatchReader : public FileBatchReader {
const std::shared_ptr<MemoryPool>& pool,
const std::shared_ptr<arrow::MemoryPool>& arrow_pool);
+ Result<int64_t> GetOrCreateSerializedDescriptorSize() const;
+
Status ReadBlobContentAt(const int64_t offset, const int64_t length,
uint8_t* content) const;
Result<std::shared_ptr<arrow::Buffer>> NextBlobOffsets(int32_t
rows_to_read) const;
@@ -186,15 +188,15 @@ class BlobFileBatchReader : public FileBatchReader {
Result<std::shared_ptr<arrow::Buffer>> BuildNullBitmap(int32_t
rows_to_read) const;
Result<std::shared_ptr<arrow::Array>> BuildContentArray(int32_t
rows_to_read) const;
Result<ArrayBlobPayload> ReadArrayBlobPayload(size_t row_index) const;
- Status AppendArrayBlobValues(const ArrayBlobPayload& payload,
- arrow::LargeBinaryBuilder* blob_builder)
const;
+ Status AppendArrayBlobDescriptors(const ArrayBlobPayload& payload,
+ arrow::LargeBinaryBuilder* blob_builder)
const;
Result<std::shared_ptr<arrow::Array>> BuildArrayBlobArray(int32_t
rows_to_read) const;
Result<MapBlobPayload> ReadMapBlobPayload(size_t row_index, int32_t
fixed_key_length) const;
Status AppendMapBlobKeys(const MapBlobPayload& payload,
const std::shared_ptr<arrow::DataType>& key_type,
arrow::ArrayBuilder* key_builder) const;
- Status AppendMapBlobValues(const MapBlobPayload& payload,
- arrow::LargeBinaryBuilder* blob_builder) const;
+ Status AppendMapBlobDescriptors(const MapBlobPayload& payload,
+ arrow::LargeBinaryBuilder* blob_builder)
const;
Result<std::shared_ptr<arrow::Array>> BuildMapBlobArray(int32_t
rows_to_read) const;
Result<std::shared_ptr<arrow::Array>> BuildTargetArray(int32_t
rows_to_read) const;
@@ -242,6 +244,7 @@ class BlobFileBatchReader : public FileBatchReader {
const bool emit_placeholder_sentinel_;
std::shared_ptr<MemoryPool> pool_;
std::shared_ptr<arrow::MemoryPool> arrow_pool_;
+ mutable int64_t serialized_descriptor_size_ = -1;
std::shared_ptr<arrow::DataType> target_type_;
std::shared_ptr<Metrics> metrics_;
