src/paimon/testing/utils/test_helper.h

/*
 * 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.
 */

#pragma once

#include <map>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "arrow/api.h"
#include "arrow/array/concatenate.h"
#include "arrow/c/bridge.h"
#include "arrow/ipc/api.h"
#include "fmt/format.h"
#include "paimon/api.h"
#include "paimon/catalog/catalog.h"
#include "paimon/commit_context.h"
#include "paimon/common/data/blob_utils.h"
#include "paimon/common/types/data_field.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/common/utils/checked_cast.h"
#include "paimon/common/utils/path_util.h"
#include "paimon/core/operation/append_only_file_store_write.h"
#include "paimon/core/operation/file_store_commit_impl.h"
#include "paimon/core/schema/schema_manager.h"
#include "paimon/core/schema/table_schema.h"
#include "paimon/core/table/sink/commit_message_impl.h"
#include "paimon/core/utils/file_store_path_factory.h"
#include "paimon/core/utils/snapshot_manager.h"
#include "paimon/data/blob.h"
#include "paimon/data/blob_descriptor.h"
#include "paimon/file_store_commit.h"
#include "paimon/file_store_write.h"
#include "paimon/fs/file_system_factory.h"
#include "paimon/record_batch.h"
#include "paimon/result.h"
#include "paimon/table/source/startup_mode.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 TestHelper {
 public:
    static Result<std::unique_ptr<TestHelper>> Create(
        const std::string& root_path, const std::shared_ptr<arrow::Schema>& schema,
        const std::vector<std::string>& partition_keys,
        const std::vector<std::string>& primary_keys,
        const std::map<std::string, std::string>& options, bool is_streaming_mode,
        bool ignore_if_exists = false, const std::string& temp_directory = "") {
        PAIMON_ASSIGN_OR_RAISE(auto catalog, Catalog::Create(root_path, options));
        PAIMON_RETURN_NOT_OK(catalog->CreateDatabase("foo", options, ignore_if_exists));
        ::ArrowSchema c_schema;
        ScopeGuard guard([schema = &c_schema]() { ArrowSchemaRelease(schema); });
        PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &c_schema));
        PAIMON_RETURN_NOT_OK(catalog->CreateTable(Identifier("foo", "bar"), &c_schema,
                                                  partition_keys, primary_keys, options,
                                                  ignore_if_exists));
        std::string table_path = PathUtil::JoinPath(root_path, "foo.db/bar");
        return Create(table_path, options, is_streaming_mode, temp_directory);
    }

    static Result<std::unique_ptr<TestHelper>> Create(
        const std::string& table_path, const std::map<std::string, std::string>& options,
        bool is_streaming_mode, const std::string& temp_directory = "") {
        std::string file_system_identifier = "local";
        auto fs_iter = options.find(Options::FILE_SYSTEM);
        if (fs_iter != options.end()) {
            file_system_identifier = StringUtils::ToLowerCase(fs_iter->second);
        }
        PAIMON_ASSIGN_OR_RAISE(auto file_system,
                               FileSystemFactory::Get(file_system_identifier, table_path, options));
        std::string commit_user = "commit_user";
        WriteContextBuilder context_builder(table_path, commit_user);
        if (!temp_directory.empty()) {
            context_builder.WithTempDirectory(temp_directory);
        }
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<WriteContext> write_context,
                               context_builder.SetOptions(options)
                                   .WithStreamingMode(is_streaming_mode)
                                   .WithWriteId(12345)
                                   .Finish());
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FileStoreWrite> write,
                               FileStoreWrite::Create(std::move(write_context)));
        CommitContextBuilder commit_context_builder(table_path, commit_user);
        PAIMON_ASSIGN_OR_RAISE(
            std::unique_ptr<CommitContext> commit_context,
            commit_context_builder.SetOptions(options).IgnoreEmptyCommit(false).Finish());
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FileStoreCommit> commit,
                               FileStoreCommit::Create(std::move(commit_context)));
        return std::unique_ptr<TestHelper>(new TestHelper(std::move(file_system), std::move(write),
                                                          std::move(commit), commit_user,
                                                          table_path, options));
    }

    static Result<std::unique_ptr<RecordBatch>> MakeRecordBatch(
        const std::shared_ptr<arrow::DataType>& data_type, const std::string& data_str,
        const std::map<std::string, std::string>& partition_map, int32_t bucket,
        const std::vector<RecordBatch::RowKind>& row_kinds) {
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
            auto array, arrow::ipc::internal::json::ArrayFromJSON(data_type, data_str));
        ::ArrowArray arrow_array;
        PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &arrow_array));
        RecordBatchBuilder batch_builder(&arrow_array);
        return batch_builder.SetPartition(partition_map)
            .SetBucket(bucket)
            .SetRowKinds(row_kinds)
            .Finish();
    }

    Result<std::vector<std::shared_ptr<CommitMessage>>> WriteAndCommit(
        std::unique_ptr<RecordBatch>&& record_batch, int64_t commit_identifier,
        const std::optional<std::vector<std::shared_ptr<CommitMessage>>>&
            expected_commit_messages) {
        std::vector<std::unique_ptr<RecordBatch>> batches;
        batches.emplace_back(std::move(record_batch));
        return WriteAndCommit(std::move(batches), commit_identifier, expected_commit_messages);
    }

    Result<std::vector<std::shared_ptr<CommitMessage>>> WriteAndCommit(
        std::vector<std::unique_ptr<RecordBatch>>&& record_batches, int64_t commit_identifier,
        const std::optional<std::vector<std::shared_ptr<CommitMessage>>>&
            expected_commit_messages) {
        for (auto& record_batch : record_batches) {
            PAIMON_RETURN_NOT_OK(write_->Write(std::move(record_batch)));
        }
        PAIMON_ASSIGN_OR_RAISE(std::vector<std::shared_ptr<CommitMessage>> commit_messages,
                               write_->PrepareCommit(/*wait_compaction=*/false, commit_identifier));
        if (expected_commit_messages) {
            CheckCommitMessages(expected_commit_messages.value(), commit_messages);
            CheckExternalPath(commit_messages);
        }
        PAIMON_RETURN_NOT_OK(commit_->Commit(commit_messages, commit_identifier));
        return commit_messages;
    }

    Result<std::vector<std::shared_ptr<Split>>> NewScan(StartupMode startup_mode,
                                                        std::optional<int64_t> snapshot_id,
                                                        bool is_streaming = true) {
        ScanContextBuilder scan_context_builder(table_path_);
        scan_context_builder.WithStreamingMode(is_streaming)
            .SetOptions(options_)
            .AddOption(Options::SCAN_MODE, startup_mode.ToString());
        if (snapshot_id) {
            scan_context_builder.AddOption(Options::SCAN_SNAPSHOT_ID,
                                           std::to_string(snapshot_id.value()));
        }
        PAIMON_ASSIGN_OR_RAISE(auto scan_context, scan_context_builder.Finish());
        PAIMON_ASSIGN_OR_RAISE(scan_, TableScan::Create(std::move(scan_context)));
        return Scan();
    }

    Result<std::vector<std::shared_ptr<Split>>> Scan() {
        if (scan_ == nullptr) {
            return Status::Invalid("need call NewScan first");
        }
        PAIMON_ASSIGN_OR_RAISE(auto result_plan, scan_->CreatePlan());
        return result_plan->Splits();
    }

    /// Builds a one-row struct array holding the serialized descriptor of the blob.
    static Result<std::shared_ptr<arrow::Array>> MakeBlobDescriptorArray(
        const std::shared_ptr<arrow::DataType>& struct_type, const std::shared_ptr<Blob>& blob,
        const std::shared_ptr<MemoryPool>& pool) {
        if (struct_type->num_fields() != 1 ||
            struct_type->field(0)->type()->id() != arrow::Type::LARGE_BINARY) {
            return Status::Invalid("struct_type must have a single large binary field");
        }
        arrow::StructBuilder struct_builder(struct_type, arrow::default_memory_pool(),
                                            {std::make_shared<arrow::LargeBinaryBuilder>()});
        auto blob_builder =
            checked_cast<arrow::LargeBinaryBuilder*>(struct_builder.field_builder(0));
        PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder.Append());
        PAIMON_UNIQUE_PTR<Bytes> descriptor = blob->ToDescriptor(pool);
        PAIMON_RETURN_NOT_OK_FROM_ARROW(
            blob_builder->Append(descriptor->data(), descriptor->size()));
        std::shared_ptr<arrow::Array> array;
        PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder.Finish(&array));
        return array;
    }

    static Result<bool> CheckBlobsEqual(const std::vector<std::shared_ptr<Blob>>& result_blobs,
                                        const std::vector<std::shared_ptr<Blob>>& expected_blobs,
                                        const std::shared_ptr<FileSystem>& fs) {
        if (result_blobs.size() != expected_blobs.size()) {
            std::cout << "[result_blobs.size]: " << result_blobs.size() << std::endl;
            std::cout << "[expected_blobs.size]: " << expected_blobs.size() << std::endl;
            return false;
        }
        for (uint32_t i = 0; i < result_blobs.size(); ++i) {
            PAIMON_ASSIGN_OR_RAISE(auto result_stream, result_blobs[i]->NewInputStream(fs));
            PAIMON_ASSIGN_OR_RAISE(auto expected_stream, expected_blobs[i]->NewInputStream(fs));

            PAIMON_ASSIGN_OR_RAISE(int64_t result_length, result_stream->Length());
            PAIMON_ASSIGN_OR_RAISE(int64_t expected_length, expected_stream->Length());
            if (result_length != expected_length) {
                auto result_descriptor_bytes = result_blobs[i]->ToDescriptor(GetDefaultPool());
                auto expected_descriptor_bytes = expected_blobs[i]->ToDescriptor(GetDefaultPool());
                PAIMON_ASSIGN_OR_RAISE(
                    auto result_descriptor,
                    BlobDescriptor::Deserialize(result_descriptor_bytes->data(),
                                                result_descriptor_bytes->size()));
                PAIMON_ASSIGN_OR_RAISE(
                    auto expected_descriptor,
                    BlobDescriptor::Deserialize(expected_descriptor_bytes->data(),
                                                expected_descriptor_bytes->size()));
                std::cout << "blobs[" << i << "]: " << std::endl;
                std::cout << "[result_length(" << result_length << ") != expected_length("
                          << expected_length << ")]" << std::endl;
                std::cout << "[result_descriptor]: " << result_descriptor->ToString() << std::endl;
                std::cout << "[expected_descriptor]: " << expected_descriptor->ToString()
                          << std::endl;
                return false;
            }

            std::vector<char> result_bytes(result_length, 0);
            std::vector<char> expected_bytes(expected_length, 0);
            PAIMON_RETURN_NOT_OK(result_stream->Read(result_bytes.data(), result_length));
            PAIMON_RETURN_NOT_OK(expected_stream->Read(expected_bytes.data(), expected_length));
            if (result_bytes != expected_bytes) {
                std::cout << "blobs[" << i << "]: " << std::endl;
                std::cout << "[result_bytes != expected_bytes]" << std::endl;
                return false;
            }
        }
        return true;
    }

    static Result<std::vector<std::shared_ptr<Blob>>> ToBlobs(
        const std::shared_ptr<arrow::StructArray>& blob_struct_array) {
        std::vector<std::shared_ptr<Blob>> result_blobs;
        auto child_array = blob_struct_array->field(0);
        assert(blob_struct_array->num_fields() == 1);
        assert(blob_struct_array->null_count() == 0);
        assert(child_array->null_count() == 0);
        assert(child_array->type_id() == arrow::Type::type::LARGE_BINARY);

        const auto& blob_array = checked_cast<const arrow::LargeBinaryArray&>(*child_array);
        for (int64_t i = 0; i < blob_array.length(); ++i) {
            std::string_view descriptor = blob_array.GetView(i);
            PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Blob> blob,
                                   Blob::FromDescriptor(descriptor.data(), descriptor.size()));
            result_blobs.emplace_back(blob);
        }
        return result_blobs;
    }

    /// Reads all rows of the given splits and returns the raw result (including the leading
    /// `_VALUE_KIND` column). Useful when the expected data cannot be expressed as JSON, e.g.
    /// binary-encoded VARIANT columns.
    Result<std::shared_ptr<arrow::ChunkedArray>> ReadResult(
        const std::vector<std::shared_ptr<Split>>& splits) {
        return ReadResult(splits, /*read_schema=*/nullptr);
    }

    /// Reads all rows of the given splits with an optional projected read schema.
    Result<std::shared_ptr<arrow::ChunkedArray>> ReadResult(
        const std::vector<std::shared_ptr<Split>>& splits,
        std::unique_ptr<::ArrowSchema> read_schema) {
        ReadContextBuilder read_context_builder(table_path_);
        read_context_builder.SetOptions(options_);
        if (read_schema != nullptr) {
            read_context_builder.SetReadSchema(std::move(read_schema));
        }
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> read_context,
                               read_context_builder.Finish());
        PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context)));
        PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(splits));
        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::ChunkedArray> collected,
                               ReadResultCollector::CollectResult(std::move(batch_reader)));
        if (collected->num_chunks() == 0) {
            return collected;
        }
        // Preserve this helper's single-chunk result contract. CollectResult destroys the reader
        // before importing the batches, so concatenating here also exercises the returned data
        // after the reader has been released.
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
            std::shared_ptr<arrow::Array> concatenated,
            arrow::Concatenate(collected->chunks(), arrow::default_memory_pool()));
        return std::make_shared<arrow::ChunkedArray>(std::move(concatenated));
    }

    Result<bool> ReadAndCheckResult(const std::shared_ptr<arrow::DataType>& data_type,
                                    const std::vector<std::shared_ptr<Split>>& splits,
                                    const std::string& expected_result) {
        ReadContextBuilder read_context_builder(table_path_);
        read_context_builder.SetOptions(options_);
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> read_context,
                               read_context_builder.Finish());
        PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context)));
        PAIMON_ASSIGN_OR_RAISE(auto batch_reader, table_read->CreateReader(splits));
        PAIMON_ASSIGN_OR_RAISE(auto read_result,
                               ReadResultCollector::CollectResult(std::move(batch_reader)));
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
            auto expected_array,
            arrow::ipc::internal::json::ArrayFromJSON(data_type, expected_result));
        auto expected_chunk_array = std::make_shared<arrow::ChunkedArray>(expected_array);

        arrow::EqualOptions equal_options = arrow::EqualOptions::Defaults();
        bool is_equal =
            expected_chunk_array->Equals(read_result, equal_options.diff_sink(&std::cout));

        if (!is_equal) {
            std::cout << "[expected_data_type]" << expected_chunk_array->type()->ToString()
                      << std::endl;
            std::cout << "[actual_data_type]" << read_result->type()->ToString() << std::endl;
            std::cout << "[expected]:" << expected_chunk_array->ToString() << std::endl;
            std::cout << "[actual]: " << read_result->ToString() << std::endl;
        }
        table_read.reset();
        return is_equal;
    }

    Result<std::optional<Snapshot>> LatestSnapshot() const {
        auto commit_impl = dynamic_cast<FileStoreCommitImpl*>(commit_.get());
        return commit_impl->snapshot_manager_->LatestSnapshotOfUser(commit_user_);
    }

    Result<std::optional<std::shared_ptr<TableSchema>>> LatestSchema() const {
        auto commit_impl = dynamic_cast<FileStoreCommitImpl*>(commit_.get());
        return commit_impl->schema_manager_->Latest();
    }

    Result<std::string> PartitionStr(const BinaryRow& partition) const {
        auto abstract_write = dynamic_cast<AbstractFileStoreWrite*>(write_.get());
        return abstract_write->file_store_path_factory_->GetPartitionString(partition);
    }

    /// Appends a new schema holding `fields` on top of the table's latest one, the way an ALTER
    /// TABLE would, so later writes and reads see the evolved schema.
    static Status WriteNextSchema(const std::shared_ptr<FileSystem>& file_system,
                                  const std::string& table_path,
                                  const std::vector<DataField>& fields, int32_t highest_field_id,
                                  const std::map<std::string, std::string>& options) {
        SchemaManager schema_manager(file_system, table_path);
        PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<TableSchema>> latest_schema_opt,
                               schema_manager.Latest());
        if (!latest_schema_opt) {
            return Status::Invalid("table schema does not exist");
        }
        auto next_schema = std::make_shared<TableSchema>(*latest_schema_opt.value());
        next_schema->id_ = latest_schema_opt.value()->Id() + 1;
        next_schema->fields_ = fields;
        next_schema->highest_field_id_ = highest_field_id;
        next_schema->options_ = options;
        PAIMON_ASSIGN_OR_RAISE(std::string schema_content, next_schema->ToJsonString());
        std::string schema_path = PathUtil::JoinPath(schema_manager.SchemaDirectory(),
                                                     fmt::format("schema-{}", next_schema->Id()));
        return file_system->AtomicStore(schema_path, schema_content);
    }

    static void CheckCommitMessages(std::vector<std::shared_ptr<CommitMessage>> expected,
                                    std::vector<std::shared_ptr<CommitMessage>> actual) {
        ASSERT_EQ(expected.size(), actual.size());
        auto commit_message_sort_function = [](const std::shared_ptr<CommitMessage>& lhs,
                                               const std::shared_ptr<CommitMessage>& rhs) -> bool {
            auto lhs_impl = std::dynamic_pointer_cast<CommitMessageImpl>(lhs);
            auto rhs_impl = std::dynamic_pointer_cast<CommitMessageImpl>(rhs);
            if (lhs_impl->Partition() == rhs_impl->Partition()) {
                return lhs_impl->Bucket() < rhs_impl->Bucket();
            }
            return lhs_impl->Partition().HashCode() < rhs_impl->Partition().HashCode();
        };
        std::stable_sort(expected.begin(), expected.end(), commit_message_sort_function);
        std::stable_sort(actual.begin(), actual.end(), commit_message_sort_function);
        for (size_t i = 0; i < actual.size(); i++) {
            auto result_impl = std::dynamic_pointer_cast<CommitMessageImpl>(actual[i]);
            auto expect_impl = std::dynamic_pointer_cast<CommitMessageImpl>(expected[i]);
            bool is_equal = result_impl->TEST_Equal(*expect_impl);
            if (!is_equal) {
                std::cout << "actual_msg:" << result_impl->ToString() << std::endl;
                std::cout << "expect_msg:" << expect_impl->ToString() << std::endl;
            }
            ASSERT_TRUE(is_equal);
        }
    }

    static int64_t CountChannelFiles(const std::shared_ptr<FileSystem>& file_system,
                                     const std::string& tmp_path) {
        std::vector<BasicFileStatus> dir_statuses;
        EXPECT_OK(file_system->ListDir(tmp_path, &dir_statuses));

        int64_t channel_file_count = 0;
        for (const auto& dir_status : dir_statuses) {
            const std::string dir_path = dir_status.GetPath();
            if (dir_status.IsDir() && dir_path.find("paimon-io-") != std::string::npos) {
                std::vector<BasicFileStatus> file_statuses;
                EXPECT_OK(file_system->ListDir(dir_path, &file_statuses));

                for (const auto& file_status : file_statuses) {
                    if (StringUtils::EndsWith(file_status.GetPath(), ".channel")) {
                        ++channel_file_count;
                    }
                }
            }
        }
        return channel_file_count;
    }

 private:
    void CheckExternalPath(const std::vector<std::shared_ptr<CommitMessage>>& actuals) {
        for (const auto& actual : actuals) {
            auto msg = std::dynamic_pointer_cast<CommitMessageImpl>(actual);
            auto files = msg->GetNewFilesIncrement().NewFiles();
            for (auto& file : files) {
                if (file->external_path) {
                    ASSERT_OK_AND_ASSIGN(bool file_exist, fs_->Exists(file->external_path.value()));
                    ASSERT_TRUE(file_exist);
                }
            }
        }
    }

    TestHelper(std::unique_ptr<FileSystem> fs, std::unique_ptr<FileStoreWrite> write,
               std::unique_ptr<FileStoreCommit> commit, const std::string& commit_user,
               const std::string& table_path, const std::map<std::string, std::string>& options)
        : fs_(std::move(fs)),
          write_(std::move(write)),
          commit_(std::move(commit)),
          commit_user_(commit_user),
          table_path_(table_path),
          options_(options) {}

    const std::string& CommitUser() const {
        return commit_user_;
    }

    std::shared_ptr<FileSystem> fs_;
    std::unique_ptr<FileStoreWrite> write_;
    std::unique_ptr<FileStoreCommit> commit_;
    std::unique_ptr<TableScan> scan_;
    std::string commit_user_;
    std::string table_path_;
    std::map<std::string, std::string> options_;
};

}  // namespace paimon::test