be/src/exprs/function/array/function_array_index.h

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.
// This file is copied from
// https://github.com/ClickHouse/ClickHouse/blob/master/src/Functions/array/arrayIndex.h
// and modified by Doris
#pragma once

#include <stddef.h>

#include <memory>
#include <utility>

#include "common/status.h"
#include "core/assert_cast.h"
#include "core/block/block.h"
#include "core/block/column_numbers.h"
#include "core/block/column_with_type_and_name.h"
#include "core/call_on_type_index.h"
#include "core/column/column.h"
#include "core/column/column_array.h"
#include "core/column/column_array_view.h"
#include "core/column/column_nullable.h"
#include "core/column/column_string.h"
#include "core/column/column_vector.h"
#include "core/data_type/data_type.h"
#include "core/data_type/data_type_array.h"
#include "core/data_type/data_type_nullable.h"
#include "core/data_type/data_type_number.h" // IWYU pragma: keep
#include "core/data_type/define_primitive_type.h"
#include "core/data_type/primitive_type.h"
#include "core/field.h"
#include "core/string_ref.h"
#include "core/types.h"
#include "exprs/function/function.h"
#include "storage/index/index_reader_helper.h"
#include "storage/index/inverted/inverted_index_query_type.h"
#include "storage/index/inverted/inverted_index_reader.h"
#include "storage/predicate/column_predicate.h"

namespace doris {
class FunctionContext;
} // namespace doris

namespace doris {

struct ArrayContainsAction {
    static constexpr auto ResultType = PrimitiveType::TYPE_BOOLEAN;
    static constexpr auto name = "array_contains";
    static constexpr const bool resume_execution = false;
    static constexpr void apply(typename PrimitiveTypeTraits<ResultType>::CppType& current,
                                size_t) noexcept {
        current = 1;
    }
};

struct ArrayPositionAction {
    static constexpr auto ResultType = PrimitiveType::TYPE_BIGINT;
    static constexpr auto name = "array_position";
    static constexpr const bool resume_execution = false;
    static constexpr void apply(typename PrimitiveTypeTraits<ResultType>::CppType& current,
                                size_t j) noexcept {
        current = j + 1;
    }
};

struct ArrayCountEqual {
    static constexpr auto ResultType = PrimitiveType::TYPE_BIGINT;
    static constexpr auto name = "countequal";
    static constexpr const bool resume_execution = true;
    static constexpr void apply(typename PrimitiveTypeTraits<ResultType>::CppType& current,
                                size_t j) noexcept {
        ++current;
    }
};

struct ParamValue {
    PrimitiveType type;
    Field value;
};

template <typename ConcreteAction>
class FunctionArrayIndex : public IFunction {
public:
    static constexpr auto ResultType = ConcreteAction::ResultType;

    static constexpr auto name = ConcreteAction::name;
    static FunctionPtr create() { return std::make_shared<FunctionArrayIndex>(); }

    /// Get function name.
    String get_name() const override { return name; }

    bool is_variadic() const override { return false; }

    size_t get_number_of_arguments() const override { return 2; }

    bool use_default_implementation_for_nulls() const override { return false; }

    Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) override {
        if (scope == FunctionContext::THREAD_LOCAL) {
            return Status::OK();
        }

        DCHECK(context->get_num_args() >= 1);
        DCHECK_EQ(context->get_arg_type(0)->get_primitive_type(), PrimitiveType::TYPE_ARRAY);
        // now we only support same
        std::shared_ptr<ParamValue> state = std::make_shared<ParamValue>();
        Field field;
        if (context->get_constant_col(1)) {
            context->get_constant_col(1)->column_ptr->get(0, field);
            state->value = field;
            state->type = context->get_arg_type(1)->get_primitive_type();
            context->set_function_state(scope, state);
        }
        return Status::OK();
    }

    Status evaluate_inverted_index(
            const ColumnsWithTypeAndName& arguments,
            const std::vector<IndexFieldNameAndTypePair>& data_type_with_names,
            std::vector<segment_v2::IndexIterator*> iterators, uint32_t num_rows,
            const InvertedIndexAnalyzerCtx* analyzer_ctx,
            segment_v2::InvertedIndexResultBitmap& bitmap_result) const override {
        DCHECK(arguments.size() == 1);
        DCHECK(data_type_with_names.size() == 1);
        DCHECK(iterators.size() == 1);
        auto* iter = iterators[0];
        auto data_type_with_name = data_type_with_names[0];
        if (iter == nullptr) {
            return Status::OK();
        }
        if (!segment_v2::IndexReaderHelper::has_string_or_bkd_index(iter)) {
            // parser is not none we can not make sure the result is correct in expr combination
            // for example, filter: !array_index(array, 'tall:120cm, weight: 35kg')
            // here we have rows [tall:120cm, weight: 35kg, hobbies: reading book] which be tokenized
            // but query is also tokenized, and FULLTEXT reader will catch this row as matched,
            // so array_index(array, 'tall:120cm, weight: 35kg') return this rowid,
            // but we expect it to be filtered, because we want row is equal to 'tall:120cm, weight: 35kg'
            return Status::OK();
        }
        Field param_value;
        arguments[0].column->get(0, param_value);
        // The current implementation for the inverted index of arrays cannot handle cases where the array contains null values,
        // meaning an item in the array is null.
        if (param_value.is_null()) {
            return Status::OK();
        }

        std::shared_ptr<roaring::Roaring> null_bitmap = std::make_shared<roaring::Roaring>();
        if (iter->has_null()) {
            segment_v2::InvertedIndexQueryCacheHandle null_bitmap_cache_handle;
            RETURN_IF_ERROR(iter->read_null_bitmap(&null_bitmap_cache_handle));
            null_bitmap = null_bitmap_cache_handle.get_bitmap();
        }
        InvertedIndexParam param;
        param.column_name = data_type_with_name.first;
        param.column_type = data_type_with_name.second;
        param.query_value = param_value;
        param.query_type = segment_v2::InvertedIndexQueryType::EQUAL_QUERY;
        param.num_rows = num_rows;
        param.roaring = std::make_shared<roaring::Roaring>();
        param.analyzer_ctx = analyzer_ctx;
        RETURN_IF_ERROR(iter->read_from_index(segment_v2::IndexParam {&param}));
        // here debug for check array_contains function really filter rows by inverted index correctly
        DBUG_EXECUTE_IF("array_func.array_contains", {
            auto result_bitmap = DebugPoints::instance()->get_debug_param_or_default<int32_t>(
                    "array_func.array_contains", "result_bitmap", 0);
            if (result_bitmap < 0) {
                return Status::Error<ErrorCode::INTERNAL_ERROR>(
                        "result_bitmap count cannot be negative");
            }
            if (param.roaring->cardinality() != result_bitmap) {
                return Status::Error<ErrorCode::INTERNAL_ERROR>(
                        "array_contains really filtered {} by inverted index not equal to expected "
                        "{}",
                        param.roaring->cardinality(), result_bitmap);
            }
        })
        segment_v2::InvertedIndexResultBitmap result(param.roaring, null_bitmap);
        bitmap_result = result;
        bitmap_result.mask_out_null();

        return Status::OK();
    }

    DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
        if (arguments[0]->is_nullable()) {
            return make_nullable(
                    std::make_shared<typename PrimitiveTypeTraits<ResultType>::DataType>());
        } else {
            return std::make_shared<typename PrimitiveTypeTraits<ResultType>::DataType>();
        }
    }

    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                        uint32_t result, size_t input_rows_count) const override {
        DBUG_EXECUTE_IF("array_func.array_contains", {
            auto req_id = DebugPoints::instance()->get_debug_param_or_default<int32_t>(
                    "array_func.array_contains", "req_id", 0);
            return Status::Error<ErrorCode::INTERNAL_ERROR>(
                    "{} has already execute inverted index req_id {} , should not execute expr "
                    "with rows: {}",
                    get_name(), req_id, input_rows_count);
        });
        return _execute_dispatch(block, arguments, result, input_rows_count);
    }

private:
    template <PrimitiveType PType>
    ColumnPtr _execute_view(const ColumnArrayView<PType>& array_view,
                            const ColumnView<PType>& right_view) const {
        // prepare return data
        auto dst = PrimitiveTypeTraits<ResultType>::ColumnType::create(array_view.size(), 0);
        auto& dst_data = dst->get_data();
        auto dst_null_column = ColumnUInt8::create(array_view.size(), 0);
        auto& dst_null_data = dst_null_column->get_data();

        // process
        for (size_t row = 0; row < array_view.size(); ++row) {
            if (array_view.is_null_at(row)) {
                dst_null_data[row] = true;
                continue;
            }
            dst_null_data[row] = false;
            typename PrimitiveTypeTraits<ResultType>::CppType res = 0;
            const auto array_data = array_view[row];
            for (size_t pos = 0; pos < array_data.size(); ++pos) {
                // match null value
                if (right_view.is_null_at(row) && array_data.is_null_at(pos)) {
                    ConcreteAction::apply(res, pos);
                    if constexpr (!ConcreteAction::resume_execution) {
                        break;
                    }
                }
                // some is null while another is not
                if (right_view.is_null_at(row) != array_data.is_null_at(pos)) {
                    continue;
                }
                if (array_data.is_null_at(pos)) {
                    continue;
                }
                if (array_data.value_at(pos) == right_view.value_at(row)) {
                    ConcreteAction::apply(res, pos);
                    if constexpr (!ConcreteAction::resume_execution) {
                        break;
                    }
                }
            }
            dst_data[row] = res;
        }

        if (!array_view.is_nullable()) {
            return dst;
        }
        return ColumnNullable::create(std::move(dst), std::move(dst_null_column));
    }

    Status _execute_dispatch(Block& block, const ColumnNumbers& arguments, uint32_t result,
                             size_t input_rows_count) const {
        if (block.get_by_position(arguments[0]).type->get_primitive_type() != TYPE_ARRAY) {
            return Status::InvalidArgument(get_name() + " first argument must be array, but got " +
                                           block.get_by_position(arguments[0]).type->get_name());
        }
        // execute
        auto array_type = remove_nullable(block.get_by_position(arguments[0]).type);
        auto left_element_type = remove_nullable(
                assert_cast<const DataTypeArray*>(array_type.get())->get_nested_type());
        auto right_type = remove_nullable(block.get_by_position(arguments[1]).type);

        auto left_element_primitive_type = left_element_type->get_primitive_type();
        auto right_primitive_type = right_type->get_primitive_type();
        ColumnPtr return_column = nullptr;
        if (right_primitive_type == left_element_primitive_type ||
            (is_string_type(right_primitive_type) && is_string_type(left_element_primitive_type))) {
            auto call = [&](const auto& type) -> bool {
                using DispatchType = std::decay_t<decltype(type)>;
                constexpr PrimitiveType PType = DispatchType::PType;
                auto array_view =
                        ColumnArrayView<PType>::create(block.get_by_position(arguments[0]).column);
                auto right_view =
                        ColumnView<PType>::create(block.get_by_position(arguments[1]).column);
                return_column = _execute_view(array_view, right_view);
                return true;
            };
            dispatch_switch_all(left_element_primitive_type, call);
        }

        if (return_column) {
            block.replace_by_position(result, std::move(return_column));
            return Status::OK();
        }
        return Status::RuntimeError("execute failed or unsupported types for function {}({}, {})",
                                    get_name(),
                                    block.get_by_position(arguments[0]).type->get_name(),
                                    block.get_by_position(arguments[1]).type->get_name());
    }
};

} // namespace doris

Coverage Report

Created: 2026-06-29 14:39