feat(read): route primary-key realtime scans

src/paimon/core/realtime/realtime_reader_test.cpp

@@ -21,6 +21,7 @@
 
 #include <memory>
 #include <optional>
+#include <utility>
 
 #include "paimon/arrow/abi.h"
 #include "paimon/testing/utils/testharness.h"
@@ -37,6 +38,8 @@ class TestingReadView : public RealtimeReadView {
 
 class TestingBatchReader : public BatchReader {
  public:
+    explicit TestingBatchReader(int32_t* close_count = nullptr) : close_count_(close_count) {}
+
     Result<ReadBatch> NextBatch() override {
         return MakeEofBatch();
     }
@@ -45,7 +48,14 @@ class TestingBatchReader : public BatchReader {
         return nullptr;
     }
 
-    void Close() override {}
+    void Close() override {
+        if (close_count_) {
+            ++(*close_count_);
+        }
+    }
+
+ private:
+    int32_t* close_count_;
 };
 
 TEST(RealtimeReaderTest, TestRejectsIncompleteReader) {
@@ -57,5 +67,19 @@ TEST(RealtimeReaderTest, TestRejectsIncompleteReader) {
                         "inner reader is null");
 }
 
+TEST(RealtimeReaderTest, TestCloseReleasesResources) {
+    int32_t close_count = 0;
+    std::shared_ptr<TestingReadView> read_view = std::make_shared<TestingReadView>();
+    std::weak_ptr<TestingReadView> weak_read_view = read_view;
+    ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<RealtimeReader> reader,
+        RealtimeReader::Create(std::move(read_view),
+                               std::make_unique<TestingBatchReader>(&close_count)));
+    ASSERT_FALSE(weak_read_view.expired());
+    reader->Close();
+    ASSERT_EQ(1, close_count);
+    ASSERT_TRUE(weak_read_view.expired());
+}
+
 }  // namespace
 }  // namespace paimon::test

src/paimon/core/table/source/append_only_table_read.cpp

@@ -77,6 +77,13 @@ Result<std::unique_ptr<BatchReader>> AppendOnlyTableRead::CreateReader(
     std::vector<std::unique_ptr<BatchReader>> readers;
     readers.reserve(splits.size());
     std::vector<std::shared_ptr<RealtimeSplit>> realtime_splits;
+    ScopeGuard cleanup_guard([&]() {
+        for (const std::unique_ptr<BatchReader>& reader : readers) {
+            if (reader) {
+                reader->Close();
+            }
+        }
+    });
     for (const std::shared_ptr<Split>& split : splits) {
         std::shared_ptr<RealtimeSplit> realtime_split =
             std::dynamic_pointer_cast<RealtimeSplit>(split);
@@ -92,8 +99,6 @@ Result<std::unique_ptr<BatchReader>> AppendOnlyTableRead::CreateReader(
         }
     }
 
-    std::unique_ptr<BatchReader> result =
-        std::make_unique<ConcatBatchReader>(std::move(readers), GetMemoryPool());
     if (!realtime_splits.empty()) {
         const std::shared_ptr<RealtimeContext> realtime_context = context_->GetRealtimeContext();
         if (!realtime_context) {
@@ -106,7 +111,7 @@ Result<std::unique_ptr<BatchReader>> AppendOnlyTableRead::CreateReader(
                 realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket()));
         }
     }
-    return result;
+    return std::make_unique<ConcatBatchReader>(std::move(readers), GetMemoryPool());
 }
 
 Result<std::unique_ptr<BatchReader>> AppendOnlyTableRead::CreateRealtimeReader(
@@ -124,6 +129,13 @@ Result<std::unique_ptr<BatchReader>> AppendOnlyTableRead::CreateRealtimeReader(
                            realtime_context_impl->ResolveReadView(realtime_split->OpaqueTicket()));
     std::vector<std::unique_ptr<BatchReader>> readers;
     readers.reserve(realtime_split->DiskSplits().size() + 1);
+    ScopeGuard readers_guard([&readers]() {
+        for (const std::unique_ptr<BatchReader>& reader : readers) {
+            if (reader) {
+                reader->Close();
+            }
+        }
+    });
     const RealtimePartitionBucket expected_partition_bucket(realtime_split->Partition(),
                                                             realtime_split->Bucket());
     if (memory.partition_bucket != expected_partition_bucket) {
@@ -150,23 +162,33 @@ Result<std::unique_ptr<BatchReader>> AppendOnlyTableRead::CreateRealtimeReader(
         std::vector<std::unique_ptr<BatchReader>> memory_readers,
         memory.store->CreateQueryReaders(memory.read_view, realtime_split->CommittedEndOffset(),
                                          query_context));
-
+    const size_t first_memory_reader = readers.size();
+    readers.reserve(readers.size() + memory_readers.size());
     for (std::unique_ptr<BatchReader>& memory_reader : memory_readers) {
+        readers.push_back(std::move(memory_reader));
+    }
+
+    for (size_t i = first_memory_reader; i < readers.size(); ++i) {
+        std::unique_ptr<BatchReader>& memory_reader = readers[i];
+        if (!memory_reader) {
+            return Status::Invalid("append-only real-time store returned a null query reader");
+        }
         if (context_->EnablePredicateFilter() && context_->GetPredicate()) {
             PAIMON_ASSIGN_OR_RAISE(memory_reader, PredicateBatchReader::Create(
                                                       std::move(memory_reader),
                                                       context_->GetPredicate(), GetMemoryPool()));
         }
         PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<RealtimeReader> realtime_reader,
                                RealtimeReader::Create(memory.read_view, std::move(memory_reader)));
-        readers.push_back(std::move(realtime_reader));
+        memory_reader = std::move(realtime_reader);
     }
-    std::unique_ptr<BatchReader> result =
-        std::make_unique<ConcatBatchReader>(std::move(readers), GetMemoryPool());
     if (release_ticket) {
         PAIMON_RETURN_NOT_OK(
             realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket()));
     }
+    std::unique_ptr<BatchReader> result =
+        std::make_unique<ConcatBatchReader>(std::move(readers), GetMemoryPool());
+    readers_guard.Release();
     return result;
 }
 

src/paimon/core/table/source/realtime_table_scan.cpp

@@ -108,7 +108,7 @@ Result<std::vector<std::shared_ptr<Split>>> RealtimeTableScan::CreateRealtimeSpl
     for (const std::shared_ptr<Split>& split : disk_splits) {
         std::shared_ptr<DataSplitImpl> data_split = std::dynamic_pointer_cast<DataSplitImpl>(split);
         if (!data_split) {
-            return Status::Invalid("real-time append scan requires process-local data splits");
+            return Status::Invalid("real-time scan requires process-local data splits");
         }
         std::vector<std::pair<std::string, std::string>> partition_values;
         PAIMON_ASSIGN_OR_RAISE(partition_values,
@@ -152,16 +152,17 @@ Result<std::vector<std::shared_ptr<Split>>> RealtimeTableScan::CreateRealtimeSpl
             continue;
         }
 
-        // Append tables can schedule all but the tail disk split independently. The tail split
-        // carries the immutable memory view so disk and memory are still concatenated by one
-        // RealtimeSplit without collapsing the whole partition-bucket into one scheduling unit.
-        auto tail_disk_split = std::prev(grouped_disk_splits.end());
-        result.insert(result.end(), grouped_disk_splits.begin(), tail_disk_split);
-        std::vector<std::shared_ptr<Split>> realtime_disk_splits;
-        realtime_disk_splits.push_back(std::move(*tail_disk_split));
         RealtimePartitionBucketView& memory = memory_iter->second;
+        if (!pk_table_) {
+            // Append tables can schedule all but the tail disk split independently. The tail split
+            // carries the immutable memory view so disk and memory are still concatenated by one
+            // RealtimeSplit without collapsing the whole partition-bucket into one scheduling unit.
+            auto tail_disk_split = std::prev(grouped_disk_splits.end());
+            result.insert(result.end(), grouped_disk_splits.begin(), tail_disk_split);
+            grouped_disk_splits.erase(grouped_disk_splits.begin(), tail_disk_split);
+        }
         PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Split> realtime_split,
-                               create_realtime_split(key, std::move(realtime_disk_splits), memory));
+                               create_realtime_split(key, std::move(grouped_disk_splits), memory));
         result.push_back(std::move(realtime_split));
         active_memory.erase(memory_iter);
     }
@@ -176,14 +177,15 @@ Result<std::vector<std::shared_ptr<Split>>> RealtimeTableScan::CreateRealtimeSpl
     return result;
 }
 
-RealtimeTableScan::RealtimeTableScan(std::unique_ptr<TableScan>&& disk_scan,
+RealtimeTableScan::RealtimeTableScan(std::unique_ptr<TableScan>&& disk_scan, bool pk_table,
                                      const std::shared_ptr<RealtimeContextImpl>& realtime_context,
                                      const std::shared_ptr<FileStorePathFactory>& path_factory,
                                      const std::shared_ptr<SnapshotManager>& snapshot_manager,
                                      const std::shared_ptr<FileSystem>& file_system,
                                      const std::shared_ptr<ScanFilter>& scan_filter,
                                      int64_t read_view_ttl_millis)
     : disk_scan_(std::move(disk_scan)),
+      pk_table_(pk_table),
       realtime_context_(realtime_context),
       path_factory_(path_factory),
       snapshot_manager_(snapshot_manager),

src/paimon/core/table/source/realtime_table_scan.h

@@ -35,10 +35,10 @@ class FileSystem;
 class ScanFilter;
 class SnapshotManager;
 
-/// Adds process-local memory splits to a normal append-table batch scan.
+/// Adds process-local memory splits to a normal data-table batch scan.
 class RealtimeTableScan : public TableScan {
  public:
-    RealtimeTableScan(std::unique_ptr<TableScan>&& disk_scan,
+    RealtimeTableScan(std::unique_ptr<TableScan>&& disk_scan, bool pk_table,
                       const std::shared_ptr<RealtimeContextImpl>& realtime_context,
                       const std::shared_ptr<FileStorePathFactory>& path_factory,
                       const std::shared_ptr<SnapshotManager>& snapshot_manager,
@@ -71,6 +71,7 @@ class RealtimeTableScan : public TableScan {
         const std::optional<int64_t>& snapshot_id) const;
 
     std::unique_ptr<TableScan> disk_scan_;
+    bool pk_table_;
     std::shared_ptr<RealtimeContextImpl> realtime_context_;
     std::shared_ptr<FileStorePathFactory> path_factory_;
     std::shared_ptr<SnapshotManager> snapshot_manager_;

src/paimon/core/table/source/table_scan.cpp

@@ -64,6 +64,7 @@
 #include "paimon/core/utils/field_mapping.h"
 #include "paimon/core/utils/file_store_path_factory.h"
 #include "paimon/core/utils/index_file_path_factories.h"
+#include "paimon/core/utils/primary_key_table_utils.h"
 #include "paimon/core/utils/snapshot_manager.h"
 #include "paimon/format/file_format.h"
 #include "paimon/realtime/realtime_context.h"
@@ -223,22 +224,27 @@ Result<std::unique_ptr<TableScan>> TableScan::Create(std::unique_ptr<ScanContext
 namespace {
 
 Status ValidateRealtimeScan(const TableSchema& table_schema, const CoreOptions& core_options,
-                            const ScanContext& context) {
+                            const ScanContext& context, bool read_optimized) {
     if (!context.GetRealtimeContext()) {
         return Status::OK();
     }
     if (!core_options.RealtimeEnabled()) {
         return Status::Invalid("real-time scan requires realtime.enabled=true");
     }
-    if (!table_schema.PrimaryKeys().empty()) {
-        return Status::Invalid("real-time union read currently supports append tables only");
-    }
     if (core_options.GetBucket() <= 0) {
         return Status::Invalid("real-time union read requires fixed bucket mode");
     }
     if (core_options.DataEvolutionEnabled()) {
         return Status::Invalid("real-time union read does not support data evolution");
     }
+    if (!table_schema.PrimaryKeys().empty()) {
+        if (read_optimized) {
+            return Status::NotImplemented(
+                "PK real-time union read does not support read-optimized scans");
+        }
+        PAIMON_RETURN_NOT_OK(
+            PrimaryKeyTableUtils::ValidateRealtimeOptions(core_options, table_schema));
+    }
     if (context.IsStreamingMode()) {
         return Status::Invalid("real-time union read currently supports batch scans only");
     }
@@ -286,7 +292,8 @@ Result<std::unique_ptr<TableScan>> NewDataTableScan(const std::shared_ptr<ScanCo
                            CoreOptions::FromMap(options, context->GetSpecificFileSystem(), {}));
     core_options.WithCache(context->GetCache());
 
-    PAIMON_RETURN_NOT_OK(ValidateRealtimeScan(*table_schema, core_options, *context));
+    PAIMON_RETURN_NOT_OK(
+        ValidateRealtimeScan(*table_schema, core_options, *context, read_optimized));
     // validate options
     if (core_options.GetBucket() == -1) {
         if (!table_schema->PrimaryKeys().empty()) {
@@ -348,7 +355,7 @@ Result<std::unique_ptr<TableScan>> NewDataTableScan(const std::shared_ptr<ScanCo
         PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<RealtimeContextImpl> realtime_context,
                                RealtimeContextImpl::Cast(context->GetRealtimeContext()));
         return std::make_unique<RealtimeTableScan>(
-            std::move(batch_scan), realtime_context, path_factory,
+            std::move(batch_scan), pk_table, realtime_context, path_factory,
             snapshot_reader->GetSnapshotManager(), core_options.GetFileSystem(),
             context->GetScanFilters(), core_options.GetRealtimeReadViewTtlMillis());
     }

src/paimon/core/table/system/read_optimized_system_table.cpp

@@ -58,6 +58,10 @@ std::map<std::string, std::string> ReadOptimizedSystemTable::ReadOptimizedOption
 
 Result<std::unique_ptr<TableScan>> ReadOptimizedSystemTable::NewScan(
     const std::shared_ptr<ScanContext>& context) const {
+    if (context->GetRealtimeContext() && !table_schema_->PrimaryKeys().empty()) {
+        return Status::NotImplemented(
+            "PK real-time union read does not support read-optimized scans");
+    }
     auto options = ReadOptimizedOptions();
     ScanContextBuilder builder(table_path_);
     builder.SetOptions(options)