| /* |
| * 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/manifest_meta_reader.h" |
| |
| #include <cstdint> |
| #include <string> |
| #include <utility> |
| #include <vector> |
| |
| #include "arrow/api.h" |
| #include "arrow/array/array_base.h" |
| #include "arrow/array/array_nested.h" |
| #include "arrow/array/array_primitive.h" |
| #include "arrow/array/util.h" |
| #include "arrow/c/abi.h" |
| #include "arrow/c/bridge.h" |
| #include "arrow/compute/cast.h" |
| #include "arrow/type.h" |
| #include "arrow/util/checked_cast.h" |
| #include "paimon/common/utils/arrow/mem_utils.h" |
| #include "paimon/common/utils/arrow/status_utils.h" |
| #include "paimon/status.h" |
| |
| namespace paimon { |
| class MemoryPool; |
| |
| ManifestMetaReader::ManifestMetaReader(std::unique_ptr<BatchReader>&& reader, |
| const std::shared_ptr<arrow::DataType>& target_type, |
| const std::shared_ptr<MemoryPool>& pool) |
| : reader_(std::move(reader)), target_type_(target_type), pool_(GetArrowPool(pool)) {} |
| |
| Result<BatchReader::ReadBatch> ManifestMetaReader::NextBatch() { |
| PAIMON_ASSIGN_OR_RAISE(ReadBatch src_result, reader_->NextBatch()); |
| auto& c_array = src_result.first; |
| auto& c_schema = src_result.second; |
| if (!c_array) { |
| return BatchReader::MakeEofBatch(); |
| } |
| PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> arrow_array, |
| arrow::ImportArray(c_array.get(), c_schema.get())); |
| |
| PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> target_array, |
| AlignArrayWithSchema(arrow_array, target_type_, pool_.get())); |
| std::unique_ptr<ArrowArray> target_c_arrow_array = std::make_unique<ArrowArray>(); |
| std::unique_ptr<ArrowSchema> target_c_schema = std::make_unique<ArrowSchema>(); |
| |
| PAIMON_RETURN_NOT_OK_FROM_ARROW( |
| arrow::ExportArray(*target_array, target_c_arrow_array.get(), target_c_schema.get())); |
| return std::make_pair(std::move(target_c_arrow_array), std::move(target_c_schema)); |
| } |
| |
| Result<std::shared_ptr<arrow::Array>> ManifestMetaReader::AlignArrayWithSchema( |
| const std::shared_ptr<arrow::Array>& src_array, |
| const std::shared_ptr<arrow::DataType>& target_type, arrow::MemoryPool* pool) { |
| const auto src_kind = src_array->type()->id(); |
| switch (src_kind) { |
| case arrow::Type::type::LIST: { |
| auto list_src_array = |
| arrow::internal::checked_pointer_cast<arrow::ListArray>(src_array); |
| auto list_target_type = |
| arrow::internal::checked_pointer_cast<arrow::ListType>(target_type); |
| if (!list_target_type) { |
| return Status::Invalid( |
| "Complete non exist field failed, target type cannot cast to a list data " |
| "type"); |
| } |
| PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> converted, |
| AlignArrayWithSchema(list_src_array->values(), |
| list_target_type->value_type(), pool)); |
| auto new_list_type = std::make_shared<arrow::ListType>(converted->type()); |
| return std::make_shared<arrow::ListArray>( |
| new_list_type, list_src_array->length(), list_src_array->value_offsets(), converted, |
| list_src_array->null_bitmap(), list_src_array->null_count(), |
| list_src_array->offset()); |
| } |
| case arrow::Type::type::MAP: { |
| auto map_src_array = arrow::internal::checked_pointer_cast<arrow::MapArray>(src_array); |
| auto map_target_type = |
| arrow::internal::checked_pointer_cast<arrow::MapType>(target_type); |
| if (!map_target_type) { |
| return Status::Invalid( |
| "Complete non exist field failed, target type cannot cast to a map data " |
| "type"); |
| } |
| PAIMON_ASSIGN_OR_RAISE( |
| std::shared_ptr<arrow::Array> key_converted, |
| AlignArrayWithSchema(map_src_array->keys(), map_target_type->key_type(), pool)); |
| PAIMON_ASSIGN_OR_RAISE( |
| std::shared_ptr<arrow::Array> value_converted, |
| AlignArrayWithSchema(map_src_array->items(), map_target_type->item_type(), pool)); |
| auto new_map_type = |
| std::make_shared<arrow::MapType>(key_converted->type(), value_converted->type()); |
| return std::make_shared<arrow::MapArray>( |
| new_map_type, map_src_array->length(), map_src_array->value_offsets(), |
| key_converted, value_converted, map_src_array->null_bitmap(), |
| map_src_array->null_count(), map_src_array->offset()); |
| } |
| case arrow::Type::type::STRUCT: { |
| auto struct_src_array = |
| arrow::internal::checked_pointer_cast<arrow::StructArray>(src_array); |
| auto struct_target_type = |
| arrow::internal::checked_pointer_cast<arrow::StructType>(target_type); |
| if (!struct_target_type) { |
| return Status::Invalid( |
| "Complete non exist field failed, target type cannot cast to a struct data " |
| "type"); |
| } |
| std::vector<std::string> field_names; |
| arrow::ArrayVector converted_array; |
| field_names.reserve(target_type->num_fields()); |
| converted_array.reserve(target_type->num_fields()); |
| for (int32_t i = 0; i < struct_target_type->num_fields(); i++) { |
| const auto& target_field = struct_target_type->field(i); |
| const std::string& target_name = target_field->name(); |
| field_names.push_back(target_name); |
| auto src_field = struct_src_array->GetFieldByName(target_name); |
| if (src_field) { |
| PAIMON_ASSIGN_OR_RAISE( |
| std::shared_ptr<arrow::Array> converted, |
| AlignArrayWithSchema(src_field, target_field->type(), pool)); |
| converted_array.push_back(converted); |
| } else { |
| PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( |
| std::shared_ptr<arrow::Array> null_array, |
| arrow::MakeArrayOfNull(target_field->type(), struct_src_array->length(), |
| pool)); |
| converted_array.push_back(null_array); |
| } |
| } |
| PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( |
| std::shared_ptr<arrow::Array> arrow_array, |
| arrow::StructArray::Make( |
| converted_array, field_names, struct_src_array->null_bitmap(), |
| struct_src_array->null_count(), struct_src_array->offset())); |
| return arrow_array; |
| } |
| case arrow::Type::type::INT32: |
| // cast for avro format, avro store int8,int16,int32 as int32, so need to cast to |
| // correct type |
| if (src_kind != target_type->id()) { |
| arrow::compute::CastOptions cast_options; |
| cast_options.allow_int_overflow = false; |
| auto int32_array = |
| arrow::internal::checked_pointer_cast<arrow::Int32Array>(src_array); |
| PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( |
| std::shared_ptr<arrow::Array> result, |
| arrow::compute::Cast(*int32_array, target_type, cast_options)); |
| return result; |
| } |
| default: |
| return src_array; |
| } |
| } |
| |
| } // namespace paimon |
