HaHaJeff · HaHaJeff · Aug 29, 2026
diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt
@@ -385,6 +385,7 @@ set(PAIMON_CORE_SRCS
     core/operation/write_restore.cpp
     core/realtime/arrow_realtime_store.cpp
     core/realtime/arrow_realtime_store_factory.cpp
+    core/realtime/primary_key_realtime_store.cpp
     core/realtime/realtime_append_only_writer.cpp
     core/realtime/realtime_context.cpp
     core/realtime/realtime_context_impl.cpp
@@ -791,6 +792,7 @@ if(PAIMON_BUILD_TESTS)
                     core/manifest/index_manifest_file_handler_test.cpp
                     core/memory/writer_memory_manager_test.cpp
                     core/realtime/arrow_realtime_store_test.cpp
+                    core/realtime/primary_key_realtime_store_test.cpp
                     core/realtime/realtime_context_test.cpp
                     core/realtime/realtime_reader_test.cpp
                     core/mergetree/levels_test.cpp

diff --git a/src/paimon/common/utils/arrow/arrow_utils.cpp b/src/paimon/common/utils/arrow/arrow_utils.cpp
@@ -364,6 +364,22 @@ void ArrowUtils::TraverseArray(const std::shared_ptr<arrow::Array>& array) {
     }
 }
 
+uint64_t ArrowUtils::GetArrayMemoryUsage(const std::shared_ptr<arrow::ArrayData>& data) {
+    uint64_t result = 0;
+    for (const std::shared_ptr<arrow::Buffer>& buffer : data->buffers) {
+        if (buffer) {
+            result += static_cast<uint64_t>(buffer->size());
+        }
+    }
+    for (const std::shared_ptr<arrow::ArrayData>& child : data->child_data) {
+        result += GetArrayMemoryUsage(child);
+    }
+    if (data->dictionary) {
+        result += GetArrayMemoryUsage(data->dictionary);
+    }
+    return result;
+}
+
 bool ArrowUtils::EqualsIgnoreNullable(const std::shared_ptr<arrow::DataType>& type,
                                       const std::shared_ptr<arrow::DataType>& other_type) {
     if (type->id() != other_type->id() || type->num_fields() != other_type->num_fields()) {

diff --git a/src/paimon/common/utils/arrow/arrow_utils.h b/src/paimon/common/utils/arrow/arrow_utils.h
@@ -19,6 +19,7 @@
 
 #pragma once
 
+#include <cstdint>
 #include <vector>
 
 #include "arrow/api.h"
@@ -48,6 +49,8 @@ class PAIMON_EXPORT ArrowUtils {
     // avoid subsequent multi-threading problems.
     static void TraverseArray(const std::shared_ptr<arrow::Array>& array);
 
+    static uint64_t GetArrayMemoryUsage(const std::shared_ptr<arrow::ArrayData>& data);
+
     static Result<std::shared_ptr<arrow::StructArray>> RemoveFieldFromStructArray(
         const std::shared_ptr<arrow::StructArray>& struct_array, const std::string& field_name);
 

diff --git a/src/paimon/common/utils/arrow/mem_utils.cpp b/src/paimon/common/utils/arrow/mem_utils.cpp
@@ -24,12 +24,31 @@
 #include <new>
 #include <string>
 
+#include "arrow/c/abi.h"
+#include "arrow/c/helpers.h"
 #include "arrow/memory_pool.h"
 #include "arrow/status.h"
 #include "fmt/format.h"
 #include "paimon/memory/memory_pool.h"
 
 namespace paimon {
+namespace {
+
+struct ArrowArrayPrivateData {
+    void (*release)(ArrowArray*);
+    void* private_data;
+    std::shared_ptr<arrow::MemoryPool> arrow_pool;
+};
+
+void ReleaseArrowArray(ArrowArray* array) {
+    std::unique_ptr<ArrowArrayPrivateData> data(
+        static_cast<ArrowArrayPrivateData*>(array->private_data));
+    array->release = data->release;
+    array->private_data = data->private_data;
+    array->release(array);
+}
+
+}  // namespace
 
 class ArrowMemPoolAdaptor : public arrow::MemoryPool {
  public:
@@ -107,4 +126,26 @@ std::unique_ptr<arrow::MemoryPool> GetArrowPool(const std::shared_ptr<MemoryPool
     return std::make_unique<ArrowMemPoolAdaptor>(pool);
 }
 
+Status RetainArrowArrayMemoryPool(ArrowArray* array,
+                                  const std::shared_ptr<arrow::MemoryPool>& arrow_pool) {
+    if (!array || !array->release) {
+        return Status::Invalid("cannot retain Arrow array memory pool");
+    }
+    if (!arrow_pool) {
+        ArrowArrayRelease(array);
+        return Status::Invalid("cannot retain Arrow array memory pool");
+    }
+    std::unique_ptr<ArrowArrayPrivateData> data;
+    try {
+        data = std::make_unique<ArrowArrayPrivateData>(
+            ArrowArrayPrivateData{array->release, array->private_data, arrow_pool});
+    } catch (const std::bad_alloc&) {
+        ArrowArrayRelease(array);
+        return Status::OutOfMemory("failed to retain Arrow array memory pool");
+    }
+    array->private_data = data.release();
+    array->release = ReleaseArrowArray;
+    return Status::OK();
+}
+
 }  // namespace paimon
diff --git a/src/paimon/common/utils/arrow/mem_utils.h b/src/paimon/common/utils/arrow/mem_utils.h
@@ -23,11 +23,17 @@
 
 #include "arrow/memory_pool.h"
 #include "paimon/memory/memory_pool.h"
+#include "paimon/status.h"
 #include "paimon/visibility.h"
 
+struct ArrowArray;
+
 namespace paimon {
 
 PAIMON_EXPORT std::unique_ptr<arrow::MemoryPool> GetArrowPool(
     const std::shared_ptr<MemoryPool>& pool);
 
+Status RetainArrowArrayMemoryPool(ArrowArray* array,
+                                  const std::shared_ptr<arrow::MemoryPool>& arrow_pool);
+
 }  // namespace paimon
diff --git a/src/paimon/core/realtime/arrow_realtime_store.cpp b/src/paimon/core/realtime/arrow_realtime_store.cpp
@@ -33,6 +33,7 @@
 #include "paimon/common/reader/complete_row_kind_batch_reader.h"
 #include "paimon/common/table/special_fields.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/projected_array.h"
@@ -43,22 +44,6 @@
 namespace paimon {
 namespace {
 
-uint64_t GetArrayMemoryUsage(const std::shared_ptr<arrow::ArrayData>& data) {
-    uint64_t result = 0;
-    for (const std::shared_ptr<arrow::Buffer>& buffer : data->buffers) {
-        if (buffer) {
-            result += static_cast<uint64_t>(buffer->size());
-        }
-    }
-    for (const std::shared_ptr<arrow::ArrayData>& child : data->child_data) {
-        result += GetArrayMemoryUsage(child);
-    }
-    if (data->dictionary) {
-        result += GetArrayMemoryUsage(data->dictionary);
-    }
-    return result;
-}
-
 bool SupportsMinMax(const std::shared_ptr<arrow::DataType>& type) {
     switch (type->id()) {
         case arrow::Type::BOOL:
@@ -393,11 +378,11 @@ Status ArrowRealtimeStore::Write(RealtimeWriteBatch&& write_batch) {
     if (building_range_ && write_batch.offset_range.begin != building_range_->end) {
         return Status::Invalid("real-time offset ranges must be contiguous");
     }
-    uint64_t memory_usage = GetArrayMemoryUsage(struct_array->data());
+    uint64_t memory_usage = ArrowUtils::GetArrayMemoryUsage(struct_array->data());
     if (statistics) {
-        memory_usage += GetArrayMemoryUsage(statistics->min_values->data()) +
-                        GetArrayMemoryUsage(statistics->max_values->data()) +
-                        GetArrayMemoryUsage(statistics->null_counts->data());
+        memory_usage += ArrowUtils::GetArrayMemoryUsage(statistics->min_values->data()) +
+                        ArrowUtils::GetArrayMemoryUsage(statistics->max_values->data()) +
+                        ArrowUtils::GetArrayMemoryUsage(statistics->null_counts->data());
     }
     building_memory_usage_ += memory_usage;
     building_batches_.push_back(

diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp
@@ -25,6 +25,7 @@
 #include "paimon/common/utils/arrow/status_utils.h"
 #include "paimon/common/utils/scope_guard.h"
 #include "paimon/core/realtime/arrow_realtime_store.h"
+#include "paimon/core/realtime/primary_key_realtime_store.h"
 #include "paimon/macros.h"
 
 namespace paimon {
@@ -48,8 +49,10 @@ Result<std::shared_ptr<RealtimeStore>> ArrowRealtimeStoreFactory::Create(
                                                         request.memory_pool, arrow_pool);
         }
         case RealtimeStoreMode::PRIMARY_KEY: {
-            return Status::NotImplemented(
-                "primary-key real-time store support is not installed");
+            PAIMON_ASSIGN_OR_RAISE(
+                std::shared_ptr<PrimaryKeyRealtimeStore> store,
+                PrimaryKeyRealtimeStore::Create(imported_schema, request.memory_pool));
+            return std::shared_ptr<RealtimeStore>(std::move(store));
         }
     }
     return Status::Invalid("invalid real-time store mode: ", static_cast<int32_t>(request.mode));