be/src/exprs/vectorized_fn_call.cpp

Source
// 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 "exprs/vectorized_fn_call.h"

#include <fmt/compile.h>
#include <fmt/format.h>
#include <fmt/ranges.h> // IWYU pragma: keep
#include <gen_cpp/Opcodes_types.h>
#include <gen_cpp/Types_types.h>

#include <memory>
#include <optional>
#include <ostream>
#include <set>

#include "common/config.h"
#include "common/exception.h"
#include "common/logging.h"
#include "common/status.h"
#include "common/utils.h"
#include "core/assert_cast.h"
#include "core/block/block.h"
#include "core/block/column_numbers.h"
#include "core/column/column.h"
#include "core/column/column_array.h"
#include "core/column/column_nullable.h"
#include "core/column/column_vector.h"
#include "core/data_type/data_type.h"
#include "core/data_type/data_type_agg_state.h"
#include "core/types.h"
#include "exec/common/util.hpp"
#include "exec/pipeline/pipeline_task.h"
#include "exprs/function/array/function_array_distance.h"
#include "exprs/function/function_agg_state.h"
#include "exprs/function/function_fake.h"
#include "exprs/function/function_java_udf.h"
#include "exprs/function/function_python_udf.h"
#include "exprs/function/function_rpc.h"
#include "exprs/function/simple_function_factory.h"
#include "exprs/function_context.h"
#include "exprs/varray_literal.h"
#include "exprs/vcast_expr.h"
#include "exprs/vexpr_context.h"
#include "exprs/virtual_slot_ref.h"
#include "exprs/vliteral.h"
#include "runtime/runtime_state.h"
#include "storage/index/ann/ann_index.h"
#include "storage/index/ann/ann_index_iterator.h"
#include "storage/index/ann/ann_search_params.h"
#include "storage/index/index_reader.h"
#include "storage/index/zone_map/zonemap_eval_context.h"
#include "storage/segment/column_reader.h"
#include "storage/segment/virtual_column_iterator.h"
#include "util/simd/parquet_kernels.h"

namespace doris {
class RowDescriptor;
class RuntimeState;
class TExprNode;
} // namespace doris

namespace doris {

const std::string AGG_STATE_SUFFIX = "_state";

// Now left child is a function call, we need to check if it is a distance function
const static std::set<std::string> DISTANCE_FUNCS = {L2DistanceApproximate::name,
                                                     InnerProductApproximate::name};
const static std::set<TExprOpcode::type> OPS_FOR_ANN_RANGE_SEARCH = {
        TExprOpcode::GE, TExprOpcode::LE, TExprOpcode::LE, TExprOpcode::GT, TExprOpcode::LT};

namespace {

using simd::RawComparisonOp;

std::optional<RawComparisonOp> raw_comparison_op(std::string_view function_name, bool reverse) {
    RawComparisonOp op;
    if (function_name == "eq") {
        op = RawComparisonOp::EQ;
    } else if (function_name == "ne") {
        op = RawComparisonOp::NE;
    } else if (function_name == "lt") {
        op = RawComparisonOp::LT;
    } else if (function_name == "le") {
        op = RawComparisonOp::LE;
    } else if (function_name == "gt") {
        op = RawComparisonOp::GT;
    } else if (function_name == "ge") {
        op = RawComparisonOp::GE;
    } else {
        return std::nullopt;
    }
    if (!reverse || op == RawComparisonOp::EQ || op == RawComparisonOp::NE) {
        return op;
    }
    switch (op) {
    case RawComparisonOp::LT:
        return RawComparisonOp::GT;
    case RawComparisonOp::LE:
        return RawComparisonOp::GE;
    case RawComparisonOp::GT:
        return RawComparisonOp::LT;
    case RawComparisonOp::GE:
        return RawComparisonOp::LE;
    case RawComparisonOp::EQ:
    case RawComparisonOp::NE:
        break;
    }
    __builtin_unreachable();
}

template <typename T, PrimitiveType PT>
void execute_raw_comparison(const uint8_t* values, size_t num_values, const Field& literal,
                            RawComparisonOp op, uint8_t* matches) {
    const T rhs = literal.get<PT>();
    simd::raw_compare(values, num_values, rhs, op, matches);
}

} // namespace

VectorizedFnCall::VectorizedFnCall(const TExprNode& node) : VExpr(node) {
    _function_name = _fn.name.function_name;
}

Status VectorizedFnCall::prepare(RuntimeState* state, const RowDescriptor& desc,
                                 VExprContext* context) {
    RETURN_IF_ERROR_OR_PREPARED(VExpr::prepare(state, desc, context));
    ColumnsWithTypeAndName argument_template;
    argument_template.reserve(_children.size());
    for (auto child : _children) {
        if (child->is_literal()) {
            // For some functions, he needs some literal columns to derive the return type.
            auto literal_node = std::dynamic_pointer_cast<VLiteral>(child);
            argument_template.emplace_back(literal_node->get_column_ptr(), child->data_type(),
                                           child->expr_name());
        } else {
            argument_template.emplace_back(nullptr, child->data_type(), child->expr_name());
        }
    }

    _expr_name = fmt::format("VectorizedFnCall[{}](arguments={},return={})", _fn.name.function_name,
                             get_child_names(), _data_type->get_name());
    if (_fn.binary_type == TFunctionBinaryType::RPC) {
        _function = FunctionRPC::create(_fn, argument_template, _data_type);
    } else if (_fn.binary_type == TFunctionBinaryType::JAVA_UDF) {
        if (config::enable_java_support) {
            if (_fn.is_udtf_function) {
                // fake function. it's no use and can't execute.
                auto builder =
                        std::make_shared<DefaultFunctionBuilder>(FunctionFake<UDTFImpl>::create());
                _function = builder->build(argument_template, std::make_shared<DataTypeUInt8>());
            } else {
                _function = JavaFunctionCall::create(_fn, argument_template, _data_type);
            }
        } else {
            return Status::InternalError(
                    "Java UDF is not enabled, you can change be config enable_java_support to true "
                    "and restart be.");
        }
    } else if (_fn.binary_type == TFunctionBinaryType::PYTHON_UDF) {
        if (config::enable_python_udf_support) {
            if (_fn.is_udtf_function) {
                // fake function. it's no use and can't execute.
                // Python UDTF is executed via PythonUDTFFunction in table function path
                auto builder =
                        std::make_shared<DefaultFunctionBuilder>(FunctionFake<UDTFImpl>::create());
                _function = builder->build(argument_template, std::make_shared<DataTypeUInt8>());
            } else {
                _function = PythonFunctionCall::create(_fn, argument_template, _data_type);
                LOG(INFO) << fmt::format(
                        "create python function call: {}, runtime version: {}, function code: {}",
                        _fn.name.function_name, _fn.runtime_version, _fn.function_code);
            }
        } else {
            return Status::InternalError(
                    "Python UDF is not enabled, you can change be config enable_python_udf_support "
                    "to true and restart be.");
        }
    } else if (_fn.binary_type == TFunctionBinaryType::AGG_STATE) {
        DataTypes argument_types;
        for (auto column : argument_template) {
            argument_types.emplace_back(column.type);
        }

        if (match_suffix(_fn.name.function_name, AGG_STATE_SUFFIX)) {
            if (_data_type->is_nullable()) {
                return Status::InternalError("State function's return type must be not nullable");
            }
            if (_data_type->get_primitive_type() != PrimitiveType::TYPE_AGG_STATE) {
                return Status::InternalError(
                        "State function's return type must be agg_state but get {}",
                        _data_type->get_family_name());
            }
            _function = FunctionAggState::create(
                    argument_types, _data_type,
                    assert_cast<const DataTypeAggState*>(_data_type.get())->get_nested_function());
        } else {
            return Status::InternalError("Function {} is not endwith '_state'", _fn.signature);
        }
    } else {
        // get the function. won't prepare function.
        _function = SimpleFunctionFactory::instance().get_function(
                _fn.name.function_name, argument_template, _data_type,
                {.new_version_unix_timestamp = state->query_options().new_version_unix_timestamp,
                 .new_version_bitmap_op_count =
                         state->query_options().__isset.new_version_bitmap_op_count &&
                         state->query_options().new_version_bitmap_op_count},
                state->be_exec_version());
    }
    if (_function == nullptr) {
        return Status::InternalError("Could not find function {}, arg {} return {} ",
                                     _fn.name.function_name, get_child_type_names(),
                                     _data_type->get_name());
    }
    VExpr::register_function_context(state, context);
    _function_name = _fn.name.function_name;
    _prepare_finished = true;

    FunctionContext* fn_ctx = context->fn_context(_fn_context_index);
    if (fn().__isset.dict_function) {
        fn_ctx->set_dict_function(fn().dict_function);
    }
    return Status::OK();
}

Status VectorizedFnCall::open(RuntimeState* state, VExprContext* context,
                              FunctionContext::FunctionStateScope scope) {
    DCHECK(_prepare_finished);
    for (auto& i : _children) {
        RETURN_IF_ERROR(i->open(state, context, scope));
    }
    RETURN_IF_ERROR(VExpr::init_function_context(state, context, scope, _function));
    if (scope == FunctionContext::FRAGMENT_LOCAL) {
        RETURN_IF_ERROR(VExpr::get_const_col(context, nullptr));
    }
    _open_finished = true;
    return Status::OK();
}

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

Status VectorizedFnCall::evaluate_inverted_index(VExprContext* context, uint32_t segment_num_rows) {
    if (get_num_children() < 1) {
        // score() and similar 0-children virtual column functions don't need
        // inverted index evaluation; return OK to skip gracefully.
        return Status::OK();
    }
    return _evaluate_inverted_index(context, _function, segment_num_rows);
}

ZoneMapFilterResult VectorizedFnCall::evaluate_zonemap_filter(const ZoneMapEvalContext& ctx) const {
    return _function->evaluate_zonemap_filter(ctx, _children);
}

bool VectorizedFnCall::can_evaluate_zonemap_filter() const {
    return _function != nullptr && !_function->is_blockable() &&
           _function->can_evaluate_zonemap_filter(_children);
}

ZoneMapFilterResult VectorizedFnCall::evaluate_dictionary_filter(
        const DictionaryEvalContext& ctx) const {
    return _function->evaluate_dictionary_filter(ctx, _children);
}

bool VectorizedFnCall::can_evaluate_dictionary_filter() const {
    return _function != nullptr && !_function->is_blockable() &&
           _function->can_evaluate_dictionary_filter(_children);
}

ZoneMapFilterResult VectorizedFnCall::evaluate_bloom_filter(
        const BloomFilterEvalContext& ctx) const {
    return _function->evaluate_bloom_filter(ctx, _children);
}

bool VectorizedFnCall::can_evaluate_bloom_filter() const {
    return _function != nullptr && !_function->is_blockable() &&
           _function->can_evaluate_bloom_filter(_children);
}

Status VectorizedFnCall::_do_execute(VExprContext* context, const Block* block,
                                     const Selector* selector, size_t count,
                                     ColumnPtr& result_column, ColumnPtr* arg_column) const {
    if (is_const_and_have_executed()) { // const have executed in open function
        result_column = get_result_from_const(count);
        return Status::OK();
    }
    if (fast_execute(context, selector, count, result_column)) {
        return Status::OK();
    }
    DBUG_EXECUTE_IF("VectorizedFnCall.must_in_slow_path", {
        if (get_child(0)->is_slot_ref()) {
            auto debug_col_name = DebugPoints::instance()->get_debug_param_or_default<std::string>(
                    "VectorizedFnCall.must_in_slow_path", "column_name", "");

            std::vector<std::string> column_names;
            boost::split(column_names, debug_col_name, boost::algorithm::is_any_of(","));

            auto* column_slot_ref = assert_cast<VSlotRef*>(get_child(0).get());
            std::string column_name = column_slot_ref->expr_name();
            auto it = std::find(column_names.begin(), column_names.end(), column_name);
            if (it == column_names.end()) {
                return Status::Error<ErrorCode::INTERNAL_ERROR>(
                        "column {} should in slow path while VectorizedFnCall::execute.",
                        column_name);
            }
        }
    })
    DCHECK(_open_finished || block == nullptr) << debug_string();

    Block temp_block;
    ColumnNumbers args(_children.size());

    for (int i = 0; i < _children.size(); ++i) {
        ColumnPtr tmp_arg_column;
        RETURN_IF_ERROR(
                _children[i]->execute_column(context, block, selector, count, tmp_arg_column));
        auto arg_type = _children[i]->execute_type(block);
        temp_block.insert({tmp_arg_column, arg_type, _children[i]->expr_name()});
        args[i] = i;

        if (arg_column != nullptr && i == 0) {
            *arg_column = tmp_arg_column;
        }
    }

    uint32_t num_columns_without_result = temp_block.columns();
    // prepare a column to save result
    temp_block.insert({nullptr, _data_type, _expr_name});

    DBUG_EXECUTE_IF("VectorizedFnCall.wait_before_execute", {
        auto possibility = DebugPoints::instance()->get_debug_param_or_default<double>(
                "VectorizedFnCall.wait_before_execute", "possibility", 0);
        if (random_bool_slow(possibility)) {
            LOG(WARNING) << "VectorizedFnCall::execute sleep 30s";
            sleep(30);
        }
    });

    RETURN_IF_ERROR(_function->execute(context->fn_context(_fn_context_index), temp_block, args,
                                       num_columns_without_result, count));
    result_column = temp_block.get_by_position(num_columns_without_result).column;
    DCHECK_EQ(result_column->size(), count);
    RETURN_IF_ERROR(result_column->column_self_check());
    return Status::OK();
}

size_t VectorizedFnCall::estimate_memory(const size_t rows) {
    if (is_const_and_have_executed()) { // const have execute in open function
        return 0;
    }

    size_t estimate_size = 0;
    for (auto& child : _children) {
        estimate_size += child->estimate_memory(rows);
    }

    if (_data_type->have_maximum_size_of_value()) {
        estimate_size += rows * _data_type->get_size_of_value_in_memory();
    } else {
        estimate_size += rows * 512; /// FIXME: estimated value...
    }
    return estimate_size;
}

Status VectorizedFnCall::execute_runtime_filter(VExprContext* context, const Block* block,
                                                const uint8_t* __restrict filter, size_t count,
                                                ColumnPtr& result_column,
                                                ColumnPtr* arg_column) const {
    return _do_execute(context, block, nullptr, count, result_column, arg_column);
}

Status VectorizedFnCall::execute_column_impl(VExprContext* context, const Block* block,
                                             const Selector* selector, size_t count,
                                             ColumnPtr& result_column) const {
    return _do_execute(context, block, selector, count, result_column, nullptr);
}

bool VectorizedFnCall::can_execute_on_raw_fixed_values(const DataTypePtr& data_type,
                                                       int column_id) const {
    if (data_type == nullptr || !raw_comparison_op(_function_name, false).has_value()) {
        return false;
    }
    auto slot_literal = expr_zonemap::extract_slot_and_literal(_children);
    if (!slot_literal.has_value() || slot_literal->slot_index != column_id ||
        slot_literal->literal.is_null()) {
        return false;
    }
    const auto raw_type = remove_nullable(data_type);
    if (!remove_nullable(slot_literal->slot_type)->equals(*raw_type) ||
        !remove_nullable(slot_literal->literal_type)->equals(*raw_type)) {
        return false;
    }
    const auto primitive_type = raw_type->get_primitive_type();
    return primitive_type == TYPE_INT || primitive_type == TYPE_BIGINT ||
           primitive_type == TYPE_FLOAT || primitive_type == TYPE_DOUBLE;
}

Status VectorizedFnCall::execute_on_raw_fixed_values(const uint8_t* values, size_t num_values,
                                                     size_t value_width,
                                                     const DataTypePtr& data_type, int column_id,
                                                     uint8_t* matches) const {
    if (!can_execute_on_raw_fixed_values(data_type, column_id)) {
        return Status::NotSupported("Expression {} cannot evaluate raw fixed-width values",
                                    expr_name());
    }
    DORIS_CHECK(values != nullptr || num_values == 0);
    DORIS_CHECK(matches != nullptr || num_values == 0);
    const auto slot_literal = expr_zonemap::extract_slot_and_literal(_children);
    DORIS_CHECK(slot_literal.has_value());
    const auto op = raw_comparison_op(_function_name, slot_literal->literal_on_left);
    DORIS_CHECK(op.has_value());
    const auto primitive_type = remove_nullable(data_type)->get_primitive_type();
    const size_t expected_width = primitive_type == TYPE_INT || primitive_type == TYPE_FLOAT
                                          ? sizeof(uint32_t)
                                          : sizeof(uint64_t);
    if (value_width != expected_width) {
        return Status::Corruption("Raw expression width {} does not match expected {}", value_width,
                                  expected_width);
    }
    switch (primitive_type) {
    case TYPE_INT:
        execute_raw_comparison<int32_t, TYPE_INT>(values, num_values, slot_literal->literal, *op,
                                                  matches);
        break;
    case TYPE_BIGINT:
        execute_raw_comparison<int64_t, TYPE_BIGINT>(values, num_values, slot_literal->literal, *op,
                                                     matches);
        break;
    case TYPE_FLOAT:
        execute_raw_comparison<float, TYPE_FLOAT>(values, num_values, slot_literal->literal, *op,
                                                  matches);
        break;
    case TYPE_DOUBLE:
        execute_raw_comparison<double, TYPE_DOUBLE>(values, num_values, slot_literal->literal, *op,
                                                    matches);
        break;
    default:
        __builtin_unreachable();
    }
    return Status::OK();
}

const std::string& VectorizedFnCall::expr_name() const {
    return _expr_name;
}

std::string VectorizedFnCall::function_name() const {
    return _function_name;
}

std::string VectorizedFnCall::debug_string() const {
    std::stringstream out;
    out << "VectorizedFn[";
    out << _expr_name;
    out << "]{";
    bool first = true;
    for (const auto& input_expr : children()) {
        if (first) {
            first = false;
        } else {
            out << ",";
        }
        out << "\n" << input_expr->debug_string();
    }
    out << "}";
    return out.str();
}

std::string VectorizedFnCall::debug_string(const std::vector<VectorizedFnCall*>& agg_fns) {
    std::stringstream out;
    out << "[";
    for (int i = 0; i < agg_fns.size(); ++i) {
        out << (i == 0 ? "" : " ") << agg_fns[i]->debug_string();
    }
    out << "]";
    return out.str();
}

bool VectorizedFnCall::can_push_down_to_index() const {
    return _function->can_push_down_to_index();
}

bool VectorizedFnCall::is_deterministic() const {
    static const std::set<std::string> NON_DETERMINISTIC_FUNCTIONS = {
            "random", "rand", "random_bytes", "uuid", "uuid_numeric"};
    return !NON_DETERMINISTIC_FUNCTIONS.contains(_function_name) && VExpr::is_deterministic();
}

bool VectorizedFnCall::is_safe_to_execute_on_selected_rows() const {
    static const std::set<std::string> TOTAL_PREDICATE_FUNCTIONS = {
            "eq", "ne", "lt", "le", "gt", "ge", "in", "not_in", "is_null_pred", "is_not_null_pred"};
    // Selected-row execution may hide data-dependent errors in rows rejected by an earlier
    // predicate. Keep function calls unsafe by default and opt in only operations that are total
    // for their input domain; child checks then reject expressions such as gt(mod(x, -1), 0).
    return TOTAL_PREDICATE_FUNCTIONS.contains(_function_name) &&
           VExpr::is_safe_to_execute_on_selected_rows();
}

bool VectorizedFnCall::equals(const VExpr& other) {
    const auto* other_ptr = dynamic_cast<const VectorizedFnCall*>(&other);
    if (!other_ptr) {
        return false;
    }
    if (this->_function_name != other_ptr->_function_name) {
        return false;
    }
    if (get_num_children() != other_ptr->get_num_children()) {
        return false;
    }
    for (uint16_t i = 0; i < get_num_children(); i++) {
        if (!this->get_child(i)->equals(*other_ptr->get_child(i))) {
            return false;
        }
    }
    return true;
}

/*
 * For ANN range search we expect a comparison expression (LE/LT/GE/GT) whose left side is either:
 *   1) a vector distance function call, or
 *   2) a cast/virtual slot that unwraps to the function call when the planner promotes float to
 *      double literals.
 *
 * Visually the logical tree looks like:
 *
 *   FunctionCall(LE/LT/GE/GT)
 *   |----------------
 *   |               |
 *   |               |
 *   VirtualSlotRef* Float32Literal/Float64Literal
 *   |
 *   |
 *   Cast(Float -> Double)*
 *   |
 *   FunctionCall(distance)
 *   |----------------
 *   |               |
 *   |               |
 *   SlotRef         ArrayLiteral/Cast(String as Array<FLOAT>)
 *
 * Items marked with * are optional and depend on literal types/virtual column usage. The helper
 * below normalizes the shape and validates distance function, slot, and constant vector inputs.
 */

void VectorizedFnCall::prepare_ann_range_search(
        const doris::VectorSearchUserParams& user_params,
        segment_v2::AnnRangeSearchRuntime& range_search_runtime, bool& suitable_for_ann_index) {
    if (!suitable_for_ann_index) {
        return;
    }

    if (OPS_FOR_ANN_RANGE_SEARCH.find(this->op()) == OPS_FOR_ANN_RANGE_SEARCH.end()) {
        suitable_for_ann_index = false;
        return;
    }

    auto mark_unsuitable = [&](const std::string& reason) {
        suitable_for_ann_index = false;
        VLOG_DEBUG << "ANN range search skipped: " << reason;
    };

    range_search_runtime.is_le_or_lt =
            (this->op() == TExprOpcode::LE || this->op() == TExprOpcode::LT);

    DCHECK(_children.size() == 2);

    auto left_child = get_child(0);
    auto right_child = get_child(1);

    // ========== Step 1: Check left child - must be a distance function ==========
    auto get_virtual_expr = [&](const VExprSPtr& expr,
                                std::shared_ptr<VirtualSlotRef>& slot_ref) -> VExprSPtr {
        auto virtual_ref = std::dynamic_pointer_cast<VirtualSlotRef>(expr);
        if (virtual_ref != nullptr) {
            DCHECK(virtual_ref->get_virtual_column_expr() != nullptr);
            slot_ref = virtual_ref;
            return virtual_ref->get_virtual_column_expr();
        }
        return expr;
    };

    std::shared_ptr<VirtualSlotRef> vir_slot_ref;
    auto normalized_left = get_virtual_expr(left_child, vir_slot_ref);

    // Try to find the distance function call, it may be wrapped in a Cast(Float->Double)
    std::shared_ptr<VectorizedFnCall> function_call =
            std::dynamic_pointer_cast<VectorizedFnCall>(normalized_left);
    bool has_float_to_double_cast = false;

    if (function_call == nullptr) {
        // Check if it's a Cast expression wrapping a function call
        auto cast_expr = std::dynamic_pointer_cast<VCastExpr>(normalized_left);
        if (cast_expr == nullptr) {
            mark_unsuitable("Left child is neither a function call nor a cast expression.");
            return;
        }
        has_float_to_double_cast = true;
        auto normalized_cast_child = get_virtual_expr(cast_expr->get_child(0), vir_slot_ref);
        function_call = std::dynamic_pointer_cast<VectorizedFnCall>(normalized_cast_child);
        if (function_call == nullptr) {
            mark_unsuitable("Left child of cast is not a function call.");
            return;
        }
    }

    // Check if it's a supported distance function
    if (DISTANCE_FUNCS.find(function_call->_function_name) == DISTANCE_FUNCS.end()) {
        mark_unsuitable(fmt::format("Left child is not a supported distance function: {}",
                                    function_call->_function_name));
        return;
    }

    // Strip the _approximate suffix to get metric type
    std::string metric_name = function_call->_function_name;
    metric_name = metric_name.substr(0, metric_name.size() - 12);
    range_search_runtime.metric_type = segment_v2::string_to_metric(metric_name);

    // ========== Step 2: Validate distance function arguments ==========
    // Identify the slot ref child and the constant query array child (ArrayLiteral or CAST to array)
    Int32 idx_of_slot_ref = -1;
    Int32 idx_of_array_expr = -1;
    auto classify_child = [&](const VExprSPtr& child, UInt16 index) {
        if (idx_of_slot_ref == -1 && std::dynamic_pointer_cast<VSlotRef>(child) != nullptr) {
            idx_of_slot_ref = index;
            return;
        }
        if (idx_of_array_expr == -1 &&
            (std::dynamic_pointer_cast<VArrayLiteral>(child) != nullptr ||
             std::dynamic_pointer_cast<VCastExpr>(child) != nullptr)) {
            idx_of_array_expr = index;
        }
    };

    for (UInt16 i = 0; i < function_call->get_num_children(); ++i) {
        classify_child(function_call->get_child(i), i);
    }

    if (idx_of_slot_ref == -1 || idx_of_array_expr == -1) {
        mark_unsuitable("slot ref or array literal/cast is missing.");
        return;
    }

    auto slot_ref = std::dynamic_pointer_cast<VSlotRef>(
            function_call->get_child(static_cast<UInt16>(idx_of_slot_ref)));
    range_search_runtime.src_col_idx = slot_ref->column_id();
    range_search_runtime.dst_col_idx = vir_slot_ref == nullptr ? -1 : vir_slot_ref->column_id();

    // Materialize the constant array expression and validate its shape and types
    auto array_expr = function_call->get_child(static_cast<UInt16>(idx_of_array_expr));
    auto extract_result = extract_query_vector(array_expr);
    if (!extract_result.has_value()) {
        mark_unsuitable("Failed to extract query vector from constant array expression.");
        return;
    }
    range_search_runtime.query_value = extract_result.value();
    range_search_runtime.dim = range_search_runtime.query_value->size();

    // ========== Step 3: Check right child - must be a float/double literal ==========
    auto right_literal = std::dynamic_pointer_cast<VLiteral>(right_child);
    if (right_literal == nullptr) {
        mark_unsuitable("Right child is not a literal.");
        return;
    }

    // Handle nullable literal gracefully - just mark as unsuitable instead of crash
    if (right_literal->is_nullable()) {
        mark_unsuitable("Right literal is nullable, not supported for ANN range search.");
        return;
    }

    auto right_type = right_literal->get_data_type();
    PrimitiveType right_primitive = right_type->get_primitive_type();
    const bool float32_literal = right_primitive == PrimitiveType::TYPE_FLOAT;
    const bool float64_literal = right_primitive == PrimitiveType::TYPE_DOUBLE;

    if (!float32_literal && !float64_literal) {
        mark_unsuitable("Right child is not a Float32Literal or Float64Literal.");
        return;
    }

    // Validate consistency: if we have Cast(Float->Double), right must be double literal
    if (has_float_to_double_cast && !float64_literal) {
        mark_unsuitable("Cast expression expects double literal on right side.");
        return;
    }

    // Extract radius value
    auto right_col = right_literal->get_column_ptr()->convert_to_full_column_if_const();
    if (float32_literal) {
        const ColumnFloat32* cf32_right = assert_cast<const ColumnFloat32*>(right_col.get());
        range_search_runtime.radius = cf32_right->get_data()[0];
    } else {
        const ColumnFloat64* cf64_right = assert_cast<const ColumnFloat64*>(right_col.get());
        range_search_runtime.radius = static_cast<float>(cf64_right->get_data()[0]);
    }

    // ========== Done: Mark as suitable for ANN range search ==========
    range_search_runtime.is_ann_range_search = true;
    range_search_runtime.user_params = user_params;
    VLOG_DEBUG << fmt::format("Ann range search params: {}", range_search_runtime.to_string());
    return;
}

Status VectorizedFnCall::evaluate_ann_range_search(
        const segment_v2::AnnRangeSearchRuntime& range_search_runtime,
        const std::vector<std::unique_ptr<segment_v2::IndexIterator>>& cid_to_index_iterators,
        const std::vector<ColumnId>& idx_to_cid,
        const std::vector<std::unique_ptr<segment_v2::ColumnIterator>>& column_iterators,
        size_t rows_of_segment, roaring::Roaring& row_bitmap,
        segment_v2::AnnIndexStats& ann_index_stats, bool enable_result_cache,
        AnnRangeSearchEvaluationResult& evaluation_result) {
    evaluation_result = {};
    if (range_search_runtime.is_ann_range_search == false) {
        return Status::OK();
    }

    VLOG_DEBUG << fmt::format("Try apply ann range search. Local search params: {}",
                              range_search_runtime.to_string());
    size_t origin_num = row_bitmap.cardinality();

    const auto idx_in_block = range_search_runtime.src_col_idx;
    DCHECK_LT(idx_in_block, idx_to_cid.size())
            << "idx_in_block: " << idx_in_block << ", idx_to_cid.size(): " << idx_to_cid.size();

    ColumnId src_col_cid = idx_to_cid[idx_in_block];
    DCHECK(src_col_cid < cid_to_index_iterators.size());
    segment_v2::IndexIterator* index_iterator = cid_to_index_iterators[src_col_cid].get();
    if (index_iterator == nullptr) {
        VLOG_DEBUG << "ANN range search skipped: "
                   << fmt::format("No index iterator for column cid {}", src_col_cid);
        ;
        return Status::OK();
    }

    segment_v2::AnnIndexIterator* ann_index_iterator =
            dynamic_cast<segment_v2::AnnIndexIterator*>(index_iterator);
    if (ann_index_iterator == nullptr) {
        VLOG_DEBUG << "ANN range search skipped: "
                   << fmt::format("Column cid {} has no ANN index iterator", src_col_cid);
        return Status::OK();
    }
    DCHECK(ann_index_iterator->get_reader(AnnIndexReaderType::ANN) != nullptr)
            << "Ann index iterator should have reader. Column cid: " << src_col_cid;
    std::shared_ptr<AnnIndexReader> ann_index_reader = std::dynamic_pointer_cast<AnnIndexReader>(
            ann_index_iterator->get_reader(segment_v2::AnnIndexReaderType::ANN));
    DCHECK(ann_index_reader != nullptr)
            << "Ann index reader should not be null. Column cid: " << src_col_cid;
    // Check if metrics type is match.
    if (ann_index_reader->get_metric_type() != range_search_runtime.metric_type) {
        VLOG_DEBUG << "ANN range search skipped: "
                   << fmt::format("Metric type mismatch. Index={} Query={}",
                                  segment_v2::metric_to_string(ann_index_reader->get_metric_type()),
                                  segment_v2::metric_to_string(range_search_runtime.metric_type));
        return Status::OK();
    }

    // Check dimension if available (>0)
    const size_t index_dim = ann_index_reader->get_dimension();
    if (index_dim > 0 && index_dim != range_search_runtime.dim) {
        return Status::InvalidArgument(
                "Ann range search query dimension {} does not match index dimension {}",
                range_search_runtime.dim, index_dim);
    }

    const auto& user_params = range_search_runtime.user_params;
    if (user_params.should_fallback_ann_index_by_small_candidate(origin_num, rows_of_segment)) {
        VLOG_DEBUG << fmt::format(
                "Ann range search input rows {} reach small candidate threshold, "
                "rows_of_segment: {}, absolute_threshold: {}, percent_threshold: {}, "
                "will not use ann index to filter",
                origin_num, rows_of_segment, user_params.ann_index_candidate_rows_threshold,
                user_params.ann_index_candidate_rows_percent_threshold);
        ann_index_stats.fall_back_brute_force_cnt += 1;
        ann_index_stats.range_fallback_by_small_candidate_cnt += 1;
        ann_index_stats.range_fallback_small_candidate_rows += origin_num;
        return Status::OK();
    }

    auto stats = std::make_unique<segment_v2::AnnIndexStats>();
    // Track load index timing
    {
        SCOPED_TIMER(&(stats->load_index_costs_ns));
        if (!ann_index_iterator->try_load_index()) {
            VLOG_DEBUG << "ANN range search skipped: "
                       << fmt::format("Failed to load ANN index for column cid {}", src_col_cid);
            ann_index_stats.fall_back_brute_force_cnt += 1;
            return Status::OK();
        }
        double load_costs_ms = static_cast<double>(stats->load_index_costs_ns.value()) / 1000000.0;
        DorisMetrics::instance()->ann_index_load_costs_ms->increment(
                static_cast<int64_t>(load_costs_ms));
    }

    AnnRangeSearchParams params = range_search_runtime.to_range_search_params();

    params.roaring = &row_bitmap;
    params.enable_result_cache = enable_result_cache;
    DCHECK(params.roaring != nullptr);
    DCHECK(params.query_value != nullptr);
    segment_v2::AnnRangeSearchResult result;
    RETURN_IF_ERROR(ann_index_iterator->range_search(params, range_search_runtime.user_params,
                                                     &result, stats.get()));

#ifndef NDEBUG
    if (range_search_runtime.is_le_or_lt == false &&
        ann_index_reader->get_metric_type() == AnnIndexMetric::L2) {
        DCHECK(result.distance == nullptr) << "Should not have distance";
    }
    if (range_search_runtime.is_le_or_lt == true &&
        ann_index_reader->get_metric_type() == AnnIndexMetric::IP) {
        DCHECK(result.distance == nullptr);
    }
#endif
    DCHECK(result.roaring != nullptr);
    row_bitmap = *result.roaring;

    // Process virtual column
    bool dist_fulfilled = false;
    if (range_search_runtime.dst_col_idx >= 0) {
        // Prepare materialization if we can use result from index.
        // Typical situation: range search and operator is LE or LT.
        if (result.distance != nullptr) {
            DCHECK(result.row_ids != nullptr);
            ColumnId dst_col_cid = idx_to_cid[range_search_runtime.dst_col_idx];
            DCHECK(dst_col_cid < column_iterators.size());
            DCHECK(column_iterators[dst_col_cid] != nullptr);
            segment_v2::ColumnIterator* column_iterator = column_iterators[dst_col_cid].get();
            DCHECK(column_iterator != nullptr);
            segment_v2::VirtualColumnIterator* virtual_column_iterator =
                    dynamic_cast<segment_v2::VirtualColumnIterator*>(column_iterator);
            DCHECK(virtual_column_iterator != nullptr);
            // Now convert distance to column
            size_t size = result.roaring->cardinality();
            auto distance_col = ColumnFloat32::create(size);
            const float* src = result.distance.get();
            float* dst = distance_col->get_data().data();
            for (size_t i = 0; i < size; ++i) {
                dst[i] = src[i];
            }
            virtual_column_iterator->prepare_materialization(std::move(distance_col),
                                                             std::move(result.row_ids));
            dist_fulfilled = true;
        } else {
            // Whether the ANN index should have produced distance depends on metric and operator:
            //  - L2: distance is produced for LE/LT; not produced for GE/GT
            //  - IP: distance is produced for GE/GT; not produced for LE/LT
#ifndef NDEBUG
            const bool should_have_distance =
                    (range_search_runtime.is_le_or_lt &&
                     range_search_runtime.metric_type == AnnIndexMetric::L2) ||
                    (!range_search_runtime.is_le_or_lt &&
                     range_search_runtime.metric_type == AnnIndexMetric::IP);
            // If we expected distance but didn't get it, assert in debug to catch logic errors.
            DCHECK(!should_have_distance) << "Expected distance from ANN index but got none";
#endif
        }
    } else {
        // Dest is not virtual column.
        dist_fulfilled = true;
    }

    evaluation_result.executed = true;
    evaluation_result.dist_fulfilled = dist_fulfilled;
    VLOG_DEBUG << fmt::format(
            "Ann range search filtered {} rows, origin {} rows, virtual column is full-filled: {}",
            origin_num - row_bitmap.cardinality(), origin_num, dist_fulfilled);

    ann_index_stats = *stats;
    return Status::OK();
}

double VectorizedFnCall::execute_cost() const {
    if (!_function) {
        throw Exception(
                Status::InternalError("Function is null in expression: {}", this->debug_string()));
    }
    double cost = _function->execute_cost();
    for (const auto& child : _children) {
        cost += child->execute_cost();
    }
    return cost;
}

} // namespace doris

Coverage Report

Created: 2026-07-27 18:41