HaHaJeff commented on code in PR #224: URL: https://github.com/apache/paimon-cpp/pull/224#discussion_r3855837185
########## src/paimon/core/realtime/prepared_key_value_reader.cpp: ########## @@ -0,0 +1,738 @@ +/* + * 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/realtime/prepared_key_value_reader.h" + +#include <algorithm> +#include <cstdint> +#include <limits> +#include <memory> +#include <optional> +#include <unordered_map> +#include <utility> +#include <vector> + +#include "arrow/array/array_base.h" +#include "arrow/array/array_nested.h" +#include "arrow/array/array_primitive.h" +#include "arrow/c/bridge.h" +#include "arrow/type.h" +#include "fmt/format.h" +#include "paimon/common/data/columnar/columnar_batch_context.h" +#include "paimon/common/data/columnar/columnar_row_ref.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/types/row_kind.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" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/macros.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/status.h" + +namespace paimon { + +namespace { + +constexpr int32_t kValueKindIndex = 0; +constexpr int32_t kSequenceNumberIndex = 1; +constexpr int32_t kRealtimeOffsetIndex = 2; +constexpr int32_t kPreparedValueStartIndex = 3; + +template <typename Reader> +void CloseReaders(const std::vector<std::unique_ptr<Reader>>& readers) { + for (const std::unique_ptr<Reader>& reader : readers) { + if (reader) { + reader->Close(); + } + } +} + +Result<std::shared_ptr<arrow::Array>> AlignArrayByPaimonIds( + const std::shared_ptr<arrow::Array>& array, const std::shared_ptr<arrow::DataType>& read_type, + arrow::MemoryPool* arrow_pool); + +class RealtimeOffsetCoverage { + public: + static Result<std::shared_ptr<RealtimeOffsetCoverage>> Create(const OffsetRange& sealed_offsets, + size_t reader_count) { + if (sealed_offsets.begin < 0 || sealed_offsets.end < sealed_offsets.begin) { + return Status::Invalid("PK real-time store returned an invalid sealed offset range"); + } + return std::shared_ptr<RealtimeOffsetCoverage>( + new RealtimeOffsetCoverage(sealed_offsets, reader_count)); + } + + Status Add(const arrow::Int64Array& offsets) { + for (int64_t row = 0; row < offsets.length(); ++row) { + const int64_t offset = offsets.Value(row); + if (offset < sealed_offsets_.begin || offset >= sealed_offsets_.end) { + return Status::Invalid( + "PK real-time store commit reader offset is outside the sealed range"); + } + min_seen_offset_ = std::min(min_seen_offset_, offset); + max_seen_offset_ = std::max(max_seen_offset_, offset); + ++seen_count_; + } + return Status::OK(); + } + + Status FinishReader() { + ++finished_reader_count_; + if (finished_reader_count_ == reader_count_ && + (seen_count_ != sealed_offsets_.Count() || + (seen_count_ > 0 && (min_seen_offset_ != sealed_offsets_.begin || + max_seen_offset_ != sealed_offsets_.end - 1)))) { + return Status::Invalid( + "PK real-time store commit readers did not cover the sealed range"); + } + return Status::OK(); + } + + private: + RealtimeOffsetCoverage(const OffsetRange& sealed_offsets, size_t reader_count) + : sealed_offsets_(sealed_offsets), reader_count_(reader_count) {} + + OffsetRange sealed_offsets_; + size_t reader_count_; + int64_t min_seen_offset_ = std::numeric_limits<int64_t>::max(); + int64_t max_seen_offset_ = std::numeric_limits<int64_t>::min(); + int64_t seen_count_ = 0; + size_t finished_reader_count_ = 0; +}; + +Status CheckPreparedField(const std::shared_ptr<arrow::Schema>& schema, int32_t field_idx, + const DataField& expected_field) { + if (schema->num_fields() <= field_idx) { + return Status::Invalid(fmt::format("prepared schema missing transport field {} at index {}", + expected_field.Name(), field_idx)); + } + const std::shared_ptr<arrow::Field>& field = schema->field(field_idx); + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(field)); + if (field->name() != expected_field.Name() || !field->type()->Equals(*expected_field.Type()) || + field->nullable() || field_id != expected_field.Id()) { + return Status::Invalid(fmt::format( + "prepared schema field {} must be non-null {}:{} with field id {}, got {}:{} " + "nullable={} field id {}", + field_idx, expected_field.Name(), expected_field.Type()->ToString(), + expected_field.Id(), field->name(), field->type()->ToString(), field->nullable(), + field_id)); + } + return Status::OK(); +} + +Result<int32_t> FindFieldIndexByPaimonId(const arrow::FieldVector& fields, int32_t field_id) { + std::optional<int32_t> matching_index; + for (int32_t i = 0; i < static_cast<int32_t>(fields.size()); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t candidate_id, + NestedProjectionUtils::GetPaimonFieldId(fields[i])); + if (candidate_id == field_id) { + if (matching_index.has_value()) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared schema", field_id)); + } + matching_index = i; + } + } + if (matching_index.has_value()) { + return matching_index.value(); + } + return Status::Invalid(fmt::format("cannot find field id {} in prepared schema", field_id)); +} + +Status ValidateProjectionType(const std::shared_ptr<arrow::DataType>& prepared_type, + const std::shared_ptr<arrow::DataType>& query_type) { + if (prepared_type->id() != query_type->id()) { + return Status::Invalid(fmt::format("prepared value type {} does not match query type {}", + prepared_type->ToString(), query_type->ToString())); + } + switch (query_type->id()) { + case arrow::Type::STRUCT: { + const arrow::FieldVector& prepared_fields = prepared_type->fields(); + for (const std::shared_ptr<arrow::Field>& query_field : query_type->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + PAIMON_ASSIGN_OR_RAISE(int32_t prepared_idx, + FindFieldIndexByPaimonId(prepared_fields, query_id)); + PAIMON_RETURN_NOT_OK(ValidateProjectionType(prepared_fields[prepared_idx]->type(), + query_field->type())); + } + return Status::OK(); + } + case arrow::Type::LIST: + return ValidateProjectionType(prepared_type->field(0)->type(), + query_type->field(0)->type()); + case arrow::Type::MAP: { + const std::shared_ptr<arrow::MapType> prepared_map = + checked_pointer_cast<arrow::MapType>(prepared_type); + const std::shared_ptr<arrow::MapType> query_map = + checked_pointer_cast<arrow::MapType>(query_type); + PAIMON_RETURN_NOT_OK( + ValidateProjectionType(prepared_map->key_type(), query_map->key_type())); + return ValidateProjectionType(prepared_map->item_type(), query_map->item_type()); + } + default: + if (!prepared_type->Equals(*query_type)) { + return Status::Invalid( + fmt::format("prepared leaf type {} does not match query type {}", + prepared_type->ToString(), query_type->ToString())); + } + return Status::OK(); + } +} + +Status ValidateProjectionSchema(const std::shared_ptr<arrow::Schema>& prepared_schema, + const std::shared_ptr<arrow::Schema>& query_schema) { + arrow::FieldVector prepared_value_fields( + prepared_schema->fields().begin() + kPreparedValueStartIndex, + prepared_schema->fields().end()); + for (const std::shared_ptr<arrow::Field>& query_field : query_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + PAIMON_ASSIGN_OR_RAISE(int32_t prepared_idx, + FindFieldIndexByPaimonId(prepared_value_fields, query_id)); + PAIMON_RETURN_NOT_OK(ValidateProjectionType(prepared_value_fields[prepared_idx]->type(), + query_field->type())); + } + return Status::OK(); +} + +Status ValidateExactCommitSchema(const std::shared_ptr<arrow::Schema>& prepared_schema, + const std::shared_ptr<arrow::Schema>& value_schema) { + if (prepared_schema->num_fields() != value_schema->num_fields() + kPreparedValueStartIndex) { + return Status::Invalid("commit requires the exact prepared writer schema"); + } + for (int32_t i = 0; i < value_schema->num_fields(); ++i) { + if (!prepared_schema->field(i + kPreparedValueStartIndex) + ->Equals(value_schema->field(i), true)) { + return Status::Invalid("commit requires the exact prepared writer schema"); + } + } + return Status::OK(); +} + +Result<std::shared_ptr<arrow::Array>> AlignStructArrayByPaimonIds( + const std::shared_ptr<arrow::StructArray>& array, + const std::shared_ptr<arrow::StructType>& read_type, arrow::MemoryPool* arrow_pool) { + const std::shared_ptr<arrow::StructType> data_type = + checked_pointer_cast<arrow::StructType>(array->type()); + std::unordered_map<int32_t, int32_t> data_field_id_to_idx; + data_field_id_to_idx.reserve(data_type->num_fields()); + for (int32_t i = 0; i < data_type->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, + NestedProjectionUtils::GetPaimonFieldId(data_type->field(i))); + if (!data_field_id_to_idx.emplace(field_id, i).second) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared value struct", field_id)); + } + } + + arrow::ArrayVector aligned_arrays; + aligned_arrays.reserve(read_type->num_fields()); + for (const std::shared_ptr<arrow::Field>& read_field : read_type->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t read_field_id, + NestedProjectionUtils::GetPaimonFieldId(read_field)); + auto data_iter = data_field_id_to_idx.find(read_field_id); + if (data_iter == data_field_id_to_idx.end()) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr<arrow::Array> null_child, + arrow::MakeArrayOfNull(read_field->type(), array->offset() + array->length(), + arrow_pool)); + aligned_arrays.push_back(std::move(null_child)); + continue; + } + std::shared_ptr<arrow::Array> child = + arrow::MakeArray(array->data()->child_data[data_iter->second]); + PAIMON_ASSIGN_OR_RAISE(child, AlignArrayByPaimonIds(child, read_field->type(), arrow_pool)); + aligned_arrays.push_back(std::move(child)); + } + + std::shared_ptr<arrow::ArrayData> aligned_data = array->data()->Copy(); + aligned_data->type = read_type; + aligned_data->child_data.clear(); + aligned_data->child_data.reserve(aligned_arrays.size()); + for (const std::shared_ptr<arrow::Array>& aligned_array : aligned_arrays) { + aligned_data->child_data.push_back(aligned_array->data()); + } + return arrow::MakeArray(std::move(aligned_data)); +} + +Result<std::shared_ptr<arrow::Array>> AlignListArrayByPaimonIds( + const std::shared_ptr<arrow::ListArray>& array, + const std::shared_ptr<arrow::ListType>& read_type, arrow::MemoryPool* arrow_pool) { + std::shared_ptr<arrow::Array> values = array->values(); + PAIMON_ASSIGN_OR_RAISE(values, + AlignArrayByPaimonIds(values, read_type->value_type(), arrow_pool)); + std::shared_ptr<arrow::ArrayData> new_data = array->data()->Copy(); + new_data->type = read_type; + new_data->child_data = {values->data()}; + return arrow::MakeArray(new_data); +} + +Result<std::shared_ptr<arrow::Array>> AlignMapArrayByPaimonIds( + const std::shared_ptr<arrow::MapArray>& array, const std::shared_ptr<arrow::MapType>& read_type, + arrow::MemoryPool* arrow_pool) { + std::shared_ptr<arrow::Array> keys = array->keys(); + PAIMON_ASSIGN_OR_RAISE(keys, AlignArrayByPaimonIds(keys, read_type->key_type(), arrow_pool)); + std::shared_ptr<arrow::Array> items = array->items(); + PAIMON_ASSIGN_OR_RAISE(items, AlignArrayByPaimonIds(items, read_type->item_type(), arrow_pool)); + + const std::shared_ptr<arrow::ArrayData>& entries_data = array->data()->child_data[0]; + std::shared_ptr<arrow::ArrayData> new_entries = entries_data->Copy(); + new_entries->type = arrow::struct_({read_type->key_field(), read_type->item_field()}); + new_entries->child_data = {keys->data(), items->data()}; + + std::shared_ptr<arrow::ArrayData> new_data = array->data()->Copy(); + new_data->type = read_type; + new_data->child_data = {std::move(new_entries)}; + return arrow::MakeArray(new_data); +} + +Result<std::shared_ptr<arrow::Array>> AlignArrayByPaimonIds( + const std::shared_ptr<arrow::Array>& array, const std::shared_ptr<arrow::DataType>& read_type, + arrow::MemoryPool* arrow_pool) { + if (array->type()->id() != read_type->id()) { + return Status::Invalid(fmt::format("prepared value type {} does not match query type {}", + array->type()->ToString(), read_type->ToString())); + } + switch (read_type->id()) { + case arrow::Type::STRUCT: + return AlignStructArrayByPaimonIds(checked_pointer_cast<arrow::StructArray>(array), + checked_pointer_cast<arrow::StructType>(read_type), + arrow_pool); + case arrow::Type::LIST: + return AlignListArrayByPaimonIds(checked_pointer_cast<arrow::ListArray>(array), + checked_pointer_cast<arrow::ListType>(read_type), + arrow_pool); + case arrow::Type::MAP: + return AlignMapArrayByPaimonIds(checked_pointer_cast<arrow::MapArray>(array), + checked_pointer_cast<arrow::MapType>(read_type), + arrow_pool); + default: + if (!array->type()->Equals(*read_type)) { + return Status::Invalid( + fmt::format("prepared leaf type {} does not match query type {}", + array->type()->ToString(), read_type->ToString())); + } + return array; + } +} + +Result<arrow::ArrayVector> ProjectFieldsByPaimonIds( + const std::shared_ptr<arrow::StructArray>& data_batch, + const std::shared_ptr<arrow::Schema>& prepared_schema, + const std::shared_ptr<arrow::Schema>& query_schema, arrow::MemoryPool* arrow_pool) { + std::unordered_map<int32_t, int32_t> prepared_field_id_to_idx; + prepared_field_id_to_idx.reserve(prepared_schema->num_fields()); + for (int32_t i = kPreparedValueStartIndex; i < prepared_schema->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, + NestedProjectionUtils::GetPaimonFieldId(prepared_schema->field(i))); + if (!prepared_field_id_to_idx.emplace(field_id, i).second) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared schema", field_id)); + } + } + + arrow::ArrayVector result; + result.reserve(query_schema->num_fields()); + for (const std::shared_ptr<arrow::Field>& query_field : query_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_field_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + auto prepared_iter = prepared_field_id_to_idx.find(query_field_id); + if (prepared_iter == prepared_field_id_to_idx.end()) { + return Status::Invalid( + fmt::format("cannot find field id {} in prepared schema", query_field_id)); + } + std::shared_ptr<arrow::Array> field_array = data_batch->field(prepared_iter->second); + PAIMON_ASSIGN_OR_RAISE(field_array, + AlignArrayByPaimonIds(field_array, query_field->type(), arrow_pool)); + result.push_back(std::move(field_array)); + } + return result; +} + +class PreparedKeyValueReader final : public KeyValueRecordReader { + public: + PreparedKeyValueReader(std::unique_ptr<BatchReader>&& reader, + const std::shared_ptr<arrow::Schema>& prepared_schema, + const std::optional<OffsetRange>& visible_offsets, + const std::shared_ptr<arrow::Schema>& key_schema, + const std::shared_ptr<arrow::Schema>& value_schema, + const std::shared_ptr<FieldsComparator>& key_comparator, + const std::shared_ptr<MemoryPool>& pool, + const std::shared_ptr<RealtimeOffsetCoverage>& offset_coverage) + : reader_(std::move(reader)), + prepared_schema_(prepared_schema), + visible_offsets_(visible_offsets), + key_schema_(key_schema), + value_schema_(value_schema), + key_comparator_(key_comparator), + pool_(pool), + arrow_pool_(GetArrowPool(pool)), + offset_coverage_(offset_coverage) {} + + ~PreparedKeyValueReader() override { + Close(); + } + + class Iterator final : public KeyValueRecordReader::Iterator { + public: + explicit Iterator(PreparedKeyValueReader* reader) : reader_(reader) {} + + Result<bool> HasNext() const override { + return cursor_ < reader_->RowCount(); + } + + Result<KeyValue> Next() override { + if (cursor_ >= reader_->RowCount()) { + return Status::Invalid("No more prepared key values in current iterator"); + } + const int64_t row = reader_->RowAt(cursor_); + std::shared_ptr<InternalRow> key = + std::make_shared<ColumnarRowRef>(reader_->key_ctx_, row); + auto value = std::make_unique<ColumnarRowRef>(reader_->value_ctx_, row); + PAIMON_ASSIGN_OR_RAISE(const RowKind* row_kind, + RowKind::FromByteValue(reader_->row_kind_array_->Value(row))); + int64_t sequence_number = reader_->sequence_number_array_->Value(row); + ++cursor_; + return KeyValue(row_kind, sequence_number, KeyValue::UNKNOWN_LEVEL, std::move(key), + std::move(value)); + } + + private: + PreparedKeyValueReader* reader_; + int64_t cursor_ = 0; + }; + + Result<std::unique_ptr<KeyValueRecordReader::Iterator>> NextBatch() override { + if (first_error_.has_value()) { + return first_error_.value(); + } + Result<std::unique_ptr<KeyValueRecordReader::Iterator>> result = NextBatchImpl(); + if (!result.ok()) { + first_error_ = result.status(); + Close(); + } + return result; + } + + std::shared_ptr<Metrics> GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + if (closed_) { + return; + } + closed_ = true; + ResetBatchState(); + reader_->Close(); + } + + private: + Result<std::unique_ptr<KeyValueRecordReader::Iterator>> NextBatchImpl() { + while (true) { + ResetBatchState(); + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + if (offset_coverage_ && !offset_coverage_finished_) { + offset_coverage_finished_ = true; + PAIMON_RETURN_NOT_OK(offset_coverage_->FinishReader()); + } + return std::unique_ptr<KeyValueRecordReader::Iterator>(); + } + 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("cannot cast prepared batch to StructArray"); + } + std::shared_ptr<arrow::StructArray> data_batch = + checked_pointer_cast<arrow::StructArray>(arrow_array); + Status transport_status = + ValidatePreparedTransportSchema(arrow::schema(data_batch->type()->fields())); + if (!transport_status.ok()) { + return Status::Invalid( + "prepared batch field does not match prepared transport " + "schema: ", + transport_status.ToString()); + } + if (visible_offsets_.has_value()) { + PAIMON_RETURN_NOT_OK(ValidateProjectionSchema( Review Comment: Updated in `a7bc03b7b23c0302ad77d0e7312c0fb4da8a334c`. PK realtime query projection now follows append realtime via `NestedProjectionUtils::AlignArrayToReadType`. The PK-specific schema reconciliation and old-store schema-evolution tests were removed; same-schema nested projection remains supported, including a real disk + sealed-memory + active-memory nested-projection test. Schema changes require recreating `RealtimeContext`/store with the new schema and replaying caller-owned WAL, so old Store null-fill is not a supported contract. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
