| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, |
| * software distributed under the License is distributed on an |
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| * KIND, either express or implied. See the License for the |
| * specific language governing permissions and limitations |
| * under the License. |
| */ |
| |
| #include "paimon/core/utils/blob_view_lookup.h" |
| |
| #include <algorithm> |
| #include <future> |
| #include <utility> |
| #include <vector> |
| |
| #include "arrow/array.h" |
| #include "arrow/c/bridge.h" |
| #include "fmt/format.h" |
| #include "paimon/catalog/catalog.h" |
| #include "paimon/common/data/blob_descriptor.h" |
| #include "paimon/common/executor/future.h" |
| #include "paimon/common/table/special_fields.h" |
| #include "paimon/common/utils/arrow/status_utils.h" |
| #include "paimon/common/utils/checked_cast.h" |
| #include "paimon/common/utils/path_util.h" |
| #include "paimon/defs.h" |
| #include "paimon/executor.h" |
| #include "paimon/global_index/bitmap_global_index_result.h" |
| #include "paimon/memory/bytes.h" |
| #include "paimon/read_context.h" |
| #include "paimon/scan_context.h" |
| #include "paimon/table/source/table_read.h" |
| #include "paimon/table/source/table_scan.h" |
| #include "paimon/utils/special_field_ids.h" |
| |
| namespace paimon { |
| BlobViewLookup::TableReadPlan::TableReadPlan(const BlobViewStruct& view_struct) |
| : identifier_(view_struct.GetIdentifier()) { |
| references_by_field_id_.insert(view_struct.FieldId()); |
| row_ranges_.push_back(view_struct.RowId()); |
| } |
| |
| void BlobViewLookup::TableReadPlan::Add(const BlobViewStruct& view_struct) { |
| references_by_field_id_.insert(view_struct.FieldId()); |
| row_ranges_.push_back(view_struct.RowId()); |
| } |
| |
| const Identifier& BlobViewLookup::TableReadPlan::GetIdentifier() const { |
| return identifier_; |
| } |
| |
| std::vector<int32_t> BlobViewLookup::TableReadPlan::GetFieldIds() const { |
| return std::vector<int32_t>(references_by_field_id_.begin(), references_by_field_id_.end()); |
| } |
| |
| std::vector<Range> BlobViewLookup::TableReadPlan::GetSortedDistinctRanges() const { |
| if (row_ranges_.empty()) { |
| return {}; |
| } |
| std::vector<int64_t> sorted = row_ranges_; |
| std::sort(sorted.begin(), sorted.end()); |
| std::vector<Range> ranges; |
| int64_t range_start = sorted[0]; |
| int64_t range_end = range_start; |
| for (size_t i = 1; i < sorted.size(); ++i) { |
| const int64_t row_id = sorted[i]; |
| if (row_id == range_end) { |
| continue; |
| } |
| if (row_id != range_end + 1) { |
| ranges.emplace_back(range_start, range_end); |
| range_start = row_id; |
| } |
| range_end = row_id; |
| } |
| ranges.emplace_back(range_start, range_end); |
| return ranges; |
| } |
| |
| Result<BlobViewResolver> BlobViewLookup::CreateResolver( |
| const std::unordered_set<BlobViewStruct>& view_structs, |
| const std::shared_ptr<CatalogContext>& catalog_context, const std::shared_ptr<MemoryPool>& pool, |
| const std::shared_ptr<Executor>& executor) { |
| PAIMON_ASSIGN_OR_RAISE(DescriptorMapping mapping, |
| PreloadDescriptors(view_structs, catalog_context, pool, executor)); |
| return BlobViewResolver([cached = std::move(mapping)](const BlobViewStruct& view_struct) |
| -> Result<std::shared_ptr<Bytes>> { |
| auto iter = cached.find(view_struct); |
| if (iter == cached.end()) { |
| return Status::Invalid(fmt::format("BlobViewStruct not found in preloaded cache: {}", |
| view_struct.ToString())); |
| } |
| return iter->second; |
| }); |
| } |
| |
| Result<BlobViewLookup::DescriptorMapping> BlobViewLookup::PreloadDescriptors( |
| const std::unordered_set<BlobViewStruct>& view_structs, |
| const std::shared_ptr<CatalogContext>& catalog_context, const std::shared_ptr<MemoryPool>& pool, |
| const std::shared_ptr<Executor>& executor) { |
| std::unordered_map<Identifier, BlobViewLookup::TableReadPlan> plan_by_identifier = |
| GroupByIdentifier(view_structs); |
| int64_t target_rows_per_task = TargetRowsPerTask(plan_by_identifier, executor->GetThreadNum()); |
| |
| std::vector<std::future<Result<DescriptorMapping>>> futures; |
| for (const auto& [identifier, table_read_plan] : plan_by_identifier) { |
| const auto& id = identifier; |
| std::vector<int32_t> field_ids = table_read_plan.GetFieldIds(); |
| std::vector<std::vector<Range>> range_chunks = |
| SplitRowRanges(table_read_plan.GetSortedDistinctRanges(), target_rows_per_task); |
| for (const auto& range_chunk : range_chunks) { |
| futures.push_back(Via( |
| executor.get(), |
| [catalog_context, id, field_ids, range_chunk, pool]() -> Result<DescriptorMapping> { |
| return LoadTableDescriptorChunk(catalog_context, id, field_ids, range_chunk, |
| pool); |
| })); |
| } |
| } |
| |
| DescriptorMapping mapping; |
| std::vector<Result<DescriptorMapping>> chunk_results = CollectAll(futures); |
| for (auto& chunk_result : chunk_results) { |
| if (!chunk_result.ok()) { |
| return chunk_result.status(); |
| } |
| for (const auto& [view_struct, descriptor] : chunk_result.value()) { |
| mapping[view_struct] = descriptor; |
| } |
| } |
| return mapping; |
| } |
| |
| Result<BlobViewLookup::DescriptorMapping> BlobViewLookup::LoadTableDescriptorChunk( |
| const std::shared_ptr<CatalogContext>& catalog_context, const Identifier& identifier, |
| const std::vector<int32_t>& field_ids, const std::vector<Range>& row_ranges, |
| const std::shared_ptr<MemoryPool>& pool) { |
| PAIMON_ASSIGN_OR_RAISE(std::optional<std::string> branch, identifier.GetBranchName()); |
| if (branch) { |
| return Status::Invalid("do not support upstream table with branch"); |
| } |
| auto file_system = catalog_context->file_system; |
| PAIMON_ASSIGN_OR_RAISE(std::string table_path, GetTableLocation(catalog_context, identifier)); |
| ScanContextBuilder scan_builder(table_path); |
| auto global_index_result = BitmapGlobalIndexResult::FromRanges(row_ranges); |
| scan_builder.SetGlobalIndexResult(global_index_result) |
| .WithMemoryPool(pool) |
| .WithFileSystem(file_system); |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ScanContext> scan_context, scan_builder.Finish()); |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableScan> table_scan, |
| TableScan::Create(std::move(scan_context))); |
| PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Plan> plan, table_scan->CreatePlan()); |
| |
| ReadContextBuilder read_builder(table_path); |
| std::vector<int32_t> read_field_ids = field_ids; |
| read_field_ids.push_back(SpecialFieldIds::ROW_ID); |
| read_builder.SetReadFieldIds(read_field_ids) |
| .AddOption(Options::BLOB_AS_DESCRIPTOR, "true") |
| .EnablePrefetch(true) |
| .WithMemoryPool(pool) |
| .WithFileSystem(file_system); |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> read_context, read_builder.Finish()); |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableRead> table_read, |
| TableRead::Create(std::move(read_context))); |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BatchReader> reader, |
| table_read->CreateReader(plan->Splits())); |
| |
| DescriptorMapping mapping; |
| PAIMON_RETURN_NOT_OK( |
| ExtractBlobDescriptors(identifier, read_field_ids, pool, reader.get(), &mapping)); |
| return mapping; |
| } |
| |
| Result<std::string> BlobViewLookup::GetTableLocation( |
| const std::shared_ptr<CatalogContext>& catalog_context, const Identifier& identifier) { |
| auto file_system = catalog_context->file_system; |
| PAIMON_ASSIGN_OR_RAISE( |
| std::unique_ptr<Catalog> catalog, |
| Catalog::Create(catalog_context->root_path, catalog_context->options, file_system)); |
| // The table path may be either xxx/test_database/test_table or |
| // xxx/test_database.db/test_table. If neither path exists or both paths exist, it means we |
| // cannot infer the table path, and an error will be reported. If only one of the paths |
| // exists, we will use that path. |
| PAIMON_ASSIGN_OR_RAISE(std::string source_table_path, catalog->GetTableLocation(identifier)); |
| std::string database_name = identifier.GetDatabaseName(); |
| PAIMON_ASSIGN_OR_RAISE(std::string data_table_name, identifier.GetDataTableName()); |
| std::string fallback_source_table_path = PathUtil::JoinPath( |
| PathUtil::JoinPath(catalog_context->root_path, database_name), data_table_name); |
| |
| PAIMON_ASSIGN_OR_RAISE(bool exist, catalog_context->file_system->Exists(source_table_path)); |
| PAIMON_ASSIGN_OR_RAISE(bool fallback_exist, file_system->Exists(fallback_source_table_path)); |
| if (exist == fallback_exist) { |
| return Status::Invalid( |
| fmt::format("Ambiguous table path: both table path {} and fallback table path {} are " |
| "present or absent", |
| source_table_path, fallback_source_table_path)); |
| } |
| std::string final_table_path = exist ? source_table_path : fallback_source_table_path; |
| return final_table_path; |
| } |
| |
| Status BlobViewLookup::ExtractBlobDescriptors(const Identifier& identifier, |
| const std::vector<int32_t>& field_ids, |
| const std::shared_ptr<MemoryPool>& pool, |
| BatchReader* reader, DescriptorMapping* mapping) { |
| if (reader == nullptr) { |
| return Status::Invalid("invalid reader in ExtractBlobDescriptors, reader is nullptr"); |
| } |
| while (true) { |
| PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); |
| if (BatchReader::IsEofBatch(batch)) { |
| break; |
| } |
| auto& [c_array, c_schema] = batch; |
| PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> arrow_array, |
| arrow::ImportArray(c_array.get(), c_schema.get())); |
| if (!arrow_array || arrow_array->type_id() != arrow::Type::STRUCT) { |
| return Status::Invalid( |
| "invalid array in ExtractBlobDescriptors, batch array is not a StructArray."); |
| } |
| auto struct_array = checked_pointer_cast<arrow::StructArray>(arrow_array); |
| // skip the _VALUE_KIND column |
| if (static_cast<size_t>(struct_array->num_fields()) - 1 != field_ids.size()) { |
| return Status::Invalid( |
| fmt::format("invalid array in ExtractBlobDescriptors, batch array fields(exclude " |
| "_VALUE_KIND) {} mismatch read field ids {}.", |
| struct_array->num_fields() - 1, field_ids.size())); |
| } |
| // get _VALUE_KIND |
| if (struct_array->struct_type()->field(0)->name() != SpecialFields::ValueKind().Name()) { |
| return Status::Invalid( |
| "invalid array in ExtractBlobDescriptors, expected _VALUE_KIND as the first " |
| "column"); |
| } |
| |
| // get _ROW_ID |
| if (struct_array->struct_type()->field(struct_array->num_fields() - 1)->name() != |
| SpecialFields::RowId().Name()) { |
| return Status::Invalid( |
| "invalid array in ExtractBlobDescriptors, expected _ROW_ID as the last column"); |
| } |
| auto row_id_array = struct_array->field(struct_array->num_fields() - 1); |
| if (!row_id_array || row_id_array->type_id() != arrow::Type::INT64) { |
| return Status::Invalid( |
| fmt::format("invalid array does not contain {} field, or it cannot be casted to " |
| "Int64Array in ExtractBlobDescriptors.", |
| SpecialFields::RowId().Name())); |
| } |
| auto typed_row_id_array = checked_pointer_cast<arrow::Int64Array>(row_id_array); |
| |
| // skip _VALUE_KIND |
| for (int32_t idx = 1; idx < struct_array->num_fields() - 1; ++idx) { |
| auto field_array = struct_array->field(idx); |
| if (!field_array || field_array->type_id() != arrow::Type::LARGE_BINARY) { |
| return Status::Invalid( |
| "invalid array in ExtractBlobDescriptors, column is not a LargeBinaryArray."); |
| } |
| auto binary_array = checked_pointer_cast<arrow::LargeBinaryArray>(field_array); |
| for (int64_t row = 0; row < binary_array->length(); ++row) { |
| BlobViewStruct blob_view_struct(identifier, field_ids[idx - 1], |
| typed_row_id_array->Value(row)); |
| if (binary_array->IsNull(row)) { |
| // null in source table |
| (*mapping)[blob_view_struct] = nullptr; |
| continue; |
| } |
| std::string_view bytes = binary_array->GetView(row); |
| PAIMON_ASSIGN_OR_RAISE(bool is_descriptor, BlobDescriptor::IsBlobDescriptor( |
| bytes.data(), bytes.size())); |
| if (!is_descriptor) { |
| return Status::Invalid( |
| "requires blob field value to be a serialized BlobDescriptor in source " |
| "table."); |
| } |
| auto descriptor_bytes = std::make_shared<Bytes>(bytes.size(), pool.get()); |
| std::memcpy(descriptor_bytes->data(), bytes.data(), bytes.size()); |
| (*mapping)[blob_view_struct] = std::move(descriptor_bytes); |
| } |
| } |
| } |
| return Status::OK(); |
| } |
| |
| std::unordered_map<Identifier, BlobViewLookup::TableReadPlan> BlobViewLookup::GroupByIdentifier( |
| const std::unordered_set<BlobViewStruct>& view_structs) { |
| std::unordered_map<Identifier, BlobViewLookup::TableReadPlan> grouped; |
| for (const auto& view_struct : view_structs) { |
| auto identifier = view_struct.GetIdentifier(); |
| auto iter = grouped.find(identifier); |
| if (iter != grouped.end()) { |
| iter->second.Add(view_struct); |
| } else { |
| grouped.emplace(identifier, BlobViewLookup::TableReadPlan(view_struct)); |
| } |
| } |
| return grouped; |
| } |
| |
| int64_t BlobViewLookup::TargetRowsPerTask( |
| const std::unordered_map<Identifier, TableReadPlan>& plan_by_identifier, uint32_t thread_num) { |
| int64_t total_rows = 0; |
| for (const auto& [identifier, table_read_plan] : plan_by_identifier) { |
| for (const auto& row_range : table_read_plan.GetSortedDistinctRanges()) { |
| total_rows += row_range.Count(); |
| } |
| } |
| int64_t balanced_rows = (total_rows + thread_num - 1) / thread_num; |
| return std::max(MIN_ROW_PER_TASK, balanced_rows); |
| } |
| |
| std::vector<std::vector<Range>> BlobViewLookup::SplitRowRanges(const std::vector<Range>& row_ranges, |
| int64_t target_rows_per_task) { |
| if (row_ranges.empty()) { |
| return {}; |
| } |
| std::vector<std::vector<Range>> chunks; |
| std::vector<Range> current_chunk; |
| int64_t current_chunk_rows = 0; |
| for (const auto& row_range : row_ranges) { |
| int64_t next_from = row_range.from; |
| while (next_from <= row_range.to) { |
| if (current_chunk_rows == target_rows_per_task) { |
| chunks.push_back(current_chunk); |
| current_chunk.clear(); |
| current_chunk_rows = 0; |
| } |
| |
| int64_t remaining_rows = target_rows_per_task - current_chunk_rows; |
| int64_t next_to = std::min(row_range.to, next_from + remaining_rows - 1); |
| current_chunk.emplace_back(next_from, next_to); |
| current_chunk_rows += next_to - next_from + 1; |
| next_from = next_to + 1; |
| } |
| } |
| if (!current_chunk.empty()) { |
| chunks.push_back(current_chunk); |
| } |
| return chunks; |
| } |
| |
| } // namespace paimon |
