feat(realtime): add primary-key in-memory support

include/paimon/realtime/arrow_realtime_store_factory.h

@@ -27,10 +27,7 @@ namespace paimon {
 class PAIMON_EXPORT ArrowRealtimeStoreFactory : public RealtimeStoreFactory {
  public:
     /// Creates an Arrow-backed store for one partition and bucket.
-    Result<std::shared_ptr<RealtimeStore>> Create(
-        std::unique_ptr<::ArrowSchema> write_schema, StatisticsMode statistics_mode,
-        const std::map<std::string, std::string>& options,
-        const std::shared_ptr<MemoryPool>& memory_pool) override;
+    Result<std::shared_ptr<RealtimeStore>> Create(RealtimeStoreCreateRequest&& request) override;
 };
 
 }  // namespace paimon

include/paimon/realtime/realtime_store.h

@@ -27,6 +27,7 @@
 #include <utility>
 #include <vector>
 
+#include "arrow/c/abi.h"
 #include "paimon/reader/batch_reader.h"
 #include "paimon/realtime/offset_range.h"
 #include "paimon/record_batch.h"
@@ -41,10 +42,33 @@ namespace paimon {
 class MemoryPool;
 class Predicate;
 
-/// A table record batch and its framework-assigned contiguous offset range.
+enum class PAIMON_EXPORT RealtimeStoreMode {
+    APPEND_ONLY,
+    PRIMARY_KEY,
+};
+
+/// Parameters used by a `RealtimeStoreFactory` to create a store.
+struct PAIMON_EXPORT RealtimeStoreCreateRequest {
+    /// Schema whose ownership is transferred to the factory. Append mode receives the complete
+    /// table write schema. Primary-key mode receives the realtime primary-key transport schema:
+    /// [_VALUE_KIND, _SEQUENCE_NUMBER, _REALTIME_OFFSET, table write fields].
+    std::unique_ptr<::ArrowSchema> write_schema;
+    /// Table options available to the store implementation.
+    std::map<std::string, std::string> options;
+    /// Memory pool for allocations retained by the store.
+    std::shared_ptr<MemoryPool> memory_pool;
+    /// Table mode implemented by the store.
+    RealtimeStoreMode mode = RealtimeStoreMode::APPEND_ONLY;
+    /// Statistics collected by append-only stores.
+    StatisticsMode statistics_mode = StatisticsMode::NONE;
+};
+
+/// A record batch and its framework-assigned contiguous offset range.
 ///
-/// The batch contains only table write fields. Row `i` is associated with
-/// `offset_range.begin + i`; the offset is progress metadata and is not a table field.
+/// Append-mode batches contain table write fields, and row `i` has offset
+/// `offset_range.begin + i`. Primary-key batches use the realtime primary-key transport schema,
+/// are sorted by full primary key then sequence number, and retain the original offset in
+/// `_REALTIME_OFFSET`.
 struct PAIMON_EXPORT RealtimeWriteBatch {
     /// Input batch whose ownership is transferred to `RealtimeStore::Write`.
     std::unique_ptr<RecordBatch> batch;
@@ -79,7 +103,11 @@ class PAIMON_EXPORT RealtimeReadView {
 
 /// Parameters used by a `RealtimeStore` to create readers for a query.
 struct PAIMON_EXPORT RealtimeQueryContext {
-    /// Requested output fields before the mandatory leading `_VALUE_KIND` field is added.
+    /// Append mode receives the requested output fields before the mandatory leading
+    /// `_VALUE_KIND` field is added. Primary-key mode receives the requested realtime primary-key
+    /// transport schema.
+    /// This schema is borrowed and remains valid only during `CreateQueryReaders`; plugins must
+    /// import or copy it synchronously.
     ::ArrowSchema* read_schema;
     /// Predicate using field indexes from `read_schema`.
     std::shared_ptr<Predicate> predicate;
@@ -116,9 +144,10 @@ class PAIMON_EXPORT RealtimeStore {
 
     /// Creates readers that expose all rows in a sealed segment for Paimon file writing.
     ///
-    /// Concatenating the returned readers must produce every sealed row exactly once and in write
-    /// order. Each output batch contains `_VALUE_KIND` followed by all fields from the factory's
-    /// `write_schema`.
+    /// The returned readers collectively expose every sealed row exactly once. Append-mode readers
+    /// preserve write order and contain `_VALUE_KIND` followed by table write fields. Primary-key
+    /// readers use the realtime primary-key transport schema; each reader's complete stream is
+    /// sorted by full primary key then sequence number.
     virtual Result<std::vector<std::unique_ptr<BatchReader>>> CreateCommitReaders(
         const std::shared_ptr<RealtimeSegmentHandle>& segment) = 0;
 
@@ -128,13 +157,16 @@ class PAIMON_EXPORT RealtimeStore {
     /// also provide a consistent snapshot when a write or seal is in progress.
     virtual Result<std::shared_ptr<RealtimeReadView>> AcquireReadView() = 0;
 
-    /// Creates readers over rows in `view` whose offsets are greater than or equal to
-    /// `offset_begin`.
+    /// Creates readers over rows in `view`. Append mode returns rows whose offsets are greater than
+    /// or equal to `offset_begin`; primary-key mode ignores `offset_begin`.
     ///
-    /// Each output batch contains `_VALUE_KIND` first, followed by the fields requested by
-    /// `context.read_schema` except a duplicate `_VALUE_KIND`. Concatenating all returned readers
-    /// must produce every matching row once. Paimon retains `view` for the lifetime of the
-    /// resulting framework reader.
+    /// Append-mode batches contain `_VALUE_KIND` followed by the requested fields except a
+    /// duplicate `_VALUE_KIND`, and collectively expose every matching row exactly once.
+    /// Primary-key batches use the requested realtime primary-key transport schema, including
+    /// nested field-ID alignment, and may contain multiple mutations per key; each reader's
+    /// complete stream is sorted by full primary key then sequence number, and the readers
+    /// collectively expose every raw mutation exactly once. Paimon retains `view` for the lifetime
+    /// of the resulting framework reader.
     virtual Result<std::vector<std::unique_ptr<BatchReader>>> CreateQueryReaders(
         const std::shared_ptr<RealtimeReadView>& view, int64_t offset_begin,
         const RealtimeQueryContext& context) = 0;
@@ -158,15 +190,9 @@ class PAIMON_EXPORT RealtimeStoreFactory {
     virtual ~RealtimeStoreFactory() = default;
 
     /// Creates a store configured with the supplied schema, statistics, options, and memory pool.
-    /// @param write_schema Complete table write schema whose ownership is transferred to the
-    /// factory. The factory may consume it or retain it in the created store.
-    /// @param statistics_mode Framework-parsed statistics collection mode.
-    /// @param options Effective table options available to the store.
-    /// @param memory_pool Memory pool provided by the write context.
-    virtual Result<std::shared_ptr<RealtimeStore>> Create(
-        std::unique_ptr<::ArrowSchema> write_schema, StatisticsMode statistics_mode,
-        const std::map<std::string, std::string>& options,
-        const std::shared_ptr<MemoryPool>& memory_pool) = 0;
+    /// Creates a store for the requested table mode.
+    /// The factory consumes `request`, including ownership of `request.write_schema`.
+    virtual Result<std::shared_ptr<RealtimeStore>> Create(RealtimeStoreCreateRequest&& request) = 0;
 };
 
 }  // namespace paimon

include/paimon/utils/special_field_ids.h

@@ -42,6 +42,8 @@ class SpecialFieldIds {
 
     /// Special field ID reserved for index score. Value: CPP_FIELD_ID_END - 1
     inline static constexpr int32_t INDEX_SCORE = CPP_FIELD_ID_END - 1;
+    /// Special field ID reserved for realtime offset. Value: CPP_FIELD_ID_END - 2
+    inline static constexpr int32_t REALTIME_OFFSET = CPP_FIELD_ID_END - 2;
 
     /// Lowest field ID reserved for system fields; IDs at or above it are excluded from the
     /// highest field ID of a schema. Value: INT32_MAX / 2

src/paimon/CMakeLists.txt

@@ -385,9 +385,12 @@ set(PAIMON_CORE_SRCS
     core/operation/write_restore.cpp
     core/realtime/arrow_realtime_store.cpp
     core/realtime/arrow_realtime_store_factory.cpp
+    core/realtime/realtime_primary_key_reader.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
+    core/realtime/realtime_primary_key_writer.cpp
     core/postpone/postpone_bucket_writer.cpp
     core/schema/arrow_schema_validator.cpp
     core/schema/schema_manager.cpp
@@ -791,6 +794,8 @@ 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_primary_key_reader_test.cpp
                     core/realtime/realtime_context_test.cpp
                     core/realtime/realtime_reader_test.cpp
                     core/mergetree/levels_test.cpp

src/paimon/common/table/special_fields.h

@@ -21,6 +21,7 @@
 #include <cstdint>
 #include <limits>
 #include <string>
+#include <utility>
 
 #include "arrow/type_fwd.h"
 #include "paimon/common/types/data_field.h"
@@ -66,13 +67,20 @@ struct SpecialFields {
         return data_field;
     }
 
+    static const DataField& RealtimeOffset() {
+        static const DataField data_field =
+            DataField(SpecialFieldIds::REALTIME_OFFSET,
+                      arrow::field("_REALTIME_OFFSET", arrow::int64(), false));
+        return data_field;
+    }
+
     static bool IsSystemField(const std::string& field_name) {
         if (StringUtils::StartsWith(field_name, KEY_FIELD_PREFIX)) {
             return true;
         }
         return field_name == SequenceNumber().Name() || field_name == ValueKind().Name() ||
                field_name == RowKind().Name() || field_name == RowId().Name() ||
-               field_name == IndexScore().Name();
+               field_name == IndexScore().Name() || field_name == RealtimeOffset().Name();
     }
 
     // TODO(xinyu.lxy): add a func to complete row-tracking fields

src/paimon/common/table/special_fields_test.cpp

@@ -55,6 +55,13 @@ TEST(SpecialFieldsTest, TestIndexScore) {
     ASSERT_EQ(SpecialFields::IndexScore().Type()->id(), arrow::Type::FLOAT);
 }
 
+TEST(SpecialFieldsTest, TestRealtimeOffset) {
+    ASSERT_EQ(SpecialFields::RealtimeOffset().Id(), SpecialFieldIds::REALTIME_OFFSET);
+    ASSERT_EQ(SpecialFields::RealtimeOffset().Name(), "_REALTIME_OFFSET");
+    ASSERT_EQ(SpecialFields::RealtimeOffset().Type()->id(), arrow::Type::INT64);
+    ASSERT_FALSE(SpecialFields::RealtimeOffset().Nullable());
+}
+
 TEST(SpecialFieldsTest, TestKeyValueSpecialFieldCount) {
     ASSERT_EQ(SpecialFields::KEY_VALUE_SPECIAL_FIELD_COUNT, 2);
 }
@@ -66,6 +73,7 @@ TEST(SpecialFieldsTest, TestIsSystemField) {
     ASSERT_TRUE(SpecialFields::IsSystemField("rowkind"));
     ASSERT_TRUE(SpecialFields::IsSystemField("_ROW_ID"));
     ASSERT_TRUE(SpecialFields::IsSystemField("_INDEX_SCORE"));
+    ASSERT_TRUE(SpecialFields::IsSystemField("_REALTIME_OFFSET"));
     ASSERT_TRUE(SpecialFields::IsSystemField("_KEY_0"));
 }
 

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()) {

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);
 

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

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

src/paimon/core/io/merged_key_value_record_reader_test.cpp

@@ -20,6 +20,7 @@
 
 #include <memory>
 #include <utility>
+#include <vector>
 
 #include "arrow/api.h"
 #include "arrow/array/array_nested.h"