be/src/format_v2/expr/delete_predicate.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 "format_v2/expr/delete_predicate.h"

#include <fmt/format.h>
#include <gen_cpp/Exprs_types.h>
#include <glog/logging.h>

#include <algorithm>
#include <cstddef>
#include <ostream>

#include "common/cast_set.h"
#include "common/status.h"
#include "core/block/block.h"
#include "core/block/column_numbers.h"
#include "core/block/column_with_type_and_name.h"
#include "core/block/columns_with_type_and_name.h"

namespace doris::format {

DeletePredicate::DeletePredicate(const std::vector<int64_t>& deleted_rows)
        : VExpr(), _deleted_rows(&deleted_rows) {
    _node_type = TExprNodeType::PREDICATE;
    _opcode = TExprOpcode::DELETE;
    _data_type = std::make_shared<DataTypeBool>();
}

DeletePredicate::DeletePredicate(const roaring::Roaring64Map& deletion_vector)
        : VExpr(), _deletion_vector(&deletion_vector) {
    _node_type = TExprNodeType::PREDICATE;
    _opcode = TExprOpcode::DELETE;
    _data_type = std::make_shared<DataTypeBool>();
}

Status DeletePredicate::prepare(RuntimeState* state, const RowDescriptor& desc,
                                VExprContext* context) {
    RETURN_IF_ERROR_OR_PREPARED(VExpr::prepare(state, desc, context));
    _expr_name = "DeletePredicate";
    _prepare_finished = true;
    return Status::OK();
}

Status DeletePredicate::open(RuntimeState* state, VExprContext* context,
                             FunctionContext::FunctionStateScope scope) {
    DCHECK(_prepare_finished);
    RETURN_IF_ERROR_OR_PREPARED(VExpr::open(state, context, scope));
    _open_finished = true;
    return Status::OK();
}

void DeletePredicate::close(VExprContext* context, FunctionContext::FunctionStateScope scope) {
    VExpr::close(context, scope);
}

/**
 * DeletePredicate is derived from 2 cases:
 * 1. All row IDs indicates deleted rows. (e.g. Delete rows with row_id in (1, 2, 3))
 * 2. Bit vector indicates whether each row is deleted or not. (e.g. Bit vector[0,1,0,0,1] indicates row 1 and row 4 are deleted)
 *
 * So DeletePredicate should have exactly 1 child expr, which is the slot of row id.
 * Row IDs should be generated by file reader as a virtual column in `block`.
 **/
Status DeletePredicate::execute(VExprContext* context, Block* block, int* result_column_id) const {
    if (_children.size() != 1) {
        return Status::InternalError(fmt::format(
                "DeletePredicate should have exactly 1 child expr, but got {}", _children.size()));
    }
    int slot = -1;
    RETURN_IF_ERROR(_children[0]->execute(context, block, &slot));
    if (slot < 0 || static_cast<size_t>(slot) >= block->columns()) {
        return Status::InternalError(
                "DeletePredicate row id child returned invalid column id {}, block has {} columns",
                slot, block->columns());
    }
    const auto& row_ids =
            assert_cast<const ColumnInt64&>(*block->get_by_position(slot).column).get_data();
    const auto count = row_ids.size();
    auto res_col = ColumnBool::create(count, 0);
    if ((_deleted_rows == nullptr || _deleted_rows->empty()) &&
        (_deletion_vector == nullptr || _deletion_vector->isEmpty())) {
        block->insert({std::move(res_col), std::make_shared<DataTypeBool>(), expr_name()});
        *result_column_id = static_cast<int>(block->get_columns().size() - 1);
        return Status::OK();
    }
    if (count == 0) {
        block->insert({std::move(res_col), std::make_shared<DataTypeBool>(), expr_name()});
        *result_column_id = static_cast<int>(block->get_columns().size() - 1);
        return Status::OK();
    }
    if (_deletion_vector != nullptr) {
        auto it = _deletion_vector->begin();
        it.move(cast_set<uint64_t>(row_ids[0]));
        const auto end = _deletion_vector->end();
        const auto last_row_id = cast_set<uint64_t>(row_ids[count - 1]);
        while (it != end && *it <= last_row_id) {
            const auto row = cast_set<int64_t>(*it);
            if (const auto row_it = std::ranges::lower_bound(row_ids, row);
                row_it != row_ids.end() && *row_it == row) {
                res_col->get_data()[row_it - row_ids.begin()] = true;
            }
            ++it;
        }
        block->insert({std::move(res_col), std::make_shared<DataTypeBool>(), expr_name()});
        *result_column_id = static_cast<int>(block->get_columns().size() - 1);
        return Status::OK();
    }

    const int64_t* delete_rows = _deleted_rows->data();
    const int64_t* delete_rows_end = delete_rows + _deleted_rows->size();
    const int64_t* start_pos = std::lower_bound(delete_rows, delete_rows_end, row_ids[0]);
    int64_t start_index = start_pos - delete_rows;
    const int64_t* end_pos = std::upper_bound(start_pos, delete_rows_end, row_ids[count - 1]);
    const int64_t end_index = end_pos - delete_rows;

    while (start_index < end_index) {
        int64_t delete_row = delete_rows[start_index];
        if (const auto it = std::ranges::lower_bound(row_ids, delete_row);
            it != row_ids.end() && *it == delete_row) {
            const size_t index = it - row_ids.begin();
            res_col->get_data()[index] = true;
        }
        ++start_index;
    }
    block->insert({std::move(res_col), std::make_shared<DataTypeBool>(), expr_name()});
    *result_column_id = static_cast<int>(block->get_columns().size() - 1);
    return Status::OK();
}

std::string DeletePredicate::debug_string() const {
    return _expr_name;
}

} // namespace doris::format