/*
 * 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 <algorithm>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <deque>
#include <iterator>
#include <map>
#include <memory>
#include <mutex>
#include <string>
#include <thread>
#include <tuple>
#include <utility>
#include <vector>

#include "arrow/api.h"
#include "arrow/c/bridge.h"
#include "arrow/ipc/json_simple.h"
#include "gtest/gtest.h"
#include "paimon/catalog/catalog.h"
#include "paimon/catalog/identifier.h"
#include "paimon/commit_context.h"
#include "paimon/common/utils/path_util.h"
#include "paimon/core/core_options.h"
#include "paimon/core/operation/commit/realtime_commit_properties.h"
#include "paimon/core/table/sink/commit_message_impl.h"
#include "paimon/core/utils/snapshot_manager.h"
#include "paimon/defs.h"
#include "paimon/file_store_commit.h"
#include "paimon/file_store_write.h"
#include "paimon/memory/memory_pool.h"
#include "paimon/predicate/predicate_builder.h"
#include "paimon/read_context.h"
#include "paimon/realtime/mem_indexer.h"
#include "paimon/realtime/realtime_context.h"
#include "paimon/record_batch.h"
#include "paimon/scan_context.h"
#include "paimon/table/source/table_read.h"
#include "paimon/table/source/table_scan.h"
#include "paimon/testing/utils/read_result_collector.h"
#include "paimon/testing/utils/testharness.h"
#include "paimon/write_context.h"

namespace paimon::test {

class ConcurrentTestState {
 public:
    void WaitForStart() {
        ready_threads_.fetch_add(1, std::memory_order_release);
        while (!start_.load(std::memory_order_acquire)) {
            std::this_thread::yield();
        }
    }

    void StartWhenReady(int32_t worker_count) {
        while (ready_threads_.load(std::memory_order_acquire) < worker_count) {
            std::this_thread::yield();
        }
        start_.store(true, std::memory_order_release);
    }

    void RecordError(const Status& status) {
        RecordError(status.ToString());
    }

    void RecordError(std::string error) {
        {
            std::lock_guard<std::mutex> lock(mutex);
            errors_.push_back(std::move(error));
        }
        stop_.store(true, std::memory_order_release);
        progress_cv.notify_all();
        snapshot_cv.notify_all();
    }

    bool RecordErrorIfNotOk(const Status& status) {
        if (status.ok()) {
            return false;
        }
        RecordError(status);
        return true;
    }

    template <typename T>
    bool RecordErrorIfNotOk(const Result<T>& result) {
        if (result.ok()) {
            return false;
        }
        RecordError(result.status());
        return true;
    }

    bool ShouldStop() const {
        return stop_.load(std::memory_order_acquire);
    }

    const std::vector<std::string>& Errors() const {
        return errors_;
    }

    std::mutex mutex;
    std::condition_variable progress_cv;
    std::condition_variable snapshot_cv;

 private:
    std::atomic<bool> start_{false};
    std::atomic<bool> stop_{false};
    std::atomic<int32_t> ready_threads_{0};
    std::vector<std::string> errors_;
};

class RealtimeWriteInteTest : public ::testing::Test {
 protected:
    using Row = std::tuple<int64_t, std::string, std::string>;

    struct CollectedReadResult {
        std::unique_ptr<BatchReader> reader;
        std::shared_ptr<arrow::ChunkedArray> data;
    };

    void SetUp() override {
        pool_ = GetDefaultPool();
        dir_ = UniqueTestDirectory::Create("local");
        ASSERT_NE(nullptr, dir_);
        table_path_ = PathUtil::JoinPath(dir_->Str(), "foo.db/bar");
        fields_ = {arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8()),
                   arrow::field("pt", arrow::utf8())};
        schema_ = arrow::schema(fields_);
        options_ = {
            {Options::MANIFEST_FORMAT, "orc"}, {Options::FILE_FORMAT, "orc"},
            {Options::FILE_SYSTEM, "local"},   {Options::BUCKET, "1"},
            {Options::BUCKET_KEY, "id"},       {Options::TARGET_FILE_SIZE, "1048576"},
        };
    }

    void TearDown() override {
        dir_.reset();
    }

    void CreateTable(const std::vector<std::string>& partition_keys) const {
        auto c_schema = std::make_unique<ArrowSchema>();
        ASSERT_TRUE(arrow::ExportSchema(*schema_, c_schema.get()).ok());
        ASSERT_OK_AND_ASSIGN(std::unique_ptr<Catalog> catalog,
                             Catalog::Create(dir_->Str(), options_));
        ASSERT_OK(catalog->CreateDatabase("foo", {}, /*ignore_if_exists=*/false));
        ASSERT_OK(catalog->CreateTable(Identifier("foo", "bar"), c_schema.get(), partition_keys,
                                       /*primary_keys=*/{}, options_,
                                       /*ignore_if_exists=*/false));
    }

    Result<std::unique_ptr<FileStoreWrite>> CreateRealtimeWriter(
        const std::shared_ptr<RealtimeContext>& realtime_context) const {
        WriteContextBuilder builder(table_path_, commit_user_);
        builder.SetOptions(options_).WithStreamingMode(true).WithRealtimeContext(realtime_context);
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<WriteContext> context, builder.Finish());
        return FileStoreWrite::Create(std::move(context));
    }

    Result<std::unique_ptr<FileStoreWrite>> CreateRealtimeWriter() const {
        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<RealtimeContext> realtime_context,
                               RealtimeContext::Create());
        return CreateRealtimeWriter(realtime_context);
    }

    Result<std::unique_ptr<RecordBatch>> MakeBatch(const std::vector<Row>& rows,
                                                   bool partitioned) const {
        return MakeBatch(rows, partitioned, /*bucket=*/0);
    }

    Result<std::unique_ptr<RecordBatch>> MakeBatch(const std::vector<Row>& rows, bool partitioned,
                                                   int32_t bucket) const {
        if (rows.empty()) {
            return Status::Invalid("cannot create an empty test batch");
        }
        const std::string& partition = std::get<2>(rows.front());
        std::string json = "[";
        for (size_t i = 0; i < rows.size(); ++i) {
            const auto& [id, payload, pt] = rows[i];
            if (pt != partition) {
                return Status::Invalid("one test batch must contain only one partition");
            }
            if (i > 0) {
                json += ",";
            }
            json += "[" + std::to_string(id) + ",\"" + payload + "\",\"" + pt + "\"]";
        }
        json += "]";

        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
            std::shared_ptr<arrow::Array> array,
            arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), json));
        ArrowArray c_array;
        PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array));
        RecordBatchBuilder builder(&c_array);
        if (partitioned) {
            builder.SetPartition({{"pt", partition}});
        }
        return builder.SetBucket(bucket).Finish();
    }

    static std::vector<Row> MakeRows(int64_t first_id, int64_t count,
                                     const std::string& partition) {
        std::vector<Row> rows;
        rows.reserve(count);
        for (int64_t i = 0; i < count; ++i) {
            int64_t id = first_id + i;
            rows.emplace_back(id, "value-" + std::to_string(id), partition);
        }
        return rows;
    }

    Result<int64_t> Commit(const std::vector<RealtimeCommitProgress>& realtime_commits,
                           int64_t commit_identifier) const {
        CommitContextBuilder builder(table_path_, commit_user_);
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<CommitContext> context,
                               builder.SetOptions(options_).Finish());
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FileStoreCommit> commit,
                               FileStoreCommit::Create(std::move(context)));
        return commit->CommitWithProgress(realtime_commits, commit_identifier,
                                          /*watermark=*/std::nullopt);
    }

    Result<std::shared_ptr<Plan>> CreatePlan(
        const std::shared_ptr<RealtimeContext>& realtime_context,
        const std::shared_ptr<Predicate>& predicate) const {
        ScanContextBuilder scan_builder(table_path_);
        if (realtime_context) {
            scan_builder.WithRealtimeContext(realtime_context);
        }
        scan_builder.SetPredicate(predicate).WithMemoryPool(pool_);
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ScanContext> scan_context,
                               scan_builder.SetOptions(options_).Finish());
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableScan> scan,
                               TableScan::Create(std::move(scan_context)));
        return scan->CreatePlan();
    }

    Result<CollectedReadResult> ReadPlan(const std::shared_ptr<Plan>& plan,
                                         const std::vector<std::string>& read_fields,
                                         const std::shared_ptr<Predicate>& predicate,
                                         bool enable_predicate_filter) const {
        ReadContextBuilder read_builder(table_path_);
        read_builder.SetOptions(options_)
            .SetReadFieldNames(read_fields)
            .SetPredicate(predicate)
            .EnablePredicateFilter(enable_predicate_filter)
            .WithMemoryPool(pool_);
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> read_context, read_builder.Finish());
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableRead> table_read,
                               TableRead::Create(std::move(read_context)));
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BatchReader> reader,
                               table_read->CreateReader(plan->Splits()));
        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::ChunkedArray> result,
                               ReadResultCollector::CollectResult(reader.get()));
        return CollectedReadResult{std::move(reader), std::move(result)};
    }

    Result<std::vector<Row>> ReadRows(
        const std::shared_ptr<RealtimeContext>& realtime_context) const {
        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Plan> plan,
                               CreatePlan(realtime_context, /*predicate=*/nullptr));
        return ReadRows(plan);
    }

    Result<std::vector<Row>> ReadRows(const std::shared_ptr<Plan>& plan) const {
        PAIMON_ASSIGN_OR_RAISE(CollectedReadResult read_result,
                               ReadPlan(plan, {"id", "payload", "pt"}, /*predicate=*/nullptr,
                                        /*enable_predicate_filter=*/false));
        const std::shared_ptr<arrow::ChunkedArray>& result = read_result.data;

        std::vector<Row> rows;
        if (!result) {
            return rows;
        }
        for (const std::shared_ptr<arrow::Array>& chunk : result->chunks()) {
            std::shared_ptr<arrow::StructArray> data =
                std::dynamic_pointer_cast<arrow::StructArray>(chunk);
            if (!data || data->num_fields() != 4) {
                return Status::Invalid("unexpected real-time test read schema");
            }
            std::shared_ptr<arrow::Int8Array> row_kinds =
                std::dynamic_pointer_cast<arrow::Int8Array>(data->field(0));
            std::shared_ptr<arrow::Int64Array> ids =
                std::dynamic_pointer_cast<arrow::Int64Array>(data->field(1));
            std::shared_ptr<arrow::StringArray> payloads =
                std::dynamic_pointer_cast<arrow::StringArray>(data->field(2));
            std::shared_ptr<arrow::StringArray> partitions =
                std::dynamic_pointer_cast<arrow::StringArray>(data->field(3));
            if (!row_kinds || !ids || !payloads || !partitions) {
                return Status::Invalid("unexpected real-time test read field type");
            }
            for (int64_t i = 0; i < data->length(); ++i) {
                if (row_kinds->IsNull(i) ||
                    row_kinds->Value(i) != static_cast<int8_t>(RecordBatch::RowKind::INSERT) ||
                    ids->IsNull(i) || payloads->IsNull(i) || partitions->IsNull(i)) {
                    return Status::Invalid("unexpected null or row kind in real-time test result");
                }
                rows.emplace_back(ids->Value(i), payloads->GetString(i), partitions->GetString(i));
            }
        }
        return rows;
    }

    Result<std::vector<Row>> ReadRows() const {
        return ReadRows(std::shared_ptr<RealtimeContext>());
    }

    Result<uint64_t> GetRealtimeMemoryUsage(
        const std::shared_ptr<RealtimeContext>& realtime_context) const {
        PAIMON_ASSIGN_OR_RAISE(std::vector<RealtimePartitionBucketView> views,
                               realtime_context->AcquireReadViews());
        uint64_t memory_usage = 0;
        for (const RealtimePartitionBucketView& view : views) {
            memory_usage += view.indexer->GetMemoryUsage();
        }
        return memory_usage;
    }

    static Status ValidateReadPrefix(const std::vector<Row>& rows, int64_t total_rows) {
        std::vector<bool> seen(static_cast<size_t>(total_rows), false);
        int64_t max_id = -1;
        for (const Row& row : rows) {
            const auto& [id, payload, partition] = row;
            if (id < 0 || id >= total_rows) {
                return Status::Invalid("real-time read id is out of range");
            }
            if (seen[static_cast<size_t>(id)]) {
                return Status::Invalid("real-time read contains duplicate ids");
            }
            if (payload != "value-" + std::to_string(id) || partition != "p0") {
                return Status::Invalid("real-time read row does not match its id");
            }
            seen[static_cast<size_t>(id)] = true;
            max_id = std::max(max_id, id);
        }
        for (int64_t id = 0; id <= max_id; ++id) {
            if (!seen[static_cast<size_t>(id)]) {
                return Status::Invalid("real-time read contains an id gap");
            }
        }
        return Status::OK();
    }

    Result<RealtimeOffsetMap> ReadCommittedOffsets() const {
        PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_));
        SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_);
        PAIMON_ASSIGN_OR_RAISE(std::optional<Snapshot> snapshot, snapshot_manager.LatestSnapshot());
        return RealtimeCommitProperties::ReadOffsets(snapshot, options.GetFileSystem());
    }

    void FinalizeCommitAndCheck(FileStoreWrite* writer,
                                std::vector<RealtimeCommitProgress> realtime_commits,
                                int64_t prepare_identifier, std::vector<Row> expected_rows) const {
        ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> final_commits,
                             writer->PrepareCommitWithProgress(prepare_identifier));
        realtime_commits.insert(realtime_commits.end(),
                                std::make_move_iterator(final_commits.begin()),
                                std::make_move_iterator(final_commits.end()));
        ASSERT_OK(Commit(realtime_commits, prepare_identifier));
        ASSERT_OK(writer->Close());

        ASSERT_OK_AND_ASSIGN(std::vector<Row> actual_rows, ReadRows());
        ASSERT_EQ(expected_rows, actual_rows);
    }

    std::unique_ptr<UniqueTestDirectory> dir_;
    std::string table_path_;
    std::string commit_user_ = "realtime_commit_user";
    arrow::FieldVector fields_;
    std::shared_ptr<arrow::Schema> schema_;
    std::map<std::string, std::string> options_;
    std::shared_ptr<MemoryPool> pool_;
};

TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer, CreateRealtimeWriter());
    std::vector<Row> rows = MakeRows(/*first_id=*/0, /*count=*/10, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch,
                         MakeBatch(rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(batch)));
    FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows);
}

TEST_F(RealtimeWriteInteTest, TestRollingFilesPreserveProgress) {
    options_[Options::TARGET_FILE_ROW_NUM] = "10";
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer, CreateRealtimeWriter());

    std::vector<Row> expected_rows;
    constexpr int64_t kBatchCount = 3;
    constexpr int64_t kRowsPerBatch = 10;
    for (int64_t batch_index = 0; batch_index < kBatchCount; ++batch_index) {
        std::vector<Row> rows =
            MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, /*partition=*/"p0");
        ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch,
                             MakeBatch(rows, /*partitioned=*/false));
        ASSERT_OK(writer->Write(std::move(batch)));
        expected_rows.insert(expected_rows.end(), rows.begin(), rows.end());
    }

    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> commits,
                         writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
    ASSERT_EQ(1, commits.size());
    ASSERT_EQ(Range(0, kBatchCount * kRowsPerBatch - 1), commits[0].offset_range);
    std::shared_ptr<CommitMessageImpl> commit_message =
        std::dynamic_pointer_cast<CommitMessageImpl>(commits[0].commit_message);
    ASSERT_NE(nullptr, commit_message);
    ASSERT_EQ(3, commit_message->GetNewFilesIncrement().NewFiles().size());
    ASSERT_OK(Commit(commits, /*commit_identifier=*/0));
    ASSERT_OK(writer->Close());

    ASSERT_OK_AND_ASSIGN(std::vector<Row> actual_rows, ReadRows());
    ASSERT_EQ(expected_rows, actual_rows);
}

TEST_F(RealtimeWriteInteTest, TestCommitOrdersPreparedOffsetRanges) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer, CreateRealtimeWriter());

    std::vector<Row> first_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> first_batch,
                         MakeBatch(first_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(first_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> commits,
                         writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
    ASSERT_EQ(1, commits.size());
    ASSERT_EQ(Range(0, 2), commits[0].offset_range);

    std::vector<Row> second_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> second_batch,
                         MakeBatch(second_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(second_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> second_commits,
                         writer->PrepareCommitWithProgress(/*commit_identifier=*/1));
    ASSERT_EQ(1, second_commits.size());
    ASSERT_EQ(Range(3, 4), second_commits[0].offset_range);

    commits.push_back(std::move(second_commits[0]));
    std::reverse(commits.begin(), commits.end());
    ASSERT_OK(Commit(commits, /*commit_identifier=*/1));
    ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, ReadCommittedOffsets());
    ASSERT_EQ(4, committed_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0)));
    ASSERT_OK(writer->Close());

    std::vector<Row> expected_rows = first_rows;
    expected_rows.insert(expected_rows.end(), second_rows.begin(), second_rows.end());
    ASSERT_OK_AND_ASSIGN(std::vector<Row> actual_rows, ReadRows());
    ASSERT_EQ(expected_rows, actual_rows);
}

TEST_F(RealtimeWriteInteTest, TestReadMemoryBeforePrepareCommit) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
                         CreateRealtimeWriter(realtime_context));
    std::vector<Row> rows = MakeRows(/*first_id=*/0, /*count=*/10, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch,
                         MakeBatch(rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(batch)));

    ASSERT_OK_AND_ASSIGN(std::vector<Row> actual_rows, ReadRows(realtime_context));
    ASSERT_EQ(rows, actual_rows);
    ASSERT_OK(writer->Close());
}

TEST_F(RealtimeWriteInteTest, TestCloseWriterKeepsContextReadable) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
                         CreateRealtimeWriter(realtime_context));
    std::vector<Row> rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch,
                         MakeBatch(rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(batch)));
    ASSERT_OK(writer->Close());

    ASSERT_OK_AND_ASSIGN(std::vector<Row> actual_rows, ReadRows(realtime_context));
    ASSERT_EQ(rows, actual_rows);
}

TEST_F(RealtimeWriteInteTest, TestPinnedPlanRemainsReadableAfterWriterClose) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
                         CreateRealtimeWriter(realtime_context));
    std::vector<Row> rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch,
                         MakeBatch(rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(batch)));
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan,
                         CreatePlan(realtime_context, /*predicate=*/nullptr));
    ASSERT_OK(writer->Close());

    ASSERT_OK_AND_ASSIGN(std::vector<Row> actual_rows, ReadRows(plan));
    ASSERT_EQ(rows, actual_rows);
}

TEST_F(RealtimeWriteInteTest, TestCloseWriterAllowsContextReuseByLaterWriter) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> first_writer,
                         CreateRealtimeWriter(realtime_context));
    std::vector<Row> first_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> first_batch,
                         MakeBatch(first_rows, /*partitioned=*/false));
    ASSERT_OK(first_writer->Write(std::move(first_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> commits,
                         first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
    ASSERT_EQ(1, commits.size());
    ASSERT_EQ(Range(0, 2), commits[0].offset_range);
    ASSERT_OK(first_writer->Close());

    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> second_writer,
                         CreateRealtimeWriter(realtime_context));
    std::vector<Row> second_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> second_batch,
                         MakeBatch(second_rows, /*partitioned=*/false));
    ASSERT_OK(second_writer->Write(std::move(second_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> second_commits,
                         second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1));
    ASSERT_EQ(1, second_commits.size());
    ASSERT_EQ(Range(3, 4), second_commits[0].offset_range);

    commits.push_back(std::move(second_commits[0]));
    ASSERT_OK(Commit(commits, /*commit_identifier=*/1));
    std::vector<Row> expected_rows = first_rows;
    expected_rows.insert(expected_rows.end(), second_rows.begin(), second_rows.end());
    ASSERT_OK_AND_ASSIGN(std::vector<Row> actual_rows, ReadRows(realtime_context));
    ASSERT_EQ(expected_rows, actual_rows);
    ASSERT_OK(second_writer->Close());
}

TEST_F(RealtimeWriteInteTest, TestReadCommittedDiskAndBuildingMemory) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
                         CreateRealtimeWriter(realtime_context));

    std::vector<Row> disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> disk_batch,
                         MakeBatch(disk_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(disk_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> disk_commits,
                         writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
    ASSERT_EQ(1, disk_commits.size());
    ASSERT_EQ(Range(0, 2), disk_commits[0].offset_range);

    std::vector<Row> memory_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> memory_batch,
                         MakeBatch(memory_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(memory_batch)));
    ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0));

    std::vector<Row> expected_rows = disk_rows;
    expected_rows.insert(expected_rows.end(), memory_rows.begin(), memory_rows.end());
    ASSERT_OK_AND_ASSIGN(std::vector<Row> actual_rows, ReadRows(realtime_context));
    ASSERT_EQ(expected_rows, actual_rows);
    ASSERT_OK(writer->Close());
}

TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
                         CreateRealtimeWriter(realtime_context));
    std::shared_ptr<Predicate> scan_predicate =
        PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT,
                                      Literal(static_cast<int64_t>(1)));
    std::shared_ptr<Predicate> read_predicate =
        PredicateBuilder::GreaterThan(/*field_index=*/1, /*field_name=*/"id", FieldType::BIGINT,
                                      Literal(static_cast<int64_t>(1)));
    const std::vector<std::string> read_fields = {"payload", "id"};
    std::shared_ptr<arrow::DataType> result_type = arrow::struct_(
        {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("payload", arrow::utf8()),
         arrow::field("id", arrow::int64())});

    std::vector<Row> disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> disk_batch,
                         MakeBatch(disk_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(disk_batch)));

    ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> memory_plan,
                         CreatePlan(realtime_context, scan_predicate));
    ASSERT_OK_AND_ASSIGN(
        CollectedReadResult memory_result,
        ReadPlan(memory_plan, read_fields, read_predicate, /*enable_predicate_filter=*/true));
    std::shared_ptr<arrow::Array> expected_memory =
        arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([
            [0, "value-2", 2]
        ])")
            .ValueOrDie();
    ASSERT_NE(nullptr, memory_result.data);
    ASSERT_TRUE(
        std::make_shared<arrow::ChunkedArray>(expected_memory)->Equals(*memory_result.data));

    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> disk_commits,
                         writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
    ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0));
    std::vector<Row> memory_rows = MakeRows(/*first_id=*/3, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> memory_batch,
                         MakeBatch(memory_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(memory_batch)));

    ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> union_plan,
                         CreatePlan(realtime_context, scan_predicate));
    ASSERT_OK_AND_ASSIGN(
        CollectedReadResult union_result,
        ReadPlan(union_plan, read_fields, read_predicate, /*enable_predicate_filter=*/true));
    std::shared_ptr<arrow::Array> expected_union =
        arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([
            [0, "value-2", 2],
            [0, "value-3", 3],
            [0, "value-4", 4],
            [0, "value-5", 5]
        ])")
            .ValueOrDie();
    ASSERT_NE(nullptr, union_result.data);
    ASSERT_TRUE(std::make_shared<arrow::ChunkedArray>(expected_union)->Equals(*union_result.data));
    ASSERT_OK(writer->Close());
}

TEST_F(RealtimeWriteInteTest, TestDiskPredicatePushdownWithoutMemoryFiltering) {
    options_[Options::FILE_FORMAT] = "parquet";
    options_[Options::WRITE_BATCH_SIZE] = "1";
    options_["parquet.page.size"] = "1";
    options_["parquet.enable-dictionary"] = "false";
    options_["parquet.write.enable-page-index"] = "true";
    options_["parquet.read.enable-page-index-filter"] = "true";
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
                         CreateRealtimeWriter(realtime_context));
    std::shared_ptr<Predicate> predicate =
        PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT,
                                Literal(static_cast<int64_t>(1)));

    std::vector<Row> disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> disk_batch,
                         MakeBatch(disk_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(disk_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> disk_commits,
                         writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
    ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0));

    std::vector<Row> memory_rows = MakeRows(/*first_id=*/3, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> memory_batch,
                         MakeBatch(memory_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(memory_batch)));

    ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan, CreatePlan(realtime_context, predicate));
    ASSERT_OK_AND_ASSIGN(CollectedReadResult result,
                         ReadPlan(plan, {"id", "payload", "pt"}, predicate,
                                  /*enable_predicate_filter=*/false));
    std::shared_ptr<arrow::DataType> result_type = arrow::struct_(
        {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("id", arrow::int64()),
         arrow::field("payload", arrow::utf8()), arrow::field("pt", arrow::utf8())});
    std::shared_ptr<arrow::Array> expected =
        arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([
            [0, 1, "value-1", "p0"],
            [0, 3, "value-3", "p0"],
            [0, 4, "value-4", "p0"],
            [0, 5, "value-5", "p0"]
        ])")
            .ValueOrDie();
    ASSERT_NE(nullptr, result.data);
    ASSERT_TRUE(std::make_shared<arrow::ChunkedArray>(expected)->Equals(*result.data))
        << result.data->ToString();
    ASSERT_OK(writer->Close());
}

TEST_F(RealtimeWriteInteTest, TestRefreshCommittedSnapshotReclaimsMemory) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
                         CreateRealtimeWriter(realtime_context));

    std::vector<Row> disk_rows = MakeRows(/*first_id=*/0, /*count=*/10, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> disk_batch,
                         MakeBatch(disk_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(disk_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> disk_commits,
                         writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
    ASSERT_EQ(1, disk_commits.size());
    ASSERT_OK_AND_ASSIGN(int64_t committed_snapshot_id,
                         Commit(disk_commits, /*commit_identifier=*/0));

    std::vector<Row> memory_rows = MakeRows(/*first_id=*/10, /*count=*/2, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> memory_batch,
                         MakeBatch(memory_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(memory_batch)));

    std::vector<Row> expected_rows = disk_rows;
    expected_rows.insert(expected_rows.end(), memory_rows.begin(), memory_rows.end());
    ASSERT_OK_AND_ASSIGN(std::vector<Row> read1, ReadRows(realtime_context));
    ASSERT_EQ(expected_rows, read1);
    ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_before_refresh,
                         GetRealtimeMemoryUsage(realtime_context));

    ASSERT_OK(writer->RefreshCommittedSnapshot(committed_snapshot_id));

    ASSERT_OK_AND_ASSIGN(std::vector<Row> read2, ReadRows(realtime_context));
    ASSERT_EQ(read1, read2);
    ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_after_refresh,
                         GetRealtimeMemoryUsage(realtime_context));
    ASSERT_LT(memory_usage_after_refresh, memory_usage_before_refresh);
    ASSERT_OK(writer->Close());
}

TEST_F(RealtimeWriteInteTest, TestPlanPinsMemoryAcrossRefresh) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
                         CreateRealtimeWriter(realtime_context));

    std::vector<Row> disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> disk_batch,
                         MakeBatch(disk_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(disk_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> disk_commits,
                         writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
    ASSERT_OK_AND_ASSIGN(int64_t committed_snapshot_id,
                         Commit(disk_commits, /*commit_identifier=*/0));

    std::vector<Row> memory_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> memory_batch,
                         MakeBatch(memory_rows, /*partitioned=*/false));
    ASSERT_OK(writer->Write(std::move(memory_batch)));
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> pinned_plan,
                         CreatePlan(realtime_context, /*predicate=*/nullptr));

    ASSERT_OK(writer->RefreshCommittedSnapshot(committed_snapshot_id));

    std::vector<Row> expected_rows = disk_rows;
    expected_rows.insert(expected_rows.end(), memory_rows.begin(), memory_rows.end());
    ASSERT_OK_AND_ASSIGN(std::vector<Row> pinned_rows, ReadRows(pinned_plan));
    ASSERT_EQ(expected_rows, pinned_rows);

    ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> refreshed_plan,
                         CreatePlan(realtime_context, /*predicate=*/nullptr));
    ASSERT_OK_AND_ASSIGN(std::vector<Row> refreshed_rows, ReadRows(refreshed_plan));
    ASSERT_EQ(pinned_rows, refreshed_rows);
    ASSERT_OK(writer->Close());
}

TEST_F(RealtimeWriteInteTest, TestRepeatedCommitReadAndRefresh) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
                         CreateRealtimeWriter(realtime_context));
    constexpr int64_t kRoundCount = 3;
    constexpr int64_t kRowsPerRound = 4;
    std::vector<Row> expected_rows;
    for (int64_t round = 0; round < kRoundCount; ++round) {
        std::vector<Row> rows = MakeRows(round * kRowsPerRound, kRowsPerRound, /*partition=*/"p0");
        ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch,
                             MakeBatch(rows, /*partitioned=*/false));
        ASSERT_OK(writer->Write(std::move(batch)));
        ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> commits,
                             writer->PrepareCommitWithProgress(/*commit_identifier=*/round));
        ASSERT_EQ(1, commits.size());
        ASSERT_EQ(Range(round * kRowsPerRound, (round + 1) * kRowsPerRound - 1),
                  commits[0].offset_range);
        ASSERT_OK_AND_ASSIGN(int64_t committed_snapshot_id,
                             Commit(commits, /*commit_identifier=*/round));
        expected_rows.insert(expected_rows.end(), rows.begin(), rows.end());

        ASSERT_OK_AND_ASSIGN(std::vector<Row> read_before_refresh, ReadRows(realtime_context));
        ASSERT_EQ(expected_rows, read_before_refresh);
        ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_before_refresh,
                             GetRealtimeMemoryUsage(realtime_context));
        ASSERT_GT(memory_usage_before_refresh, 0);

        ASSERT_OK(writer->RefreshCommittedSnapshot(committed_snapshot_id));

        ASSERT_OK_AND_ASSIGN(std::vector<Row> read_after_refresh, ReadRows(realtime_context));
        ASSERT_EQ(read_before_refresh, read_after_refresh);
        ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_after_refresh,
                             GetRealtimeMemoryUsage(realtime_context));
        ASSERT_EQ(0, memory_usage_after_refresh);
    }
    ASSERT_OK(writer->Close());
}

TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) {
    CreateTable(/*partition_keys=*/{});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
                         CreateRealtimeWriter(realtime_context));
    constexpr int32_t kPrepareThreadCount = 4;
    constexpr int32_t kReadThreadCount = 4;
    constexpr int64_t kBatchCount = 12;
    constexpr int64_t kRowsPerBatch = 2;
    constexpr int64_t kTotalRows = kBatchCount * kRowsPerBatch;

    std::atomic<bool> writer_done{false};
    std::atomic<bool> prepare_done{false};
    std::atomic<bool> commit_done{false};
    std::atomic<bool> refresh_done{false};
    std::atomic<int64_t> next_prepare_identifier{0};
    std::atomic<int32_t> commit_count{0};
    std::atomic<int32_t> refresh_count{0};
    ConcurrentTestState state;
    std::map<int64_t, RealtimeCommitProgress> pending_commits;
    std::deque<int64_t> pending_snapshot_ids;
    std::vector<int32_t> prepare_call_counts(kPrepareThreadCount, 0);
    std::vector<int32_t> read_call_counts(kReadThreadCount, 0);

    auto enqueue_prepared_commits = [&](std::vector<RealtimeCommitProgress>&& commits) {
        std::string error;
        {
            std::lock_guard<std::mutex> lock(state.mutex);
            for (RealtimeCommitProgress& commit : commits) {
                int64_t offset_from = commit.offset_range.from;
                if (!pending_commits.emplace(offset_from, std::move(commit)).second) {
                    error = "duplicate prepared real-time offset range";
                    break;
                }
            }
        }
        if (!error.empty()) {
            state.RecordError(error);
        }
        state.progress_cv.notify_all();
    };

    std::thread write_thread([&]() {
        state.WaitForStart();
        for (int64_t batch_index = 0; batch_index < kBatchCount && !state.ShouldStop();
             ++batch_index) {
            std::vector<Row> rows =
                MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, /*partition=*/"p0");
            Result<std::unique_ptr<RecordBatch>> batch_result =
                MakeBatch(rows, /*partitioned=*/false);
            if (state.RecordErrorIfNotOk(batch_result)) {
                break;
            }
            Status status = writer->Write(std::move(batch_result).value());
            if (state.RecordErrorIfNotOk(status)) {
                break;
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(2));
            // Pause midway until one refresh completes to guarantee write and refresh overlap.
            if (batch_index + 1 == kBatchCount / 2) {
                auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30);
                while (refresh_count.load(std::memory_order_acquire) == 0 && !state.ShouldStop() &&
                       std::chrono::steady_clock::now() < deadline) {
                    std::this_thread::sleep_for(std::chrono::milliseconds(1));
                }
                if (refresh_count.load(std::memory_order_acquire) == 0 && !state.ShouldStop()) {
                    state.RecordError("timed out waiting for a refresh while writing");
                    break;
                }
            }
        }
        writer_done.store(true, std::memory_order_release);
    });

    std::vector<std::thread> prepare_threads;
    prepare_threads.reserve(kPrepareThreadCount);
    for (int32_t thread_index = 0; thread_index < kPrepareThreadCount; ++thread_index) {
        prepare_threads.emplace_back([&, thread_index]() {
            state.WaitForStart();
            do {
                int64_t identifier = next_prepare_identifier.fetch_add(1);
                Result<std::vector<RealtimeCommitProgress>> result =
                    writer->PrepareCommitWithProgress(identifier);
                ++prepare_call_counts[thread_index];
                if (state.RecordErrorIfNotOk(result)) {
                    break;
                }
                enqueue_prepared_commits(std::move(result).value());
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
            } while (!writer_done.load(std::memory_order_acquire) && !state.ShouldStop());
        });
    }

    std::thread commit_thread([&]() {
        state.WaitForStart();
        int64_t next_offset = 0;
        int64_t commit_identifier = 0;
        while (!state.ShouldStop()) {
            std::optional<RealtimeCommitProgress> next_commit;
            {
                std::unique_lock<std::mutex> lock(state.mutex);
                state.progress_cv.wait(lock, [&]() {
                    return state.ShouldStop() || pending_commits.count(next_offset) > 0 ||
                           prepare_done.load(std::memory_order_acquire);
                });
                if (state.ShouldStop()) {
                    break;
                }
                auto iter = pending_commits.find(next_offset);
                if (iter == pending_commits.end()) {
                    if (prepare_done.load(std::memory_order_acquire)) {
                        if (!pending_commits.empty()) {
                            lock.unlock();
                            state.RecordError("prepared real-time offset ranges contain a gap");
                        }
                        break;
                    }
                    continue;
                }
                next_commit = std::move(iter->second);
                pending_commits.erase(iter);
            }

            std::vector<RealtimeCommitProgress> commits;
            commits.push_back(std::move(next_commit).value());
            int64_t committed_offset = commits[0].offset_range.to;
            Result<int64_t> commit_result = Commit(commits, commit_identifier++);
            if (state.RecordErrorIfNotOk(commit_result)) {
                break;
            }
            next_offset = committed_offset + 1;
            {
                std::lock_guard<std::mutex> lock(state.mutex);
                pending_snapshot_ids.push_back(std::move(commit_result).value());
            }
            ++commit_count;
            state.snapshot_cv.notify_all();
        }
        commit_done.store(true, std::memory_order_release);
        state.snapshot_cv.notify_all();
    });

    std::thread refresh_thread([&]() {
        state.WaitForStart();
        while (!state.ShouldStop()) {
            std::optional<int64_t> snapshot_id;
            {
                std::unique_lock<std::mutex> lock(state.mutex);
                state.snapshot_cv.wait(lock, [&]() {
                    return state.ShouldStop() || !pending_snapshot_ids.empty() ||
                           commit_done.load(std::memory_order_acquire);
                });
                if (state.ShouldStop()) {
                    break;
                }
                if (pending_snapshot_ids.empty()) {
                    if (commit_done.load(std::memory_order_acquire)) {
                        break;
                    }
                    continue;
                }
                snapshot_id = pending_snapshot_ids.front();
                pending_snapshot_ids.pop_front();
            }
            Status status = writer->RefreshCommittedSnapshot(snapshot_id.value());
            if (state.RecordErrorIfNotOk(status)) {
                break;
            }
            ++refresh_count;
        }
        refresh_done.store(true, std::memory_order_release);
    });

    std::vector<std::thread> read_threads;
    read_threads.reserve(kReadThreadCount);
    for (int32_t thread_index = 0; thread_index < kReadThreadCount; ++thread_index) {
        read_threads.emplace_back([&, thread_index]() {
            state.WaitForStart();
            while (!refresh_done.load(std::memory_order_acquire) && !state.ShouldStop()) {
                Result<std::vector<Row>> result = ReadRows(realtime_context);
                ++read_call_counts[thread_index];
                if (state.RecordErrorIfNotOk(result)) {
                    break;
                }
                Status status = ValidateReadPrefix(result.value(), kTotalRows);
                if (state.RecordErrorIfNotOk(status)) {
                    break;
                }
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
            }
        });
    }

    constexpr int32_t kWorkerCount = 1 + kPrepareThreadCount + 1 + 1 + kReadThreadCount;
    state.StartWhenReady(kWorkerCount);

    write_thread.join();
    for (std::thread& thread : prepare_threads) {
        thread.join();
    }
    if (!state.ShouldStop()) {
        Result<std::vector<RealtimeCommitProgress>> final_result =
            writer->PrepareCommitWithProgress(next_prepare_identifier.fetch_add(1));
        if (!state.RecordErrorIfNotOk(final_result)) {
            enqueue_prepared_commits(std::move(final_result).value());
        }
    }
    prepare_done.store(true, std::memory_order_release);
    state.progress_cv.notify_all();

    commit_thread.join();
    refresh_thread.join();
    for (std::thread& thread : read_threads) {
        thread.join();
    }

    ASSERT_TRUE(state.Errors().empty()) << (state.Errors().empty() ? "" : state.Errors().front());
    for (int32_t call_count : prepare_call_counts) {
        ASSERT_GT(call_count, 0);
    }
    for (int32_t call_count : read_call_counts) {
        ASSERT_GT(call_count, 0);
    }
    ASSERT_GE(commit_count.load(), 2);
    ASSERT_GE(refresh_count.load(), 2);
    ASSERT_OK_AND_ASSIGN(std::vector<Row> final_rows, ReadRows(realtime_context));
    ASSERT_EQ(kTotalRows, static_cast<int64_t>(final_rows.size()));
    ASSERT_OK(ValidateReadPrefix(final_rows, kTotalRows));
    ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, ReadCommittedOffsets());
    ASSERT_EQ(kTotalRows - 1,
              committed_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0)));
    ASSERT_OK_AND_ASSIGN(uint64_t memory_usage, GetRealtimeMemoryUsage(realtime_context));
    ASSERT_EQ(0, memory_usage);
    ASSERT_OK(writer->Close());
}

TEST_F(RealtimeWriteInteTest, TestMultiplePartitions) {
    CreateTable(/*partition_keys=*/{"pt"});
    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer,
                         CreateRealtimeWriter(realtime_context));
    std::vector<std::vector<Row>> disk_rows;
    for (int64_t partition_index = 0; partition_index < 2; ++partition_index) {
        std::string partition = "p" + std::to_string(partition_index);
        std::vector<Row> rows = MakeRows(partition_index * 10, /*count=*/10, partition);
        ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch,
                             MakeBatch(rows, /*partitioned=*/true));
        ASSERT_OK(writer->Write(std::move(batch)));
        disk_rows.push_back(std::move(rows));
    }
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> disk_commits,
                         writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
    ASSERT_EQ(2, disk_commits.size());
    ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0));

    std::vector<Row> p0_memory_rows = MakeRows(/*first_id=*/20, /*count=*/5, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> p0_memory_batch,
                         MakeBatch(p0_memory_rows, /*partitioned=*/true));
    ASSERT_OK(writer->Write(std::move(p0_memory_batch)));
    std::vector<Row> p2_memory_rows = MakeRows(/*first_id=*/30, /*count=*/5, /*partition=*/"p2");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> p2_memory_batch,
                         MakeBatch(p2_memory_rows, /*partitioned=*/true));
    ASSERT_OK(writer->Write(std::move(p2_memory_batch)));

    // p0 has disk and memory rows, p1 is disk-only, and p2 is memory-only.
    std::vector<Row> expected_rows = disk_rows[0];
    expected_rows.insert(expected_rows.end(), p0_memory_rows.begin(), p0_memory_rows.end());
    expected_rows.insert(expected_rows.end(), disk_rows[1].begin(), disk_rows[1].end());
    expected_rows.insert(expected_rows.end(), p2_memory_rows.begin(), p2_memory_rows.end());
    ASSERT_OK_AND_ASSIGN(std::vector<Row> actual_rows, ReadRows(realtime_context));
    ASSERT_EQ(expected_rows, actual_rows);

    ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, ReadCommittedOffsets());
    ASSERT_EQ(2, committed_offsets.size());
    for (int64_t partition_index = 0; partition_index < 2; ++partition_index) {
        RealtimePartitionBucket partition_bucket({{"pt", "p" + std::to_string(partition_index)}},
                                                 /*bucket=*/0);
        ASSERT_EQ(9, committed_offsets.at(partition_bucket));
    }
    ASSERT_EQ(committed_offsets.end(),
              committed_offsets.find(RealtimePartitionBucket({{"pt", "p2"}}, /*bucket=*/0)));
    ASSERT_OK(writer->Close());
}

TEST_F(RealtimeWriteInteTest, TestMultipleBucketsRestoreIndependentOffsets) {
    options_[Options::BUCKET] = "2";
    CreateTable(/*partition_keys=*/{});

    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> first_writer, CreateRealtimeWriter());
    std::vector<Row> bucket0_disk_rows = MakeRows(/*first_id=*/0, /*count=*/2, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> bucket0_disk_batch,
                         MakeBatch(bucket0_disk_rows, /*partitioned=*/false, /*bucket=*/0));
    ASSERT_OK(first_writer->Write(std::move(bucket0_disk_batch)));
    std::vector<Row> bucket1_disk_rows = MakeRows(/*first_id=*/10, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> bucket1_disk_batch,
                         MakeBatch(bucket1_disk_rows, /*partitioned=*/false, /*bucket=*/1));
    ASSERT_OK(first_writer->Write(std::move(bucket1_disk_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> first_commits,
                         first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
    ASSERT_EQ(2, first_commits.size());
    ASSERT_OK(Commit(first_commits, /*commit_identifier=*/0));
    ASSERT_OK(first_writer->Close());

    ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context,
                         RealtimeContext::Create());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> second_writer,
                         CreateRealtimeWriter(realtime_context));
    std::vector<Row> bucket0_memory_rows =
        MakeRows(/*first_id=*/2, /*count=*/1, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> bucket0_memory_batch,
                         MakeBatch(bucket0_memory_rows, /*partitioned=*/false, /*bucket=*/0));
    ASSERT_OK(second_writer->Write(std::move(bucket0_memory_batch)));
    std::vector<Row> bucket1_memory_rows =
        MakeRows(/*first_id=*/13, /*count=*/1, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> bucket1_memory_batch,
                         MakeBatch(bucket1_memory_rows, /*partitioned=*/false, /*bucket=*/1));
    ASSERT_OK(second_writer->Write(std::move(bucket1_memory_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> second_commits,
                         second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1));
    ASSERT_EQ(2, second_commits.size());
    std::map<int32_t, Range> prepared_ranges;
    for (const RealtimeCommitProgress& commit : second_commits) {
        prepared_ranges.emplace(commit.partition_bucket.bucket, commit.offset_range);
    }
    ASSERT_EQ(Range(2, 2), prepared_ranges.at(0));
    ASSERT_EQ(Range(3, 3), prepared_ranges.at(1));

    ASSERT_OK_AND_ASSIGN(std::vector<Row> actual_rows, ReadRows(realtime_context));
    std::vector<Row> expected_rows = bucket0_disk_rows;
    expected_rows.insert(expected_rows.end(), bucket0_memory_rows.begin(),
                         bucket0_memory_rows.end());
    expected_rows.insert(expected_rows.end(), bucket1_disk_rows.begin(), bucket1_disk_rows.end());
    expected_rows.insert(expected_rows.end(), bucket1_memory_rows.begin(),
                         bucket1_memory_rows.end());
    std::sort(expected_rows.begin(), expected_rows.end());
    std::sort(actual_rows.begin(), actual_rows.end());
    ASSERT_EQ(expected_rows, actual_rows);

    ASSERT_OK(Commit(second_commits, /*commit_identifier=*/1));
    ASSERT_OK(second_writer->Close());
    ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, ReadCommittedOffsets());
    ASSERT_EQ(2, committed_offsets.size());
    ASSERT_EQ(2, committed_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0)));
    ASSERT_EQ(3, committed_offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/1)));
}

TEST_F(RealtimeWriteInteTest, TestRestoreOffsetFromCommittedSnapshot) {
    CreateTable(/*partition_keys=*/{});

    std::vector<Row> first_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> first_writer, CreateRealtimeWriter());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> first_batch,
                         MakeBatch(first_rows, /*partitioned=*/false));
    ASSERT_OK(first_writer->Write(std::move(first_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> first_commits,
                         first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0));
    ASSERT_EQ(1, first_commits.size());
    ASSERT_EQ(Range(0, 2), first_commits[0].offset_range);
    ASSERT_OK(Commit(first_commits, /*commit_identifier=*/0));
    ASSERT_OK(first_writer->Close());

    std::vector<Row> second_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0");
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> second_writer, CreateRealtimeWriter());
    ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> second_batch,
                         MakeBatch(second_rows, /*partitioned=*/false));
    ASSERT_OK(second_writer->Write(std::move(second_batch)));
    ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> second_commits,
                         second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1));
    ASSERT_EQ(1, second_commits.size());
    ASSERT_EQ(Range(3, 4), second_commits[0].offset_range);

    std::vector<Row> expected_rows = MakeRows(/*first_id=*/0, /*count=*/5, /*partition=*/"p0");
    FinalizeCommitAndCheck(second_writer.get(), std::move(second_commits),
                           /*prepare_identifier=*/1, std::move(expected_rows));

    RealtimePartitionBucket partition_bucket(/*partition=*/{}, /*bucket=*/0);
    ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap second_committed_offsets, ReadCommittedOffsets());
    ASSERT_EQ(4, second_committed_offsets.at(partition_bucket));
}

}  // namespace paimon::test