| /* |
| * 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/field_mapping.h" |
| |
| #include <algorithm> |
| #include <cassert> |
| #include <cstddef> |
| #include <set> |
| |
| #include "arrow/type.h" |
| #include "fmt/format.h" |
| #include "paimon/common/predicate/compound_predicate_impl.h" |
| #include "paimon/common/predicate/leaf_predicate_impl.h" |
| #include "paimon/common/utils/field_type_utils.h" |
| #include "paimon/common/utils/object_utils.h" |
| #include "paimon/core/casting/cast_executor_factory.h" |
| #include "paimon/core/casting/casting_utils.h" |
| #include "paimon/defs.h" |
| #include "paimon/predicate/literal.h" |
| #include "paimon/predicate/predicate_builder.h" |
| #include "paimon/predicate/predicate_utils.h" |
| #include "paimon/status.h" |
| |
| namespace paimon { |
| |
| Result<std::shared_ptr<arrow::Schema>> FieldMapping::GetPartitionSchema( |
| const std::shared_ptr<arrow::Schema>& schema, const std::vector<std::string>& partition_keys) { |
| arrow::FieldVector partition_fields; |
| partition_fields.reserve(partition_keys.size()); |
| for (const auto& partition_key : partition_keys) { |
| auto field = schema->GetFieldByName(partition_key); |
| if (!field) { |
| return Status::Invalid( |
| fmt::format("get partition schema failed, cannot find partition key {} in schema", |
| partition_key)); |
| } |
| partition_fields.push_back(field); |
| } |
| return arrow::schema(partition_fields); |
| } |
| |
| Result<std::unique_ptr<FieldMappingBuilder>> FieldMappingBuilder::Create( |
| const std::shared_ptr<arrow::Schema>& read_schema, |
| const std::vector<std::string>& partition_keys, const std::shared_ptr<Predicate>& predicate) { |
| PAIMON_ASSIGN_OR_RAISE(std::vector<DataField> read_fields, |
| DataField::ConvertArrowSchemaToDataFields(read_schema)); |
| return std::unique_ptr<FieldMappingBuilder>( |
| new FieldMappingBuilder(read_fields, partition_keys, predicate)); |
| } |
| |
| Result<std::unique_ptr<FieldMapping>> FieldMappingBuilder::CreateFieldMapping( |
| const std::shared_ptr<arrow::Schema>& data_schema) const { |
| PAIMON_ASSIGN_OR_RAISE(std::vector<DataField> data_fields, |
| DataField::ConvertArrowSchemaToDataFields(data_schema)); |
| return CreateFieldMapping(data_fields); |
| } |
| |
| Result<std::unique_ptr<FieldMapping>> FieldMappingBuilder::CreateFieldMapping( |
| const std::vector<DataField>& data_fields) const { |
| // generate non-exist field info |
| std::optional<NonExistFieldInfo> non_exist_field_info = CreateNonExistFieldInfo(data_fields); |
| |
| // generate exist field info |
| ExistFieldInfo exist_field_info = CreateExistFieldInfo(data_fields); |
| |
| // key: partition key, value: partition idx |
| std::map<std::string, int32_t> partition_key_to_idx = |
| ObjectUtils::CreateIdentifierToIndexMap(partition_keys_); |
| |
| PAIMON_ASSIGN_OR_RAISE( |
| NonPartitionInfo non_partition_info, |
| CreateNonPartitionInfo(data_fields, exist_field_info, partition_key_to_idx)); |
| PAIMON_ASSIGN_OR_RAISE(std::optional<PartitionInfo> partition_info, |
| CreatePartitionInfo(exist_field_info, partition_key_to_idx)); |
| return std::make_unique<FieldMapping>(partition_info, non_partition_info, non_exist_field_info); |
| } |
| |
| ExistFieldInfo FieldMappingBuilder::CreateExistFieldInfo( |
| const std::vector<DataField>& data_fields) const { |
| // key:field id, value: {target_idx, read field} |
| std::map<int32_t, std::pair<int32_t, DataField>> field_id_to_read_fields; |
| for (size_t i = 0; i < read_fields_.size(); i++) { |
| const auto& read_field = read_fields_[i]; |
| field_id_to_read_fields.emplace(read_field.Id(), std::make_pair(i, read_field)); |
| } |
| |
| ExistFieldInfo exist_field_info; |
| for (const auto& data_field : data_fields) { |
| auto iter = field_id_to_read_fields.find(data_field.Id()); |
| if (iter != field_id_to_read_fields.end()) { |
| const auto& [target_idx, read_field] = iter->second; |
| exist_field_info.exist_read_schema.push_back(read_field); |
| exist_field_info.exist_data_schema.push_back(data_field); |
| exist_field_info.idx_in_target_read_schema.push_back(target_idx); |
| } |
| } |
| return exist_field_info; |
| } |
| |
| std::optional<NonExistFieldInfo> FieldMappingBuilder::CreateNonExistFieldInfo( |
| const std::vector<DataField>& data_fields) const { |
| // key: field id, value: data field |
| std::map<int32_t, DataField> field_id_to_data_fields; |
| for (const auto& data_field : data_fields) { |
| field_id_to_data_fields.emplace(data_field.Id(), data_field); |
| } |
| |
| NonExistFieldInfo non_exist_field_info; |
| for (size_t i = 0; i < read_fields_.size(); i++) { |
| const auto& read_field = read_fields_[i]; |
| auto iter = field_id_to_data_fields.find(read_field.Id()); |
| if (iter == field_id_to_data_fields.end()) { |
| non_exist_field_info.non_exist_read_schema.push_back(read_field); |
| non_exist_field_info.idx_in_target_read_schema.push_back(i); |
| } |
| } |
| if (non_exist_field_info.idx_in_target_read_schema.empty()) { |
| return std::nullopt; |
| } |
| return non_exist_field_info; |
| } |
| |
| Result<std::vector<std::shared_ptr<CastExecutor>>> FieldMappingBuilder::CreateDataCastExecutors( |
| const std::vector<DataField>& read_fields, const std::vector<DataField>& data_fields) { |
| assert(read_fields.size() == data_fields.size()); |
| std::vector<std::shared_ptr<CastExecutor>> cast_executors; |
| cast_executors.reserve(read_fields.size()); |
| for (size_t i = 0; i < read_fields.size(); i++) { |
| PAIMON_ASSIGN_OR_RAISE(FieldType read_type, |
| FieldTypeUtils::ConvertToFieldType(read_fields[i].Type()->id())); |
| PAIMON_ASSIGN_OR_RAISE(FieldType data_type, |
| FieldTypeUtils::ConvertToFieldType(data_fields[i].Type()->id())); |
| |
| if (!read_fields[i].Type()->Equals(data_fields[i].Type())) { |
| if (read_type == FieldType::MAP || read_type == FieldType::ARRAY || |
| read_type == FieldType::STRUCT) { |
| return Status::Invalid("Only support column type evolution in atomic data type."); |
| } |
| auto executor_factory = CastExecutorFactory::GetCastExecutorFactory(); |
| auto cast_executor = |
| executor_factory->GetCastExecutor(/*src=*/data_type, /*target=*/read_type); |
| if (!cast_executor) { |
| return Status::Invalid( |
| fmt::format("CreateDataCastExecutors failed: cannot find cast " |
| "executor for field {} from type {} to {}", |
| read_fields[i].Name(), FieldTypeUtils::FieldTypeToString(data_type), |
| FieldTypeUtils::FieldTypeToString(read_type))); |
| } |
| cast_executors.push_back(cast_executor); |
| } else { |
| cast_executors.push_back(nullptr); |
| } |
| } |
| return cast_executors; |
| } |
| |
| Result<NonPartitionInfo> FieldMappingBuilder::CreateNonPartitionInfo( |
| const std::vector<DataField>& data_fields, const ExistFieldInfo& exist_field_info, |
| const std::map<std::string, int32_t>& partition_keys) const { |
| NonPartitionInfo non_partition_info; |
| for (size_t i = 0; i < exist_field_info.exist_data_schema.size(); i++) { |
| const auto& data_field = exist_field_info.exist_data_schema[i]; |
| const auto& read_field = exist_field_info.exist_read_schema[i]; |
| auto iter = partition_keys.find(read_field.Name()); |
| if (iter == partition_keys.end()) { |
| non_partition_info.non_partition_read_schema.push_back(read_field); |
| non_partition_info.non_partition_data_schema.push_back(data_field); |
| non_partition_info.idx_in_target_read_schema.push_back( |
| exist_field_info.idx_in_target_read_schema[i]); |
| } |
| } |
| PAIMON_ASSIGN_OR_RAISE(non_partition_info.cast_executors, |
| CreateDataCastExecutors(non_partition_info.non_partition_read_schema, |
| non_partition_info.non_partition_data_schema)); |
| // prepare predicate: exclude the push down predicate with partition key and non-exist fields |
| PAIMON_ASSIGN_OR_RAISE(non_partition_info.non_partition_filter, CreateDataFilters(data_fields)); |
| return non_partition_info; |
| } |
| |
| Result<std::shared_ptr<Predicate>> FieldMappingBuilder::CreateDataFilters( |
| const std::vector<DataField>& data_fields) const { |
| if (!predicate_) { |
| return predicate_; |
| } |
| // data schema mapping: id -> {data field idx, data field} |
| std::map<int32_t, std::pair<int32_t, DataField>> field_id_to_data_fields; |
| for (size_t i = 0; i < data_fields.size(); i++) { |
| const auto& data_field = data_fields[i]; |
| field_id_to_data_fields.emplace(data_field.Id(), std::make_pair(i, data_field)); |
| } |
| // table schema mapping: name -> read field |
| std::map<std::string, DataField> field_name_to_read_fields; |
| for (const auto& read_field : read_fields_) { |
| field_name_to_read_fields.emplace(read_field.Name(), read_field); |
| } |
| // reconstruct predicate (e.g., modify field idx, cast literal) |
| auto split_predicates = PredicateUtils::SplitAnd(predicate_); |
| std::vector<std::shared_ptr<Predicate>> converted_predicates; |
| converted_predicates.reserve(split_predicates.size()); |
| for (const auto& predicate : split_predicates) { |
| PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<Predicate>> converted, |
| ReconstructPredicateWithDataFields( |
| predicate, field_name_to_read_fields, field_id_to_data_fields)); |
| if (converted != std::nullopt) { |
| converted_predicates.push_back(converted.value()); |
| } |
| } |
| // exclude partition fields in predicate |
| std::set<std::string> partition_key_field_name_set(partition_keys_.begin(), |
| partition_keys_.end()); |
| PAIMON_ASSIGN_OR_RAISE(std::vector<std::shared_ptr<Predicate>> remain_predicates, |
| PredicateUtils::ExcludePredicateWithFields( |
| converted_predicates, partition_key_field_name_set)); |
| if (remain_predicates.empty()) { |
| return std::shared_ptr<Predicate>(); |
| } |
| return PredicateBuilder::And(remain_predicates); |
| } |
| |
| Result<std::optional<std::shared_ptr<Predicate>>> FieldMappingBuilder::ReconstructLeafPredicate( |
| const std::shared_ptr<LeafPredicateImpl>& leaf_predicate, |
| const std::map<std::string, DataField>& field_name_to_read_fields, |
| const std::map<int32_t, std::pair<int32_t, DataField>>& field_id_to_data_fields) { |
| auto read_field_iter = field_name_to_read_fields.find(leaf_predicate->FieldName()); |
| if (read_field_iter == field_name_to_read_fields.end()) { |
| return Status::Invalid( |
| fmt::format("invalid predicate, cannot find field {}", leaf_predicate->FieldName())); |
| } |
| auto data_field_iter = field_id_to_data_fields.find(read_field_iter->second.Id()); |
| if (data_field_iter == field_id_to_data_fields.end()) { |
| return std::optional<std::shared_ptr<Predicate>>(); |
| } |
| std::vector<Literal> new_literals; |
| PAIMON_ASSIGN_OR_RAISE(FieldType read_type, FieldTypeUtils::ConvertToFieldType( |
| read_field_iter->second.Type()->id())); |
| PAIMON_ASSIGN_OR_RAISE(FieldType data_type, FieldTypeUtils::ConvertToFieldType( |
| data_field_iter->second.second.Type()->id())); |
| if (read_type != data_type || |
| (!read_field_iter->second.Type()->Equals(data_field_iter->second.second.Type()))) { |
| // read type and data type are different, only push down integer predicates |
| if (FieldTypeUtils::IsIntegerNumeric(read_type) && |
| FieldTypeUtils::IsIntegerNumeric(data_type) && |
| FieldTypeUtils::IntegerScaleLargerThan(read_type, data_type)) { |
| auto executor_factory = CastExecutorFactory::GetCastExecutorFactory(); |
| auto cast_executor = |
| executor_factory->GetCastExecutor(/*src=*/read_type, /*target=*/data_type); |
| if (!cast_executor) { |
| return Status::Invalid(fmt::format( |
| "ReconstructLeafPredicate failed: cannot find cast " |
| "executor for field {} from type {} to {}", |
| read_field_iter->second.Name(), FieldTypeUtils::FieldTypeToString(read_type), |
| FieldTypeUtils::FieldTypeToString(data_type))); |
| } |
| new_literals.reserve(leaf_predicate->Literals().size()); |
| for (const auto& literal : leaf_predicate->Literals()) { |
| PAIMON_ASSIGN_OR_RAISE( |
| Literal casted, |
| cast_executor->Cast(literal, data_field_iter->second.second.Type())); |
| // ignore if any literal is overflowed. |
| if (CastingUtils::IsIntegerLiteralCastedOverflow(literal, casted)) { |
| return std::optional<std::shared_ptr<Predicate>>(); |
| } |
| new_literals.push_back(casted); |
| } |
| } else { |
| return std::optional<std::shared_ptr<Predicate>>(); |
| } |
| } else { |
| new_literals = leaf_predicate->Literals(); |
| } |
| |
| const auto& [data_field_idx, data_field] = data_field_iter->second; |
| return std::optional<std::shared_ptr<Predicate>>(std::make_shared<LeafPredicateImpl>( |
| leaf_predicate->GetLeafFunction(), /*data field idx*/ data_field_idx, |
| /*data field name*/ data_field.Name(), |
| /*data field type*/ data_type, new_literals)); |
| } |
| |
| Result<std::optional<std::shared_ptr<Predicate>>> |
| FieldMappingBuilder::ReconstructPredicateWithDataFields( |
| const std::shared_ptr<Predicate>& predicate, |
| const std::map<std::string, DataField>& field_name_to_read_fields, |
| const std::map<int32_t, std::pair<int32_t, DataField>>& field_id_to_data_fields) { |
| if (auto leaf_predicate = std::dynamic_pointer_cast<LeafPredicateImpl>(predicate)) { |
| return ReconstructLeafPredicate(leaf_predicate, field_name_to_read_fields, |
| field_id_to_data_fields); |
| } else if (auto compound_predicate = |
| std::dynamic_pointer_cast<CompoundPredicateImpl>(predicate)) { |
| std::vector<std::shared_ptr<Predicate>> converted_children; |
| for (const auto& child : compound_predicate->Children()) { |
| PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<Predicate>> converted_child, |
| ReconstructPredicateWithDataFields( |
| child, field_name_to_read_fields, field_id_to_data_fields)); |
| if (converted_child == std::nullopt) { |
| return std::optional<std::shared_ptr<Predicate>>(); |
| } |
| assert(converted_child.value()); |
| converted_children.push_back(converted_child.value()); |
| } |
| return std::optional<std::shared_ptr<Predicate>>( |
| compound_predicate->NewCompoundPredicate(converted_children)); |
| } |
| return Status::Invalid( |
| "invalid predicate, cannot convert to leaf predicate or compound predicate"); |
| } |
| |
| Result<std::optional<PartitionInfo>> FieldMappingBuilder::CreatePartitionInfo( |
| const ExistFieldInfo& exist_field_info, |
| const std::map<std::string, int32_t>& partition_keys) const { |
| if (partition_keys.empty()) { |
| return std::optional<PartitionInfo>(); |
| } |
| PartitionInfo partition_info; |
| for (size_t i = 0; i < exist_field_info.exist_read_schema.size(); i++) { |
| const auto& read_field = exist_field_info.exist_read_schema[i]; |
| auto iter = partition_keys.find(read_field.Name()); |
| if (iter != partition_keys.end()) { |
| partition_info.partition_read_schema.push_back(read_field); |
| partition_info.idx_in_target_read_schema.push_back( |
| exist_field_info.idx_in_target_read_schema[i]); |
| partition_info.idx_in_partition.push_back(iter->second); |
| } |
| } |
| if (partition_info.idx_in_target_read_schema.empty()) { |
| return std::optional<PartitionInfo>(); |
| } |
| PAIMON_ASSIGN_OR_RAISE(partition_info.partition_filter, |
| PredicateUtils::CreatePickedFieldFilter(predicate_, partition_keys)); |
| return std::optional<PartitionInfo>(partition_info); |
| } |
| |
| } // namespace paimon |