be/src/exec/sink/viceberg_delete_sink.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 <gen_cpp/DataSinks_types.h>
#include <gen_cpp/PlanNodes_types.h>

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

#include "common/status.h"
#include "core/block/block.h"
#include "exec/sink/writer/async_result_writer.h"
#include "exec/sink/writer/iceberg/viceberg_delete_file_writer.h"
#include "exprs/vexpr_fwd.h"
#include "roaring/roaring64map.hh"
#include "runtime/runtime_profile.h"

namespace doris {

class ObjectPool;
class RuntimeState;
class Dependency;

class VIcebergDeleteFileWriter;

struct IcebergFileDeletion {
    IcebergFileDeletion() = default;
    IcebergFileDeletion(int32_t spec_id, std::string data_json)
            : partition_spec_id(spec_id), partition_data_json(std::move(data_json)) {}

    int32_t partition_spec_id = 0;
    std::string partition_data_json;
    roaring::Roaring64Map rows_to_delete;
};

/**
 * Sink for writing Iceberg position delete files.
 *
 * This sink receives blocks containing a $row_id column with
 * (file_path, row_position, partition_spec_id, partition_data).
 * It groups delete records by file_path and writes position delete files.
 */
class VIcebergDeleteSink final : public AsyncResultWriter {
public:
    VIcebergDeleteSink(const TDataSink& t_sink, const VExprContextSPtrs& output_exprs,
                       std::shared_ptr<Dependency> dep, std::shared_ptr<Dependency> fin_dep);

    ~VIcebergDeleteSink() override = default;

    Status init_properties(ObjectPool* pool);

    Status open(RuntimeState* state, RuntimeProfile* profile) override;

    Status write(RuntimeState* state, Block& block) override;

    Status close(Status) override;

private:
    /**
     * Extract $row_id column from block and group by file_path.
     */
    Status _collect_position_deletes(const Block& block,
                                     std::map<std::string, IcebergFileDeletion>& file_deletions);

    /**
     * Write grouped position deletes to delete files
     */
    Status _write_position_delete_files(
            const std::map<std::string, IcebergFileDeletion>& file_deletions);

    Status _write_deletion_vector_files(
            const std::map<std::string, IcebergFileDeletion>& file_deletions);

    struct DeletionVectorBlob {
        std::string referenced_data_file;
        int32_t partition_spec_id = 0;
        std::string partition_data_json;
        int64_t delete_count = 0; // The number of rows deleted in this delete operation.
        int64_t merged_count =
                0; // The number of rows after merging the old deletion vector and position delete.
        int64_t content_offset = 0;
        int64_t content_size_in_bytes = 0;
        std::vector<char> blob_data;
    };

    Status _write_puffin_file(const std::string& puffin_path,
                              std::vector<DeletionVectorBlob>* blobs, int64_t* out_file_size);

    std::string _build_puffin_footer_json(const std::vector<DeletionVectorBlob>& blobs);

    std::string _generate_puffin_file_path();

    /**
     * Generate unique delete file path
     */
    std::string _generate_delete_file_path(const std::string& referenced_data_file = "");

    /**
     * Get $row_id column index from block
     */
    int _get_row_id_column_index(const Block& block);

    /**
     * Build a block for position delete (file_path, pos)
     */
    Status _build_position_delete_block(const std::string& file_path,
                                        const std::vector<int64_t>& positions, Block& output_block);
    Status _init_position_delete_output_exprs();
    std::string _get_file_extension() const;

    TDataSink _t_sink;
    RuntimeState* _state = nullptr;

    int32_t _format_version = 2;
    TFileContent::type _delete_type = TFileContent::POSITION_DELETES;

    // Writers for delete files
    std::vector<std::unique_ptr<VIcebergDeleteFileWriter>> _writers;

    // Collected commit data from all writers
    std::vector<TIcebergCommitData> _commit_data_list;
    // TODO: All deletions are held in memory until close(). Consider flushing
    //  per-file when the upstream guarantees file_path ordering, or flushing
    //  when estimated memory exceeds a threshold, to reduce peak memory usage.
    std::map<std::string, IcebergFileDeletion> _file_deletions;
    std::map<std::string, std::vector<TIcebergDeleteFileDesc>> _rewritable_delete_files;

    // Hadoop configuration
    std::map<std::string, std::string> _hadoop_conf;

    // File format settings
    TFileFormatType::type _file_format_type = TFileFormatType::FORMAT_PARQUET;
    TFileCompressType::type _compress_type = TFileCompressType::SNAPPYBLOCK;

    // Output directory for delete files
    std::string _output_path;
    std::string _table_location;

    TFileType::type _file_type = TFileType::FILE_HDFS;
    std::vector<TNetworkAddress> _broker_addresses;

    // Partition information
    int32_t _partition_spec_id = 0;
    std::string _partition_data_json;

    // Counters
    size_t _row_count = 0;
    size_t _delete_file_count = 0;

    // Profile counters
    RuntimeProfile::Counter* _written_rows_counter = nullptr;
    RuntimeProfile::Counter* _send_data_timer = nullptr;
    RuntimeProfile::Counter* _write_delete_files_timer = nullptr;
    RuntimeProfile::Counter* _delete_file_count_counter = nullptr;
    RuntimeProfile::Counter* _open_timer = nullptr;
    RuntimeProfile::Counter* _close_timer = nullptr;

    VExprContextSPtrs _position_delete_output_expr_ctxs;
};

} // namespace doris