lxy-9602 commented on code in PR #224:
URL: https://github.com/apache/paimon-cpp/pull/224#discussion_r3870547951
##########
src/paimon/common/table/special_fields.h:
##########
@@ -35,6 +36,10 @@ struct SpecialFields {
static constexpr char KEY_FIELD_PREFIX[] = "_KEY_";
static constexpr int32_t KEY_VALUE_SPECIAL_FIELD_COUNT = 2;
+ static constexpr int32_t kPreparedKeyValueValueKindIndex = 0;
+ static constexpr int32_t kPreparedKeyValueSequenceNumberIndex = 1;
+ static constexpr int32_t kPreparedKeyValueRealtimeOffsetIndex = 2;
+ static constexpr int32_t kPreparedKeyValueValueStartIndex = 3;
Review Comment:
Hmm, maybe we shouldn’t put this `kPreparedKeyxxxx` into `SpecialFields` ?
##########
src/paimon/common/table/special_fields.h:
##########
@@ -85,6 +97,16 @@ struct SpecialFields {
target_fields.insert(target_fields.end(), schema->fields().begin(),
schema->fields().end());
return arrow::schema(target_fields);
}
+
+ static std::shared_ptr<arrow::Schema> PreparedKeyValueSchema(
+ const arrow::FieldVector& value_fields) {
Review Comment:
Could we rename this function? The current name is a bit misleading and
makes it hard to tell that it’s intended for the realtime use case rather than
general usage.
##########
src/paimon/core/realtime/prepared_key_value_reader.h:
##########
@@ -0,0 +1,55 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "arrow/type_fwd.h"
+#include "paimon/core/io/key_value_record_reader.h"
+#include "paimon/realtime/offset_range.h"
+#include "paimon/result.h"
+
+namespace paimon {
+class BatchReader;
+class MemoryPool;
+
+class PreparedKeyValueReaderFactory {
+ public:
Review Comment:
There are quite a few places in the code using names like `PreparedXXX`, but
that term feels a bit too generic and doesn’t clearly convey the connection to
streaming primary-key logic. Could we use a more explicit name?
##########
src/paimon/core/operation/merge_file_split_read.cpp:
##########
@@ -453,38 +611,90 @@ Result<std::unique_ptr<BatchReader>>
MergeFileSplitRead::CreateReaderForSection(
} else {
predicate = context_->GetPredicate();
}
- PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<SortMergeReader> sort_merge_reader,
- CreateSortMergeReaderForSection(section, partition,
dv_factory,
- predicate,
data_file_path_factory,
-
/*drop_delete=*/!force_keep_delete_));
- // KeyValueProjectionReader converts KeyValue objects to arrow array
according to projection
- if (!context_->EnableMultiThreadRowToBatch()) {
- return KeyValueProjectionReader::Create(std::move(sort_merge_reader),
raw_read_schema_,
- projection_,
options_.GetReadBatchSize(), pool_);
- }
- int32_t thread_number = context_->GetRowToBatchThreadNumber();
- assert(thread_number > 0);
- return std::make_unique<AsyncKeyValueProjectionReader>(
- std::move(sort_merge_reader), raw_read_schema_, projection_,
options_.GetReadBatchSize(),
- thread_number, pool_);
+ PAIMON_ASSIGN_OR_RAISE(
+ std::unique_ptr<SortMergeReader> sort_merge_reader,
+ CreateSortMergeReaderForSection(section, partition, dv_factory,
predicate,
+ data_file_path_factory,
/*drop_delete=*/false));
+ return CreateProjectedReader(std::move(sort_merge_reader),
/*predicate=*/nullptr,
+ /*complete_row_kind=*/false);
}
-Result<std::unique_ptr<SortMergeReader>>
MergeFileSplitRead::CreateSortMergeReaderForSection(
+Status MergeFileSplitRead::CreateDiskSections(
+ const std::vector<std::shared_ptr<DataFileMeta>>& data_files,
+ const std::vector<std::optional<DeletionFile>>& deletion_files,
+ DeletionVector::Factory* dv_factory, std::vector<std::vector<SortedRun>>*
sections) const {
+ *dv_factory = DeletionVector::CreateFactory(
+ options_.GetFileSystem(),
DeletionVector::CreateDeletionFileMap(data_files, deletion_files),
+ pool_);
+ *sections = IntervalPartition(data_files, key_comparator_).Partition();
+ return Status::OK();
+}
+
+Result<std::vector<std::unique_ptr<KeyValueRecordReader>>>
+MergeFileSplitRead::CreateRecordReadersForSection(
const std::vector<SortedRun>& section, const BinaryRow& partition,
DeletionVector::Factory dv_factory, const std::shared_ptr<Predicate>&
predicate,
- const std::shared_ptr<DataFilePathFactory>& data_file_path_factory, bool
drop_delete) {
- // with overlap in one section
+ const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const {
Review Comment:
I understand this was part of the refactor, but could we keep the original
comments instead of removing them?
##########
src/paimon/core/realtime/realtime_primary_key_writer.cpp:
##########
@@ -0,0 +1,279 @@
+/*
+ * 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/realtime_primary_key_writer.h"
+
+#include <limits>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "arrow/api.h"
+#include "arrow/c/bridge.h"
+#include "arrow/compute/api.h"
+#include "paimon/common/table/special_fields.h"
+#include "paimon/common/types/row_kind.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/core/io/merged_key_value_record_reader.h"
+#include "paimon/core/mergetree/compact/deduplicate_merge_function.h"
+#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h"
+#include "paimon/core/mergetree/merge_tree_writer.h"
+#include "paimon/core/realtime/prepared_key_value_reader.h"
+#include "paimon/core/realtime/realtime_context_impl.h"
+#include "paimon/core/utils/commit_increment.h"
+#include "paimon/macros.h"
+
+namespace paimon {
+
+namespace {
+
+Result<std::shared_ptr<arrow::StructArray>> PrepareBatch(
+ std::unique_ptr<RecordBatch>&& batch, const
std::shared_ptr<arrow::Schema>& write_schema,
+ const std::shared_ptr<arrow::Schema>& prepared_schema,
+ const std::vector<std::string>& trimmed_primary_keys, int64_t
first_sequence_number,
+ int64_t first_offset, arrow::MemoryPool* arrow_pool) {
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ std::shared_ptr<arrow::Array> input,
+ arrow::ImportArray(batch->GetData(),
arrow::struct_(write_schema->fields())));
+ if (!input || input->type_id() != arrow::Type::STRUCT) {
+ return Status::Invalid("PK real-time write data is not a StructArray");
+ }
+ std::shared_ptr<arrow::StructArray> values =
checked_pointer_cast<arrow::StructArray>(input);
+ const int64_t count = values->length();
+ arrow::Int8Builder kinds(arrow_pool);
+ arrow::Int64Builder sequences(arrow_pool);
+ arrow::Int64Builder offsets(arrow_pool);
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Reserve(count));
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Reserve(count));
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Reserve(count));
+ const std::vector<RecordBatch::RowKind>& row_kinds = batch->GetRowKind();
+ for (int64_t row = 0; row < count; ++row) {
+ const RecordBatch::RowKind kind =
+ row_kinds.empty() ? RecordBatch::RowKind::INSERT : row_kinds[row];
+ kinds.UnsafeAppend(static_cast<int8_t>(kind));
+ sequences.UnsafeAppend(first_sequence_number + row);
+ offsets.UnsafeAppend(first_offset + row);
+ }
+ std::shared_ptr<arrow::Array> kind_array;
+ std::shared_ptr<arrow::Array> sequence_array;
+ std::shared_ptr<arrow::Array> offset_array;
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Finish(&kind_array));
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Finish(&sequence_array));
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Finish(&offset_array));
+ arrow::ArrayVector columns = {std::move(kind_array),
std::move(sequence_array),
+ std::move(offset_array)};
+ columns.insert(columns.end(), values->fields().begin(),
values->fields().end());
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ std::shared_ptr<arrow::StructArray> prepared,
+ arrow::StructArray::Make(std::move(columns),
prepared_schema->fields()));
+
+ std::vector<arrow::compute::SortKey> sort_keys;
+ sort_keys.reserve(trimmed_primary_keys.size() + 1);
+ for (const std::string& key : trimmed_primary_keys) {
+ sort_keys.emplace_back(key, arrow::compute::SortOrder::Ascending);
+ }
+ sort_keys.emplace_back(SpecialFields::SequenceNumber().Name(),
+ arrow::compute::SortOrder::Ascending);
+ arrow::compute::ExecContext context(arrow_pool);
+ arrow::compute::SortOptions options(sort_keys,
arrow::compute::NullPlacement::AtStart);
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ arrow::Datum indices,
+ arrow::compute::SortIndices(arrow::Datum(prepared), options,
&context));
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ arrow::Datum sorted,
+ arrow::compute::Take(arrow::Datum(prepared), indices,
+ arrow::compute::TakeOptions::NoBoundsCheck(),
&context));
+ return checked_pointer_cast<arrow::StructArray>(sorted.make_array());
+}
+
+} // namespace
+
+Result<std::shared_ptr<RealtimePrimaryKeyWriter>>
RealtimePrimaryKeyWriter::Create(
+ const std::map<std::string, std::string>& partition, int32_t bucket,
+ const std::shared_ptr<arrow::Schema>& write_schema,
+ const std::shared_ptr<arrow::Schema>& prepared_schema,
+ const std::vector<std::string>& trimmed_primary_keys,
+ const std::shared_ptr<FieldsComparator>& key_comparator,
+ const std::shared_ptr<RealtimeContextImpl>& realtime_context,
+ const RealtimeStoreState& store_state, int64_t
restored_max_sequence_number,
+ const std::shared_ptr<MergeTreeWriter>& merge_tree_writer,
+ const std::shared_ptr<MemoryPool>& memory_pool) {
+ if (restored_max_sequence_number < -1 ||
+ restored_max_sequence_number == std::numeric_limits<int64_t>::max()) {
+ return Status::Invalid("PK restored sequence number is invalid");
+ }
+ arrow::FieldVector key_fields;
+ key_fields.reserve(trimmed_primary_keys.size());
+ for (const std::string& key : trimmed_primary_keys) {
+ std::shared_ptr<arrow::Field> field =
write_schema->GetFieldByName(key);
+ if (!field) {
+ return Status::Invalid("PK field is missing from write schema: ",
key);
+ }
+ key_fields.push_back(std::move(field));
+ }
+ const RealtimePartitionBucket partition_bucket(partition, bucket);
+ PAIMON_ASSIGN_OR_RAISE(int64_t initial_max_sequence_number,
+
realtime_context->AdvanceMaterializedMaxSequenceNumber(
+ partition_bucket,
restored_max_sequence_number));
+ return std::shared_ptr<RealtimePrimaryKeyWriter>(new
RealtimePrimaryKeyWriter(
+ store_state.store, merge_tree_writer, realtime_context,
partition_bucket, write_schema,
+ prepared_schema, arrow::schema(std::move(key_fields)),
trimmed_primary_keys, key_comparator,
+ store_state.initial_offset, initial_max_sequence_number, memory_pool));
+}
+
+RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter(
+ const std::shared_ptr<RealtimeStore>& realtime_store,
+ const std::shared_ptr<MergeTreeWriter>& merge_tree_writer,
+ const std::shared_ptr<RealtimeContextImpl>& realtime_context,
+ const RealtimePartitionBucket& partition_bucket,
+ const std::shared_ptr<arrow::Schema>& write_schema,
+ const std::shared_ptr<arrow::Schema>& prepared_schema,
+ const std::shared_ptr<arrow::Schema>& key_schema,
+ const std::vector<std::string>& trimmed_primary_keys,
+ const std::shared_ptr<FieldsComparator>& key_comparator, int64_t
next_offset,
+ int64_t last_sequence_number, const std::shared_ptr<MemoryPool>&
memory_pool)
+ : memory_pool_(memory_pool),
+ arrow_pool_(GetArrowPool(memory_pool)),
+ realtime_store_(realtime_store),
+ merge_tree_writer_(merge_tree_writer),
+ realtime_context_(realtime_context),
+ partition_bucket_(partition_bucket),
+ write_schema_(write_schema),
+ prepared_schema_(prepared_schema),
+ key_schema_(key_schema),
+ trimmed_primary_keys_(trimmed_primary_keys),
+ key_comparator_(key_comparator),
+ next_offset_(next_offset),
+ last_sequence_number_(last_sequence_number) {}
+
+Status RealtimePrimaryKeyWriter::Write(std::unique_ptr<RecordBatch>&& batch) {
+ if (!batch || !batch->GetData()) {
+ return Status::Invalid("PK real-time write batch is null");
+ }
+ const int64_t count = batch->GetData()->length;
+ if (count == 0) {
+ return Status::OK();
+ }
+ const std::vector<RecordBatch::RowKind>& row_kinds = batch->GetRowKind();
+ if (!row_kinds.empty() && static_cast<int64_t>(row_kinds.size()) != count)
{
+ return Status::Invalid("PK real-time row-kind count does not match
batch row count");
+ }
+ for (RecordBatch::RowKind row_kind : row_kinds) {
+ PAIMON_ASSIGN_OR_RAISE(const RowKind* validated,
+
RowKind::FromByteValue(static_cast<int8_t>(row_kind)));
+ static_cast<void>(validated);
+ }
+ std::lock_guard<std::mutex> lock(realtime_store_mutex_);
+ if (count > std::numeric_limits<int64_t>::max() - next_offset_) {
+ return Status::Invalid("real-time offset range exceeds INT64_MAX");
+ }
+ // Reserve INT64_MAX as the exhausted sequence-number sentinel.
+ if (last_sequence_number_ >= std::numeric_limits<int64_t>::max() - count) {
+ return Status::Invalid("PK sequence range exceeds INT64_MAX");
+ }
+ const int64_t first_sequence = last_sequence_number_ + 1;
+ PAIMON_ASSIGN_OR_RAISE(
+ std::shared_ptr<arrow::StructArray> prepared,
+ PrepareBatch(std::move(batch), write_schema_, prepared_schema_,
trimmed_primary_keys_,
+ first_sequence, next_offset_, arrow_pool_.get()));
+ auto output = std::make_unique<ArrowArray>();
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*prepared,
output.get()));
+ PAIMON_RETURN_NOT_OK(RetainArrowArrayMemoryPool(output.get(),
arrow_pool_));
+ RecordBatchBuilder builder(output.get());
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<RecordBatch> prepared_batch,
builder.Finish());
+ PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{
+ std::move(prepared_batch), OffsetRange(next_offset_, next_offset_ +
count)}));
+ next_offset_ += count;
+ last_sequence_number_ += count;
+
PAIMON_RETURN_NOT_OK(realtime_context_->AdvanceMaterializedMaxSequenceNumber(
+ partition_bucket_, last_sequence_number_));
+ return Status::OK();
+}
+
+Result<CommitIncrement> RealtimePrimaryKeyWriter::PrepareCommit(bool
wait_compaction) {
+ std::lock_guard<std::mutex> prepare_lock(prepare_mutex_);
+ std::optional<std::shared_ptr<RealtimeSegmentHandle>> segment;
+ {
+ std::lock_guard<std::mutex> store_lock(realtime_store_mutex_);
+
PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<RealtimeSegmentHandle>>
sealed,
+ realtime_store_->SealForCommit());
+ segment = std::move(sealed);
+ }
+ if (segment && !segment.value()) {
+ return Status::Invalid("PK real-time store sealed a null segment");
+ }
+ std::optional<OffsetRange> sealed_range;
+ if (segment) {
+ sealed_range = segment.value()->GetOffsetRange();
+ if (sealed_range->begin < 0 || sealed_range->end <
sealed_range->begin) {
+ return Status::Invalid("PK real-time store returned an invalid
sealed offset range");
+ }
+ PAIMON_RETURN_NOT_OK(FlushSegment(segment.value(),
sealed_range.value()));
+ }
+ PAIMON_ASSIGN_OR_RAISE(CommitIncrement increment,
+ merge_tree_writer_->PrepareCommit(wait_compaction));
+ if (segment) {
+ increment.SetRealtimeOffsetRange(sealed_range.value());
+ }
+ return increment;
+}
+
+Status RealtimePrimaryKeyWriter::FlushSegment(const
std::shared_ptr<RealtimeSegmentHandle>& segment,
+ const OffsetRange&
sealed_offsets) {
+ PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<BatchReader>> readers,
+ realtime_store_->CreateCommitReaders(segment));
+ PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<KeyValueRecordReader>>
prepared_readers,
+ PreparedKeyValueReaderFactory::CreateForCommit(
+ std::move(readers), prepared_schema_,
sealed_offsets, key_schema_,
+ write_schema_, memory_pool_));
+ std::vector<std::unique_ptr<KeyValueRecordReader>> sorted_readers;
+ sorted_readers.reserve(prepared_readers.size());
+ for (std::unique_ptr<KeyValueRecordReader>& prepared_reader :
prepared_readers) {
+ auto merge_function =
std::make_unique<DeduplicateMergeFunction>(/*ignore_delete=*/false);
+ sorted_readers.push_back(std::make_unique<MergedKeyValueRecordReader>(
Review Comment:
I’d prefer to use `CreateMergeFunction` instead of directly instantiating
`DeduplicateMergeFunction`. We can validate at the writer construction that the
merge function must currently be `deduplicate`. That way, when we support other
merge functions in the future, we can simply relax that validation, rather than
having to hunt down and update internal hardcoded usages.
##########
src/paimon/core/realtime/prepared_key_value_reader.cpp:
##########
@@ -0,0 +1,577 @@
+/*
+ * 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 <cstdint>
+#include <memory>
+#include <optional>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "arrow/array/array_base.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/status_utils.h"
+#include "paimon/common/utils/checked_cast.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"
+#include "paimon/utils/roaring_bitmap64.h"
+
+namespace paimon {
+
+namespace {
+
+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();
+ }
+ }
+}
+
+class RealtimeOffsetCoverage {
+ public:
+ static Result<std::shared_ptr<RealtimeOffsetCoverage>> Create(const
OffsetRange& offsets,
+ size_t
reader_count,
+ bool
allow_committed_prefix) {
+ if (offsets.begin < 0 || offsets.end < offsets.begin) {
+ return Status::Invalid("PK real-time store returned an invalid
offset range");
+ }
+ return std::shared_ptr<RealtimeOffsetCoverage>(
+ new RealtimeOffsetCoverage(offsets, reader_count,
allow_committed_prefix));
+ }
+
+ Status Add(const arrow::Int64Array& offsets) {
+ for (int64_t row = 0; row < offsets.length(); ++row) {
+ const int64_t offset = offsets.Value(row);
+ if (allow_committed_prefix_ && offset < 0) {
+ return Status::Invalid("PK real-time store reader offset must
be non-negative");
+ }
+ if (allow_committed_prefix_ && offset < offsets_.begin) {
+ continue;
+ }
+ if (offset < offsets_.begin || offset >= offsets_.end) {
+ return Status::Invalid(
+ allow_committed_prefix_
+ ? "PK real-time store query reader offset is outside
the visible range"
+ : "PK real-time store commit reader offset is outside
the sealed range");
+ }
+ if (!seen_offsets_.CheckedAdd(offset)) {
+ return CoverageError();
+ }
+ }
+ return Status::OK();
+ }
+
+ Status FinishReader() {
+ ++finished_reader_count_;
+ if (finished_reader_count_ == reader_count_ &&
+ seen_offsets_.Cardinality() != offsets_.Count()) {
+ return CoverageError();
+ }
+ return Status::OK();
+ }
+
+ private:
+ RealtimeOffsetCoverage(const OffsetRange& offsets, size_t reader_count,
+ bool allow_committed_prefix)
+ : offsets_(offsets),
+ reader_count_(reader_count),
+ allow_committed_prefix_(allow_committed_prefix) {}
+
+ Status CoverageError() const {
+ return Status::Invalid(
+ allow_committed_prefix_
+ ? "PK real-time store query readers did not cover the visible
range"
+ : "PK real-time store commit readers did not cover the sealed
range");
+ }
+
+ OffsetRange offsets_;
+ size_t reader_count_;
+ bool allow_committed_prefix_;
+ RoaringBitmap64 seen_offsets_;
+ 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<std::vector<int32_t>> ResolveFieldIndexes(
+ const std::shared_ptr<arrow::Schema>& prepared_schema,
+ const std::unordered_map<int32_t, int32_t>& field_indexes,
+ const std::shared_ptr<arrow::Schema>& row_schema) {
+ std::vector<int32_t> result;
+ result.reserve(row_schema->num_fields());
+ for (const std::shared_ptr<arrow::Field>& row_field :
row_schema->fields()) {
+ PAIMON_ASSIGN_OR_RAISE(int32_t field_id,
+
NestedProjectionUtils::GetPaimonFieldId(row_field));
+ auto field_index = field_indexes.find(field_id);
+ if (field_index == field_indexes.end()) {
+ return Status::Invalid(
+ fmt::format("cannot find field id {} in prepared schema",
field_id));
+ }
+ const std::shared_ptr<arrow::Field>& prepared_field =
+ prepared_schema->field(field_index->second);
+ if (!prepared_field->type()->Equals(row_field->type())) {
+ return Status::Invalid(fmt::format(
+ "prepared field id {} type {} does not match row "
+ "type {}",
+ field_id, prepared_field->type()->ToString(),
row_field->type()->ToString()));
+ }
+ result.push_back(field_index->second);
+ }
+ return result;
+}
+
+Status ValidateReaderParameters(const std::shared_ptr<arrow::Schema>&
prepared_schema,
+ const std::shared_ptr<arrow::Schema>&
key_schema,
+ const std::shared_ptr<arrow::Schema>&
value_schema,
+ const std::shared_ptr<MemoryPool>&
memory_pool) {
+
PAIMON_RETURN_NOT_OK(PreparedKeyValueReaderFactory::ValidateTransportSchema(prepared_schema));
+ if (!key_schema) {
+ return Status::Invalid("prepared key schema cannot be null");
+ }
+ if (!value_schema) {
+ return Status::Invalid("prepared value schema cannot be null");
+ }
+ if (!memory_pool) {
+ return Status::Invalid("prepared reader memory pool cannot be null");
+ }
+ 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() +
SpecialFields::kPreparedKeyValueValueStartIndex) {
+ 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 +
SpecialFields::kPreparedKeyValueValueStartIndex)
+ ->Equals(value_schema->field(i), true)) {
+ return Status::Invalid("commit requires the exact prepared writer
schema");
+ }
+ }
+ return Status::OK();
+}
+
+class PreparedReaderPlan {
+ public:
+ static Result<std::shared_ptr<const PreparedReaderPlan>> Create(
+ const std::shared_ptr<arrow::Schema>& prepared_schema,
+ const std::shared_ptr<arrow::Schema>& key_schema,
+ const std::shared_ptr<arrow::Schema>& value_schema) {
+ std::unordered_map<int32_t, int32_t> field_indexes;
+ field_indexes.reserve(prepared_schema->num_fields() -
+ SpecialFields::kPreparedKeyValueValueStartIndex);
+ for (int32_t i = SpecialFields::kPreparedKeyValueValueStartIndex;
+ i < prepared_schema->num_fields(); ++i) {
+ PAIMON_ASSIGN_OR_RAISE(int32_t field_id,
NestedProjectionUtils::GetPaimonFieldId(
+
prepared_schema->field(i)));
+ if (!field_indexes.emplace(field_id, i).second) {
+ return Status::Invalid(
+ fmt::format("duplicate field id {} in prepared schema",
field_id));
+ }
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::vector<int32_t> key_field_indexes,
+ ResolveFieldIndexes(prepared_schema,
field_indexes, key_schema));
+ PAIMON_ASSIGN_OR_RAISE(std::vector<int32_t> value_field_indexes,
+ ResolveFieldIndexes(prepared_schema,
field_indexes, value_schema));
+ return std::shared_ptr<const PreparedReaderPlan>(new
PreparedReaderPlan(
+ prepared_schema, std::move(key_field_indexes),
std::move(value_field_indexes)));
+ }
+
+ const std::shared_ptr<arrow::Schema>& PreparedSchema() const {
+ return prepared_schema_;
+ }
+
+ const std::vector<int32_t>& KeyFieldIndexes() const {
+ return key_field_indexes_;
+ }
+
+ const std::vector<int32_t>& ValueFieldIndexes() const {
+ return value_field_indexes_;
+ }
+
+ private:
+ PreparedReaderPlan(const std::shared_ptr<arrow::Schema>& schema,
+ std::vector<int32_t>&& key_indexes,
std::vector<int32_t>&& value_indexes)
+ : prepared_schema_(schema),
+ key_field_indexes_(std::move(key_indexes)),
+ value_field_indexes_(std::move(value_indexes)) {}
+
+ const std::shared_ptr<arrow::Schema> prepared_schema_;
+ const std::vector<int32_t> key_field_indexes_;
+ const std::vector<int32_t> value_field_indexes_;
+};
+
+class PreparedKeyValueReader final : public KeyValueRecordReader {
+ public:
+ PreparedKeyValueReader(std::unique_ptr<BatchReader>&& reader,
+ const std::shared_ptr<const PreparedReaderPlan>&
plan,
+ const std::optional<OffsetRange>& visible_offsets,
+ const std::shared_ptr<MemoryPool>& pool,
+ const std::shared_ptr<RealtimeOffsetCoverage>&
offset_coverage)
+ : reader_(std::move(reader)),
+ plan_(plan),
+ visible_offsets_(visible_offsets),
+ pool_(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();
+ BatchReader::ReadBatchWithBitmap batch_with_bitmap;
+ if (visible_offsets_.has_value()) {
+ PAIMON_ASSIGN_OR_RAISE(batch_with_bitmap,
reader_->NextBatchWithBitmap());
+ } else {
+ PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch,
reader_->NextBatch());
+ batch_with_bitmap.first = std::move(batch);
+ }
Review Comment:
This feels a bit too coupled. `visible_offsets_` doesn’t seem inherently
related to whether we should call `NextBatchWithBitmap` or `NextBatch`.
--
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]
