src/paimon/core/append/append_only_writer.cpp

/*
 * 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 "paimon/core/append/append_only_writer.h"

#include <functional>
#include <string>
#include <utility>

#include "arrow/c/abi.h"
#include "arrow/c/bridge.h"
#include "arrow/c/helpers.h"
#include "arrow/type.h"
#include "paimon/common/metrics/metrics_impl.h"
#include "paimon/common/types/row_kind.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/common/utils/long_counter.h"
#include "paimon/common/utils/scope_guard.h"
#include "paimon/core/io/compact_increment.h"
#include "paimon/core/io/data_file_path_factory.h"
#include "paimon/core/io/data_file_writer.h"
#include "paimon/core/io/data_increment.h"
#include "paimon/core/io/multiple_blob_file_writer.h"
#include "paimon/core/io/rolling_blob_file_writer.h"
#include "paimon/core/io/rolling_file_writer.h"
#include "paimon/core/io/single_file_writer.h"
#include "paimon/core/manifest/file_source.h"
#include "paimon/core/utils/commit_increment.h"
#include "paimon/format/file_format.h"
#include "paimon/format/file_format_factory.h"
#include "paimon/format/writer_builder.h"
#include "paimon/macros.h"
#include "paimon/metrics.h"
#include "paimon/record_batch.h"

namespace paimon {

class MemoryPool;
class FormatStatsExtractor;

AppendOnlyWriter::AppendOnlyWriter(const CoreOptions& options, int64_t schema_id,
                                   const std::shared_ptr<arrow::Schema>& write_schema,
                                   const std::optional<std::vector<std::string>>& write_cols,
                                   int64_t max_sequence_number,
                                   const std::shared_ptr<DataFilePathFactory>& path_factory,
                                   const std::shared_ptr<CompactManager>& compact_manager,
                                   const std::shared_ptr<MemoryPool>& memory_pool)
    : options_(options),
      schema_id_(schema_id),
      write_schema_(write_schema),
      write_cols_(write_cols),
      seq_num_counter_(std::make_shared<LongCounter>(max_sequence_number + 1)),
      path_factory_(path_factory),
      compact_manager_(compact_manager),
      memory_pool_(memory_pool),
      metrics_(std::make_shared<MetricsImpl>()) {}

AppendOnlyWriter::~AppendOnlyWriter() = default;

Status AppendOnlyWriter::Write(std::unique_ptr<RecordBatch>&& batch) {
    for (const auto& row_kind : batch->GetRowKind()) {
        if (PAIMON_UNLIKELY(row_kind != RecordBatch::RowKind::INSERT)) {
            PAIMON_ASSIGN_OR_RAISE(const RowKind* kind,
                                   RowKind::FromByteValue(static_cast<int8_t>(row_kind)));
            return Status::Invalid("Append only writer can not accept record batch with RowKind ",
                                   kind->Name());
        }
    }
    if (writer_ == nullptr) {
        PAIMON_ASSIGN_OR_RAISE(writer_, CreateRollingRowWriter());
    }
    return writer_->Write(batch->GetData());
}

Result<CommitIncrement> AppendOnlyWriter::PrepareCommit(bool wait_compaction) {
    PAIMON_RETURN_NOT_OK(
        Flush(/*wait_for_latest_compaction=*/false, /*forced_full_compaction=*/false));
    PAIMON_RETURN_NOT_OK(TrySyncLatestCompaction(wait_compaction || options_.CommitForceCompact()));
    return DrainIncrement();
}

Result<CommitIncrement> AppendOnlyWriter::DrainIncrement() {
    DataIncrement data_increment(std::move(new_files_), std::move(deleted_files_), {});
    CompactIncrement compact_increment(std::move(compact_before_), std::move(compact_after_), {});
    auto drain_deletion_file = compact_deletion_file_;

    new_files_.clear();
    deleted_files_.clear();
    compact_before_.clear();
    compact_after_.clear();
    compact_deletion_file_ = nullptr;

    return CommitIncrement(data_increment, compact_increment, drain_deletion_file);
}

Status AppendOnlyWriter::TrySyncLatestCompaction(bool blocking) {
    PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<CompactResult>> result,
                           compact_manager_->GetCompactionResult(blocking));
    if (result.has_value()) {
        const auto& compaction_result = result.value();
        const auto& before = compaction_result->Before();
        compact_before_.insert(compact_before_.end(), before.begin(), before.end());
        const auto& after = compaction_result->After();
        compact_after_.insert(compact_after_.end(), after.begin(), after.end());
        PAIMON_RETURN_NOT_OK(UpdateCompactDeletionFile(compaction_result->DeletionFile()));
    }
    return Status::OK();
}

Status AppendOnlyWriter::UpdateCompactDeletionFile(
    const std::shared_ptr<CompactDeletionFile>& new_deletion_file) {
    if (new_deletion_file) {
        if (compact_deletion_file_ == nullptr) {
            compact_deletion_file_ = new_deletion_file;
        } else {
            PAIMON_ASSIGN_OR_RAISE(compact_deletion_file_,
                                   new_deletion_file->MergeOldFile(compact_deletion_file_));
        }
    }
    return Status::OK();
}

Status AppendOnlyWriter::Flush(bool wait_for_latest_compaction, bool forced_full_compaction) {
    std::vector<std::shared_ptr<DataFileMeta>> flushed_files;
    if (writer_) {
        PAIMON_RETURN_NOT_OK(writer_->Close());
        PAIMON_ASSIGN_OR_RAISE(flushed_files, writer_->GetResult());
    }
    // add new generated files
    for (const auto& flushed_file : flushed_files) {
        PAIMON_RETURN_NOT_OK(compact_manager_->AddNewFile(flushed_file));
    }
    PAIMON_RETURN_NOT_OK(TrySyncLatestCompaction(wait_for_latest_compaction));
    PAIMON_RETURN_NOT_OK(compact_manager_->TriggerCompaction(forced_full_compaction));
    new_files_.insert(new_files_.end(), flushed_files.begin(), flushed_files.end());
    if (writer_) {
        metrics_->Merge(writer_->GetMetrics());
        writer_.reset();
    }
    return Status::OK();
}

AppendOnlyWriter::RollingFileWriterResult AppendOnlyWriter::CreateRollingRowWriter() const {
    auto schemas = BlobUtils::SeparateBlobSchema(write_schema_);
    if (schemas.blob_schema && schemas.blob_schema->num_fields() > 0) {
        return CreateRollingBlobWriter(schemas);
    } else {
        return std::make_unique<RollingFileWriter<::ArrowArray*, std::shared_ptr<DataFileMeta>>>(
            options_.GetTargetFileSize(/*has_primary_key=*/false),
            GetDataFileWriterCreator(write_schema_, write_cols_));
    }
}

AppendOnlyWriter::SingleFileWriterCreator AppendOnlyWriter::GetDataFileWriterCreator(
    const std::shared_ptr<arrow::Schema>& schema,
    const std::optional<std::vector<std::string>>& write_cols) const {
    return
        [this, schema, write_cols]()
            -> Result<
                std::unique_ptr<SingleFileWriter<::ArrowArray*, std::shared_ptr<DataFileMeta>>>> {
            ::ArrowSchema arrow_schema;
            ScopeGuard guard([&arrow_schema]() { ArrowSchemaRelease(&arrow_schema); });
            PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &arrow_schema));
            auto format = options_.GetFileFormat();
            PAIMON_ASSIGN_OR_RAISE(
                std::shared_ptr<WriterBuilder> writer_builder,
                format->CreateWriterBuilder(&arrow_schema, options_.GetWriteBatchSize()));
            writer_builder->WithMemoryPool(memory_pool_);

            PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &arrow_schema));
            PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<FormatStatsExtractor> stats_extractor,
                                   format->CreateStatsExtractor(&arrow_schema));
            auto writer = std::make_unique<DataFileWriter>(
                options_.GetFileCompression(), std::function<Status(ArrowArray*, ArrowArray*)>(),
                schema_id_, seq_num_counter_, FileSource::Append(), stats_extractor,
                path_factory_->IsExternalPath(), write_cols, memory_pool_);
            PAIMON_RETURN_NOT_OK(
                writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), writer_builder));
            return writer;
        };
}

AppendOnlyWriter::SingleFileWriterCreator AppendOnlyWriter::GetBlobFileWriterCreator(
    const std::shared_ptr<WriterBuilder>& writer_builder,
    const std::shared_ptr<FormatStatsExtractor>& stats_extractor,
    const std::optional<std::vector<std::string>>& write_cols) const {
    return
        [this, writer_builder, stats_extractor, write_cols]()
            -> Result<
                std::unique_ptr<SingleFileWriter<::ArrowArray*, std::shared_ptr<DataFileMeta>>>> {
            auto writer = std::make_unique<DataFileWriter>(
                /*compression=*/"none", std::function<Status(ArrowArray*, ArrowArray*)>(),
                schema_id_, seq_num_counter_, FileSource::Append(), stats_extractor,
                path_factory_->IsExternalPath(), write_cols, memory_pool_);
            PAIMON_RETURN_NOT_OK(writer->Init(options_.GetFileSystem(),
                                              path_factory_->NewBlobPath(), writer_builder));
            return writer;
        };
}

AppendOnlyWriter::RollingFileWriterResult AppendOnlyWriter::CreateRollingBlobWriter(
    const BlobUtils::SeparatedSchemas& schemas) const {
    // Multiple blob fields are supported. Each blob field gets its own rolling file writer
    // via MultipleBlobFileWriter.
    auto blob_schema = schemas.blob_schema;
    auto blob_writer_creator = [this, blob_schema](const std::string& blob_field_name)
        -> Result<
            std::unique_ptr<RollingFileWriter<::ArrowArray*, std::shared_ptr<DataFileMeta>>>> {
        // Create a single-field schema for this blob field
        auto field = blob_schema->GetFieldByName(blob_field_name);
        if (!field) {
            return Status::Invalid(
                fmt::format("Blob field '{}' not found in blob schema", blob_field_name));
        }
        auto single_field_schema = arrow::schema({field});
        ::ArrowSchema arrow_schema;
        ScopeGuard guard([&arrow_schema]() { ArrowSchemaRelease(&arrow_schema); });
        PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*single_field_schema, &arrow_schema));
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FileFormat> format,
                               FileFormatFactory::Get("blob", options_.ToMap()));
        PAIMON_ASSIGN_OR_RAISE(
            std::shared_ptr<WriterBuilder> writer_builder,
            format->CreateWriterBuilder(&arrow_schema, options_.GetWriteBatchSize()));
        writer_builder->WithMemoryPool(memory_pool_);
        PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*single_field_schema, &arrow_schema));
        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<FormatStatsExtractor> stats_extractor,
                               format->CreateStatsExtractor(&arrow_schema));

        std::vector<std::string> write_cols = {blob_field_name};
        auto single_blob_file_writer_creator =
            GetBlobFileWriterCreator(writer_builder, stats_extractor, write_cols);
        return std::make_unique<RollingFileWriter<::ArrowArray*, std::shared_ptr<DataFileMeta>>>(
            options_.GetBlobTargetFileSize(), single_blob_file_writer_creator);
    };

    return std::make_unique<RollingBlobFileWriter>(
        options_.GetTargetFileSize(/*has_primary_key=*/false),
        GetDataFileWriterCreator(schemas.main_schema, schemas.main_schema->field_names()),
        blob_schema, blob_writer_creator, arrow::struct_(write_schema_->fields()));
}

Status AppendOnlyWriter::Sync() {
    return TrySyncLatestCompaction(/*blocking=*/true);
}

Status AppendOnlyWriter::Close() {
    // Request cancellation and wait for running compaction to exit.
    // This avoids reusing cancellation state while an old task is still running.
    compact_manager_->CancelAndWaitCompaction();
    PAIMON_RETURN_NOT_OK(Sync());

    PAIMON_RETURN_NOT_OK(compact_manager_->Close());
    auto fs = options_.GetFileSystem();
    for (const auto& file : compact_after_) {
        // AppendOnlyCompactManager will rewrite the file and no file upgrade will occur, so we
        // can directly delete the file in compact_after_.
        [[maybe_unused]] auto s = fs->Delete(path_factory_->ToPath(file));
    }

    if (writer_) {
        writer_->Abort();
        writer_.reset();
    }

    if (compact_deletion_file_ != nullptr) {
        compact_deletion_file_->Clean();
    }
    return Status::OK();
}

}  // namespace paimon