Skip to content

feat(realtime): attach primary-key realtime writers #11

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
HaHaJeff wants to merge 1 commit into
base: jeff/paimon-cpp-serving-realtime-pr5
Choose a base branch
from
Draft

feat(realtime): attach primary-key realtime writers #11

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 35 additions & 13 deletions src/paimon/core/operation/file_store_write.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,27 @@ struct KeyValue;
template <typename T>
class MergeFunctionWrapper;

namespace {

Status RestoreRealtimeCommittedProgress(const std::shared_ptr<RealtimeContext>& realtime_context,
const std::shared_ptr<SnapshotManager>& snapshot_manager,
const CoreOptions& options) {
PAIMON_ASSIGN_OR_RAISE(std::optional<Snapshot> latest_snapshot,
snapshot_manager->LatestSnapshot());
if (latest_snapshot) {
PAIMON_ASSIGN_OR_RAISE(
RealtimeOffsetMap realtime_committed_offsets,
RealtimeCommitProperties::ReadOffsets(latest_snapshot, options.GetFileSystem()));
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<RealtimeContextImpl> realtime_context_impl,
RealtimeContextImpl::Cast(realtime_context));
PAIMON_RETURN_NOT_OK(realtime_context_impl->AdvanceCommittedProgress(
latest_snapshot->Id(), realtime_committed_offsets));
}
return Status::OK();
}

} // namespace

Result<std::vector<RealtimeCommitProgress>> FileStoreWrite::PrepareCommitWithProgress(int64_t) {
return Status::Invalid("prepare commit with progress requires a real-time writer");
}
Expand Down Expand Up @@ -143,17 +164,8 @@ Result<std::unique_ptr<FileStoreWrite>> FileStoreWrite::Create(std::unique_ptr<W
if (options.DeletionVectorsEnabled()) {
return Status::Invalid("real-time append write does not support deletion vectors");
}
PAIMON_ASSIGN_OR_RAISE(std::optional<Snapshot> latest_snapshot,
snapshot_manager->LatestSnapshot());
if (latest_snapshot) {
PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetMap realtime_committed_offsets,
RealtimeCommitProperties::ReadOffsets(
latest_snapshot, options.GetFileSystem()));
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<RealtimeContextImpl> realtime_context_impl,
RealtimeContextImpl::Cast(ctx->GetRealtimeContext()));
PAIMON_RETURN_NOT_OK(realtime_context_impl->AdvanceCommittedProgress(
latest_snapshot->Id(), realtime_committed_offsets));
}
PAIMON_RETURN_NOT_OK(RestoreRealtimeCommittedProgress(ctx->GetRealtimeContext(),
snapshot_manager, options));
}
std::shared_ptr<arrow::Schema> write_schema = arrow_schema;
const auto& write_field_names = ctx->GetWriteSchema();
Expand Down Expand Up @@ -197,7 +209,16 @@ Result<std::unique_ptr<FileStoreWrite>> FileStoreWrite::Create(std::unique_ptr<W
} else {
// pk table
if (ctx->GetRealtimeContext()) {
return Status::Invalid("real-time write currently supports append tables only");
PAIMON_RETURN_NOT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *schema));
if (ignore_previous_files) {
return Status::NotImplemented(
"PK realtime requires restore from the latest snapshot");
}
if (!ctx->GetWriteSchema().empty()) {
return Status::NotImplemented("PK realtime does not support a custom write schema");
}
PAIMON_RETURN_NOT_OK(RestoreRealtimeCommittedProgress(ctx->GetRealtimeContext(),
snapshot_manager, options));
}
if (options.GetBucket() == BucketModeDefine::POSTPONE_BUCKET) {
return PostponeBucketFileStoreWrite::Create(
Expand Down Expand Up @@ -253,7 +274,8 @@ Result<std::unique_ptr<FileStoreWrite>> FileStoreWrite::Create(std::unique_ptr<W
ctx->GetRootPath(), schema, arrow_schema, partition_schema, dv_maintainer_factory,
io_manager, key_comparator, sequence_fields_comparator, merge_function_wrapper, options,
ignore_previous_files, ctx->IsStreamingMode(), ctx->IgnoreNumBucketCheck(),
ctx->EnableMultiThreadSpill(), ctx->GetExecutor(), ctx->GetMemoryPool());
ctx->EnableMultiThreadSpill(), ctx->GetRealtimeContext(), ctx->GetExecutor(),
ctx->GetMemoryPool());
}
}

Expand Down
83 changes: 72 additions & 11 deletions src/paimon/core/operation/key_value_file_store_write.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,21 +18,30 @@

#include "paimon/core/operation/key_value_file_store_write.h"

#include <optional>
#include <vector>

#include "arrow/c/bridge.h"
#include "paimon/common/data/binary_row.h"
#include "paimon/common/table/special_fields.h"
#include "paimon/core/compact/noop_compact_manager.h"
#include "paimon/core/core_options.h"
#include "paimon/core/io/data_file_meta.h"
#include "paimon/core/manifest/manifest_file.h"
#include "paimon/core/manifest/manifest_list.h"
#include "paimon/core/mergetree/levels.h"
#include "paimon/core/mergetree/merge_tree_writer.h"
#include "paimon/core/operation/commit/realtime_commit_properties.h"
#include "paimon/core/operation/file_store_scan.h"
#include "paimon/core/operation/key_value_file_store_scan.h"
#include "paimon/core/realtime/realtime_context_impl.h"
#include "paimon/core/realtime/realtime_primary_key_reader.h"
#include "paimon/core/realtime/realtime_primary_key_writer.h"
#include "paimon/core/schema/table_schema.h"
#include "paimon/core/utils/file_store_path_factory.h"
#include "paimon/core/utils/primary_key_table_utils.h"
#include "paimon/core/utils/snapshot_manager.h"
#include "paimon/realtime/realtime_context.h"

namespace arrow {
class Schema;
Expand Down Expand Up @@ -60,21 +69,27 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite(
const std::shared_ptr<MergeFunctionWrapper<KeyValue>>& merge_function_wrapper,
const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode,
bool ignore_num_bucket_check, bool enable_multi_thread_spill,
const std::shared_ptr<RealtimeContext>& realtime_context,
const std::shared_ptr<Executor>& executor, const std::shared_ptr<MemoryPool>& pool)
: AbstractFileStoreWrite(file_store_path_factory, snapshot_manager, schema_manager, commit_user,
root_path, table_schema, schema, /*write_schema=*/schema,
partition_schema, dv_maintainer_factory, io_manager, options,
ignore_previous_files, is_streaming_mode, ignore_num_bucket_check,
executor, pool),
enable_multi_thread_spill_(enable_multi_thread_spill),
realtime_context_(realtime_context),
key_comparator_(key_comparator),
user_defined_seq_comparator_(user_defined_seq_comparator),
merge_function_wrapper_(merge_function_wrapper),
compact_manager_factory_(std::make_unique<MergeTreeCompactManagerFactory>(
options_, key_comparator_, user_defined_seq_comparator_, compaction_metrics_,
table_schema_, schema_, schema_manager_, io_manager_, cache_manager_,
file_store_path_factory_, root_path_, ignore_previous_files, pool_)),
logger_(Logger::GetLogger("KeyValueFileStoreWrite")) {}
logger_(Logger::GetLogger("KeyValueFileStoreWrite")) {
if (realtime_context_) {
writer_memory_manager_ = std::make_unique<NoopWriterMemoryManager>();
}
}

Result<std::unique_ptr<FileStoreScan>> KeyValueFileStoreWrite::CreateFileStoreScan(
const std::shared_ptr<ScanFilter>& scan_filter) const {
Expand Down Expand Up @@ -106,22 +121,68 @@ Result<std::shared_ptr<BatchWriter>> KeyValueFileStoreWrite::CreateWriter(
file_store_path_factory_->CreateDataFilePathFactory(partition, bucket));
PAIMON_ASSIGN_OR_RAISE(std::vector<std::string> trimmed_primary_keys,
table_schema_->TrimmedPrimaryKeys());
PAIMON_ASSIGN_OR_RAISE(
std::shared_ptr<Levels> levels,
Levels::Create(key_comparator_, restore_data_files, options_.GetNumLevels()));
auto compact_strategy = compact_manager_factory_->CreateCompactStrategy();
PAIMON_ASSIGN_OR_RAISE(
std::shared_ptr<CompactManager> compact_manager,
compact_manager_factory_->CreateCompactManager(partition, bucket, compact_strategy,
compact_executor_, levels, dv_maintainer));
std::map<std::string, std::string> partition_map;
std::shared_ptr<CompactManager> compact_manager;
std::shared_ptr<RealtimeContextImpl> realtime_context_impl;
std::optional<RealtimeStoreState> realtime_store_state;
std::shared_ptr<arrow::Schema> transport_schema;
if (realtime_context_) {
std::vector<std::pair<std::string, std::string>> partition_values;
PAIMON_ASSIGN_OR_RAISE(partition_values,
file_store_path_factory_->GeneratePartitionVector(partition));
partition_map =
std::map<std::string, std::string>(partition_values.begin(), partition_values.end());
PAIMON_ASSIGN_OR_RAISE(realtime_context_impl, RealtimeContextImpl::Cast(realtime_context_));
transport_schema = RealtimePrimaryKeyLayout::CreateSchema(schema_->fields());
auto c_write_schema = std::make_unique<ArrowSchema>();
PAIMON_RETURN_NOT_OK_FROM_ARROW(
arrow::ExportSchema(*transport_schema, c_write_schema.get()));
PAIMON_ASSIGN_OR_RAISE(
RealtimeStoreState store_state,
realtime_context_impl->GetOrCreateRealtimeStore(
RealtimeStoreCreateRequest{std::move(c_write_schema), options_.ToMap(), pool_,
RealtimeStoreMode::PRIMARY_KEY},
RealtimePartitionBucket(partition_map, bucket)));
realtime_store_state = std::move(store_state);
compact_manager = std::make_shared<NoopCompactManager>();
} else {
PAIMON_ASSIGN_OR_RAISE(
std::shared_ptr<Levels> levels,
Levels::Create(key_comparator_, restore_data_files, options_.GetNumLevels()));
auto compact_strategy = compact_manager_factory_->CreateCompactStrategy();
PAIMON_ASSIGN_OR_RAISE(compact_manager, compact_manager_factory_->CreateCompactManager(
partition, bucket, compact_strategy,
compact_executor_, levels, dv_maintainer));
}

PAIMON_ASSIGN_OR_RAISE(
std::shared_ptr<MergeTreeWriter> writer,
MergeTreeWriter::Create(
restore_max_seq_number, trimmed_primary_keys, data_file_path_factory, key_comparator_,
user_defined_seq_comparator_, merge_function_wrapper_, table_schema_->Id(), schema_,
options_, compact_manager, io_manager_, enable_multi_thread_spill_, pool_));
return writer;
options_, compact_manager, realtime_context_ ? nullptr : io_manager_,
enable_multi_thread_spill_, pool_));
if (!realtime_context_) {
return std::shared_ptr<BatchWriter>(std::move(writer));
}
return RealtimePrimaryKeyWriter::Create(partition_map, bucket, schema_, transport_schema,
trimmed_primary_keys, key_comparator_, options_,
realtime_context_impl, realtime_store_state.value(),
restore_max_seq_number, writer, pool_);
}

Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) {
if (!realtime_context_) {
return Status::Invalid("refresh committed snapshot requires a real-time writer");
}
PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(snapshot_id));
PAIMON_ASSIGN_OR_RAISE(
RealtimeOffsetMap committed_offsets,
RealtimeCommitProperties::ReadOffsets(std::optional<Snapshot>(std::move(snapshot)),
options_.GetFileSystem()));
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<RealtimeContextImpl> realtime_context_impl,
RealtimeContextImpl::Cast(realtime_context_));
return realtime_context_impl->AdvanceCommittedProgress(snapshot_id, committed_offsets);
}

Status KeyValueFileStoreWrite::Close() {
Expand Down
8 changes: 8 additions & 0 deletions src/paimon/core/operation/key_value_file_store_write.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ class SnapshotManager;
class SchemaManager;
class TableSchema;
class IOManager;
class RealtimeContext;
struct KeyValue;
template <typename T>
class MergeFunctionWrapper;
Expand All @@ -65,8 +66,10 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite {
const std::shared_ptr<MergeFunctionWrapper<KeyValue>>& merge_function_wrapper,
const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode,
bool ignore_num_bucket_check, bool enable_multi_thread_spill,
const std::shared_ptr<RealtimeContext>& realtime_context,
const std::shared_ptr<Executor>& executor, const std::shared_ptr<MemoryPool>& pool);

Status RefreshCommittedSnapshot(int64_t snapshot_id) override;
Status Close() override;

private:
Expand All @@ -79,8 +82,13 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite {
Result<std::unique_ptr<FileStoreScan>> CreateFileStoreScan(
const std::shared_ptr<ScanFilter>& filter) const override;

bool IsRealtimeWrite() const override {
return realtime_context_ != nullptr;
}

private:
bool enable_multi_thread_spill_;
std::shared_ptr<RealtimeContext> realtime_context_;
std::shared_ptr<FieldsComparator> key_comparator_;
std::shared_ptr<FieldsComparator> user_defined_seq_comparator_;
std::shared_ptr<MergeFunctionWrapper<KeyValue>> merge_function_wrapper_;
Expand Down
Loading
Loading