| /* |
| * 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/common/data/shredding/map_shared_shredding_batch_converter.h" |
| |
| #include <map> |
| #include <string> |
| #include <utility> |
| |
| #include "arrow/array.h" |
| #include "arrow/builder.h" |
| #include "arrow/c/bridge.h" |
| #include "arrow/type.h" |
| #include "fmt/format.h" |
| #include "paimon/common/data/shredding/lru_map_shared_shredding_column_allocator.h" |
| #include "paimon/common/data/shredding/map_shared_shredding_context.h" |
| #include "paimon/common/data/shredding/map_shared_shredding_utils.h" |
| #include "paimon/common/data/shredding/map_shredding_defs.h" |
| #include "paimon/common/data/shredding/plain_map_shared_shredding_column_allocator.h" |
| #include "paimon/common/data/shredding/sequential_map_shared_shredding_column_allocator.h" |
| #include "paimon/common/utils/arrow/mem_utils.h" |
| #include "paimon/common/utils/arrow/status_utils.h" |
| #include "paimon/core/core_options.h" |
| namespace paimon { |
| /// Checks that a dynamic_cast result is not null, returning Status::Invalid on failure. |
| #define PAIMON_CHECK_NOT_NULL(ptr, msg) \ |
| do { \ |
| if (PAIMON_UNLIKELY((ptr) == nullptr)) { \ |
| return Status::Invalid(msg); \ |
| } \ |
| } while (false) |
| |
| namespace { |
| |
| Result<std::unique_ptr<MapSharedShreddingColumnAllocator>> CreateMapSharedShreddingColumnAllocator( |
| int32_t num_columns, MapSharedShreddingColumnPlacementPolicy placement_policy) { |
| switch (placement_policy) { |
| case MapSharedShreddingColumnPlacementPolicy::PLAIN: |
| return std::make_unique<PlainMapSharedShreddingColumnAllocator>(num_columns); |
| case MapSharedShreddingColumnPlacementPolicy::SEQUENTIAL: |
| return std::make_unique<SequentialMapSharedShreddingColumnAllocator>(num_columns); |
| case MapSharedShreddingColumnPlacementPolicy::LRU: |
| return std::make_unique<LruMapSharedShreddingColumnAllocator>(num_columns); |
| } |
| return Status::Invalid("unknown shared-shredding column placement policy"); |
| } |
| |
| } // namespace |
| |
| Result<std::shared_ptr<MapSharedShreddingBatchConverter>> MapSharedShreddingBatchConverter::Create( |
| const std::shared_ptr<arrow::Schema>& logical_schema, |
| const std::shared_ptr<MapSharedShreddingContext>& context, const CoreOptions& options, |
| const std::shared_ptr<MemoryPool>& pool) { |
| std::map<std::string, int32_t> field_to_num_columns = context->ComputeNextK(); |
| PAIMON_ASSIGN_OR_RAISE( |
| std::shared_ptr<arrow::Schema> physical_schema, |
| MapSharedShreddingUtils::LogicalToPhysicalSchema(logical_schema, field_to_num_columns)); |
| std::vector<ColumnContext> contexts; |
| std::vector<std::string> shredding_field_names; |
| contexts.reserve(field_to_num_columns.size()); |
| shredding_field_names.reserve(field_to_num_columns.size()); |
| // Iterate in schema field order (not map order) so that shredding_field_names_ |
| // matches the order in which shredding columns appear in the schema. |
| // This is critical for the sequential matching logic in Convert(). |
| for (int32_t i = 0; i < logical_schema->num_fields(); ++i) { |
| const std::string& name = logical_schema->field(i)->name(); |
| auto it = field_to_num_columns.find(name); |
| if (it != field_to_num_columns.end()) { |
| int32_t num_columns = it->second; |
| PAIMON_ASSIGN_OR_RAISE(MapSharedShreddingColumnPlacementPolicy placement_policy, |
| options.GetMapSharedShreddingColumnPlacementPolicy(name)); |
| PAIMON_ASSIGN_OR_RAISE( |
| std::unique_ptr<MapSharedShreddingColumnAllocator> allocator, |
| CreateMapSharedShreddingColumnAllocator(num_columns, placement_policy)); |
| contexts.emplace_back(name, num_columns, std::move(allocator)); |
| shredding_field_names.push_back(name); |
| } |
| } |
| return std::shared_ptr<MapSharedShreddingBatchConverter>( |
| new MapSharedShreddingBatchConverter(logical_schema, physical_schema, std::move(contexts), |
| std::move(shredding_field_names), pool)); |
| } |
| |
| MapSharedShreddingBatchConverter::MapSharedShreddingBatchConverter( |
| const std::shared_ptr<arrow::Schema>& logical_schema, |
| const std::shared_ptr<arrow::Schema>& physical_schema, std::vector<ColumnContext>&& contexts, |
| std::vector<std::string>&& shredding_field_names, const std::shared_ptr<MemoryPool>& pool) |
| : logical_schema_(logical_schema), |
| physical_schema_(physical_schema), |
| contexts_(std::move(contexts)), |
| shredding_field_names_(std::move(shredding_field_names)), |
| pool_(GetArrowPool(pool)) {} |
| |
| const std::shared_ptr<arrow::Schema>& MapSharedShreddingBatchConverter::GetPhysicalSchema() const { |
| return physical_schema_; |
| } |
| |
| Result<std::unique_ptr<ArrowArray>> MapSharedShreddingBatchConverter::Convert( |
| ArrowArray* logical_batch) { |
| std::shared_ptr<arrow::DataType> logical_type = arrow::struct_(logical_schema_->fields()); |
| PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> logical_array, |
| arrow::ImportArray(logical_batch, logical_type)); |
| auto logical_struct = std::dynamic_pointer_cast<arrow::StructArray>(logical_array); |
| PAIMON_CHECK_NOT_NULL(logical_struct, |
| "MapSharedShreddingBatchConverter: input is not a StructArray"); |
| |
| int32_t num_fields = logical_schema_->num_fields(); |
| arrow::ArrayVector physical_columns; |
| physical_columns.reserve(num_fields); |
| size_t context_idx = 0; |
| for (int32_t col = 0; col < num_fields; ++col) { |
| auto column = logical_struct->field(col); |
| const std::string& field_name = logical_schema_->field(col)->name(); |
| if (context_idx < shredding_field_names_.size() && |
| shredding_field_names_[context_idx] == field_name) { |
| auto physical_struct_type = physical_schema_->field(col)->type(); |
| PAIMON_ASSIGN_OR_RAISE( |
| std::shared_ptr<arrow::Array> physical_column, |
| ConvertOneColumn(column, physical_struct_type, &contexts_[context_idx])); |
| physical_columns.push_back(std::move(physical_column)); |
| ++context_idx; |
| } else { |
| physical_columns.push_back(column); |
| } |
| } |
| |
| PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( |
| std::shared_ptr<arrow::Array> physical_struct, |
| arrow::StructArray::Make(physical_columns, physical_schema_->field_names())); |
| |
| std::unique_ptr<ArrowArray> result = std::make_unique<ArrowArray>(); |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*physical_struct, result.get())); |
| return result; |
| } |
| |
| Result<std::shared_ptr<arrow::Array>> MapSharedShreddingBatchConverter::ConvertOneColumn( |
| const std::shared_ptr<arrow::Array>& map_column, |
| const std::shared_ptr<arrow::DataType>& physical_struct_type, ColumnContext* context) const { |
| auto map_array = std::dynamic_pointer_cast<arrow::MapArray>(map_column); |
| PAIMON_CHECK_NOT_NULL(map_array, "MapSharedShreddingBatchConverter: column is not a MapArray"); |
| |
| int64_t num_rows = map_array->length(); |
| int32_t num_cols = context->num_columns; |
| |
| auto keys_array = std::dynamic_pointer_cast<arrow::StringArray>(map_array->keys()); |
| PAIMON_CHECK_NOT_NULL(keys_array, |
| "MapSharedShreddingBatchConverter: MAP keys are not StringArray"); |
| auto values_array = map_array->items(); |
| |
| // Create StructBuilder from physical struct type — it owns all child builders. |
| PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::unique_ptr<arrow::ArrayBuilder> struct_builder_base, |
| arrow::MakeBuilder(physical_struct_type, pool_.get())); |
| auto* struct_builder = dynamic_cast<arrow::StructBuilder*>(struct_builder_base.get()); |
| PAIMON_CHECK_NOT_NULL(struct_builder, |
| "MapSharedShreddingBatchConverter: failed to create StructBuilder"); |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->Reserve(num_rows)); |
| |
| // Extract child builders: [field_mapping, col_0..K-1, overflow] |
| auto* field_mapping_builder = |
| dynamic_cast<arrow::ListBuilder*>(struct_builder->field_builder(0)); |
| PAIMON_CHECK_NOT_NULL(field_mapping_builder, |
| "MapSharedShreddingBatchConverter: field_mapping is not a ListBuilder"); |
| auto* field_mapping_value_builder = |
| dynamic_cast<arrow::Int32Builder*>(field_mapping_builder->value_builder()); |
| PAIMON_CHECK_NOT_NULL( |
| field_mapping_value_builder, |
| "MapSharedShreddingBatchConverter: field_mapping value is not Int32Builder"); |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(field_mapping_builder->Reserve(num_rows)); |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(field_mapping_value_builder->Reserve(num_rows * num_cols)); |
| |
| std::vector<arrow::ArrayBuilder*> col_builders_raw; |
| col_builders_raw.reserve(num_cols); |
| for (int32_t c = 0; c < num_cols; ++c) { |
| arrow::ArrayBuilder* col_builder = struct_builder->field_builder(1 + c); |
| PAIMON_CHECK_NOT_NULL(col_builder, "MapSharedShreddingBatchConverter: col builder is null"); |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(col_builder->Reserve(num_rows)); |
| col_builders_raw.push_back(col_builder); |
| } |
| |
| int32_t overflow_field_idx = 1 + num_cols; |
| auto* overflow_builder = |
| dynamic_cast<arrow::MapBuilder*>(struct_builder->field_builder(overflow_field_idx)); |
| PAIMON_CHECK_NOT_NULL(overflow_builder, |
| "MapSharedShreddingBatchConverter: overflow is not a MapBuilder"); |
| auto* overflow_key_builder = |
| dynamic_cast<arrow::Int32Builder*>(overflow_builder->key_builder()); |
| PAIMON_CHECK_NOT_NULL(overflow_key_builder, |
| "MapSharedShreddingBatchConverter: overflow key is not Int32Builder"); |
| arrow::ArrayBuilder* overflow_value_builder = overflow_builder->item_builder(); |
| PAIMON_CHECK_NOT_NULL(overflow_value_builder, |
| "MapSharedShreddingBatchConverter: overflow value builder is null"); |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(overflow_builder->Reserve(num_rows)); |
| |
| // Process each row |
| for (int64_t row = 0; row < num_rows; ++row) { |
| if (map_array->IsNull(row)) { |
| // StructBuilder::AppendNull() auto-appends empty values to all children. |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->AppendNull()); |
| continue; |
| } |
| |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->Append()); |
| |
| int64_t start = map_array->value_offset(row); |
| int64_t length = map_array->value_length(row); |
| |
| // Extract field ids and build lookup map |
| std::vector<int32_t> field_ids; |
| std::unordered_map<int32_t, int64_t> field_id_to_value_index; |
| ExtractRowFields(keys_array, start, length, &context->dict, &field_ids, |
| &field_id_to_value_index); |
| |
| // Allocate columns |
| RowAllocation allocation = context->allocator->AllocateRow(field_ids); |
| |
| // Fill sub-columns |
| PAIMON_RETURN_NOT_OK(AppendFieldMapping(allocation, num_cols, field_mapping_builder, |
| field_mapping_value_builder)); |
| PAIMON_RETURN_NOT_OK(AppendColumnValues(values_array, allocation, field_id_to_value_index, |
| num_cols, col_builders_raw)); |
| PAIMON_RETURN_NOT_OK(AppendOverflow(values_array, allocation, field_id_to_value_index, |
| overflow_builder, overflow_key_builder, |
| overflow_value_builder)); |
| } |
| |
| // Finalize |
| std::shared_ptr<arrow::StructArray> result; |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->Finish(&result)); |
| return result; |
| } |
| |
| void MapSharedShreddingBatchConverter::ExtractRowFields( |
| const std::shared_ptr<arrow::StringArray>& keys_array, int64_t start, int64_t length, |
| MapSharedShreddingFieldDict* dict, std::vector<int32_t>* field_ids_out, |
| std::unordered_map<int32_t, int64_t>* field_id_to_value_index_out) const { |
| field_ids_out->clear(); |
| field_ids_out->reserve(length); |
| field_id_to_value_index_out->clear(); |
| field_id_to_value_index_out->reserve(length); |
| for (int64_t j = 0; j < length; ++j) { |
| std::string key_str = keys_array->GetString(start + j); |
| int32_t field_id = dict->GetOrAssign(key_str); |
| field_ids_out->push_back(field_id); |
| (*field_id_to_value_index_out)[field_id] = start + j; |
| } |
| } |
| |
| Status MapSharedShreddingBatchConverter::AppendFieldMapping( |
| const RowAllocation& allocation, int32_t num_cols, arrow::ListBuilder* list_builder, |
| arrow::Int32Builder* value_builder) const { |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(list_builder->Append()); |
| for (int32_t c = 0; c < num_cols; ++c) { |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(value_builder->Append(allocation.col_to_field[c])); |
| } |
| return Status::OK(); |
| } |
| |
| Status MapSharedShreddingBatchConverter::AppendColumnValues( |
| const std::shared_ptr<arrow::Array>& values_array, const RowAllocation& allocation, |
| const std::unordered_map<int32_t, int64_t>& field_id_to_value_index, int32_t num_cols, |
| const std::vector<arrow::ArrayBuilder*>& col_builders) const { |
| for (int32_t c = 0; c < num_cols; ++c) { |
| int32_t assigned_field_id = allocation.col_to_field[c]; |
| if (assigned_field_id == -1) { |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(col_builders[c]->AppendNull()); |
| } else { |
| auto it = field_id_to_value_index.find(assigned_field_id); |
| if (PAIMON_UNLIKELY(it == field_id_to_value_index.end())) { |
| return Status::Invalid( |
| fmt::format("MapSharedShreddingBatchConverter: field_id {} assigned to col {} " |
| "but not found in current row", |
| assigned_field_id, c)); |
| } |
| PAIMON_RETURN_NOT_OK_FROM_ARROW( |
| col_builders[c]->AppendArraySlice(*values_array->data(), it->second, 1)); |
| } |
| } |
| return Status::OK(); |
| } |
| |
| Status MapSharedShreddingBatchConverter::AppendOverflow( |
| const std::shared_ptr<arrow::Array>& values_array, const RowAllocation& allocation, |
| const std::unordered_map<int32_t, int64_t>& field_id_to_value_index, |
| arrow::MapBuilder* overflow_builder, arrow::Int32Builder* overflow_key_builder, |
| arrow::ArrayBuilder* overflow_value_builder) const { |
| if (allocation.overflow_fields.empty()) { |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(overflow_builder->AppendNull()); |
| return Status::OK(); |
| } |
| |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(overflow_builder->Append()); |
| for (int32_t overflow_field_id : allocation.overflow_fields) { |
| auto it = field_id_to_value_index.find(overflow_field_id); |
| if (PAIMON_UNLIKELY(it == field_id_to_value_index.end())) { |
| return Status::Invalid(fmt::format( |
| "MapSharedShreddingBatchConverter: overflow field_id {} not found in current row", |
| overflow_field_id)); |
| } |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(overflow_key_builder->Append(overflow_field_id)); |
| PAIMON_RETURN_NOT_OK_FROM_ARROW( |
| overflow_value_builder->AppendArraySlice(*values_array->data(), it->second, 1)); |
| } |
| return Status::OK(); |
| } |
| |
| Result<MapSharedShreddingFieldMeta> MapSharedShreddingBatchConverter::BuildFieldMeta( |
| const std::string& field_name) const { |
| for (const auto& context : contexts_) { |
| if (context.field_name == field_name) { |
| MapSharedShreddingFieldMeta meta; |
| meta.name_to_id = context.dict.GetNameToId(); |
| // Convert set<int32_t> -> vector<int32_t> for field_to_columns |
| for (const auto& [field_id, col_set] : context.allocator->GetFieldToColumns()) { |
| meta.field_to_columns[field_id] = |
| std::vector<int32_t>(col_set.begin(), col_set.end()); |
| } |
| meta.overflow_field_set = context.allocator->GetOverflowFieldSet(); |
| meta.num_columns = context.num_columns; |
| meta.max_row_width = context.allocator->GetMaxRowWidth(); |
| return meta; |
| } |
| } |
| return Status::Invalid(fmt::format( |
| "cannot find field_name '{}' in MapSharedShreddingBatchConverter contexts", field_name)); |
| } |
| |
| Result<int32_t> MapSharedShreddingBatchConverter::GetMaxRowWidth( |
| const std::string& field_name) const { |
| for (const auto& context : contexts_) { |
| if (context.field_name == field_name) { |
| return context.allocator->GetMaxRowWidth(); |
| } |
| } |
| return Status::Invalid(fmt::format( |
| "cannot find field_name '{}' in MapSharedShreddingBatchConverter contexts", field_name)); |
| } |
| |
| const std::vector<std::string>& MapSharedShreddingBatchConverter::GetShreddingColumnNames() const { |
| return shredding_field_names_; |
| } |
| |
| } // namespace paimon |
