be/src/exprs/function/array/function_array_element.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/arrayElement.cpp
// and modified by Doris

#pragma once

#include <glog/logging.h>
#include <string.h>

#include <algorithm>
#include <boost/iterator/iterator_facade.hpp>
#include <memory>
#include <optional>
#include <ostream>
#include <string>
#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/block/columns_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_decimal.h"
#include "core/column/column_map.h"
#include "core/column/column_nullable.h"
#include "core/column/column_string.h"
#include "core/column/column_struct.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_map.h"
#include "core/data_type/data_type_nullable.h"
#include "core/data_type/data_type_number.h"
#include "core/data_type/data_type_struct.h"
#include "core/data_type/primitive_type.h"
#include "core/types.h"
#include "exprs/function/function.h"
#include "exprs/function/function_helpers.h"

namespace doris {
class FunctionContext;
} // namespace doris

namespace doris {

class FunctionArrayElement : public IFunction {
public:
    using MapIndiceDataType = DataTypeInt64;

    static constexpr auto name = "element_at";
    static FunctionPtr create() { return std::make_shared<FunctionArrayElement>(); }

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

    bool is_variadic() const override { return false; }

    bool use_default_implementation_for_nulls() const override { return false; }

    size_t get_number_of_arguments() const override { return 2; }

    // A struct field access (element_at(struct, const) / the struct_element alias) resolves to a
    // different return type depending on which field the constant index selects, so we need the
    // index column here. Array/map return types depend only on the argument types and fall through
    // to the DataTypes-based overload below.
    DataTypePtr get_return_type_impl(const ColumnsWithTypeAndName& arguments) const override {
        DataTypePtr arg_0 = remove_nullable(arguments[0].type);
        if (arg_0->get_primitive_type() == TYPE_STRUCT) {
            const auto* struct_type = check_and_get_data_type<DataTypeStruct>(arg_0.get());
            size_t index = 0;
            // Throw the concrete error (field not found / out of bound) instead of returning
            // nullptr, which the framework would report as an opaque "return type check failed".
            Status st = get_struct_element_index(*struct_type, arguments[1].column,
                                                 arguments[1].type, &index);
            if (!st.ok()) {
                throw doris::Exception(st);
            }
            return make_nullable(struct_type->get_elements()[index]);
        }
        DataTypes data_types(arguments.size());
        for (size_t i = 0; i < arguments.size(); ++i) {
            data_types[i] = arguments[i].type;
        }
        return get_return_type_impl(data_types);
    }

    DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
        DataTypePtr arg_0 = remove_nullable(arguments[0]);
        DCHECK(arg_0->get_primitive_type() == TYPE_ARRAY || arg_0->get_primitive_type() == TYPE_MAP)
                << "first argument for function: " << name
                << " should be DataTypeArray or DataTypeMap, but it is " << arg_0->get_name();
        if (arg_0->get_primitive_type() == TYPE_ARRAY) {
            DCHECK(is_int_or_bool(arguments[1]->get_primitive_type()))
                    << "second argument for function: " << name
                    << " should be Integer for array element";
            return make_nullable(
                    check_and_get_data_type<DataTypeArray>(arg_0.get())->get_nested_type());
        } else if (arg_0->get_primitive_type() == TYPE_MAP) {
            return make_nullable(
                    check_and_get_data_type<DataTypeMap>(arg_0.get())->get_value_type());
        } else {
            throw doris::Exception(
                    ErrorCode::INVALID_ARGUMENT,
                    fmt::format("element_at only support array and map so far, but got {}",
                                arg_0->get_name()));
        }
    }

    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                        uint32_t result, size_t input_rows_count) const override {
        if (remove_nullable(block.get_by_position(arguments[0]).type)->get_primitive_type() ==
            TYPE_STRUCT) {
            return _execute_struct(block, arguments, result, input_rows_count);
        }
        auto dst_null_column = ColumnUInt8::create(input_rows_count, 0);
        UInt8* dst_null_map = dst_null_column->get_data().data();
        const UInt8* src_null_map = nullptr;
        ColumnsWithTypeAndName args;
        ColumnPtr res_column = nullptr;

        auto col_left_raw = block.get_by_position(arguments[0]);
        // Map element lookup requires row-aligned offsets; keep original path for map type.
        if (remove_nullable(col_left_raw.type)->get_primitive_type() == TYPE_MAP) {
            block.replace_by_position(arguments[0],
                                      col_left_raw.column->convert_to_full_column_if_const());
            auto col_left = block.get_by_position(arguments[0]);
            if (col_left.column->is_nullable()) {
                const auto* null_col = assert_cast<const ColumnNullable*>(col_left.column.get());
                src_null_map = null_col->get_null_map_column().get_data().data();
                args = {{null_col->get_nested_column_ptr(), remove_nullable(col_left.type),
                         col_left.name},
                        block.get_by_position(arguments[1])};
            } else {
                args = {col_left, block.get_by_position(arguments[1])};
            }
            res_column = _execute_map(args, input_rows_count, src_null_map, dst_null_map);
        } else {
            // Array element access: avoid materializing a constant array column.
            // A literal like [[1],[2]] becomes ColumnConst with a single-row inner column;
            // unpack_if_const() gives us that inner column plus a constancy flag so inner
            // loops can use index_check_const() instead of expanding N copies.
            auto [unpacked_col, is_const_array] = unpack_if_const(col_left_raw.column);
            if (unpacked_col->is_nullable()) {
                const auto* null_col = assert_cast<const ColumnNullable*>(unpacked_col.get());
                src_null_map = null_col->get_null_map_column().get_data().data();
                args = {{null_col->get_nested_column_ptr(), remove_nullable(col_left_raw.type),
                         col_left_raw.name},
                        block.get_by_position(arguments[1])};
            } else {
                args = {{unpacked_col, col_left_raw.type, col_left_raw.name},
                        block.get_by_position(arguments[1])};
            }
            res_column = _execute_nullable(args, input_rows_count, src_null_map, dst_null_map,
                                           is_const_array);
        }
        if (!res_column) {
            return Status::RuntimeError("unsupported types for function {}({}, {})", get_name(),
                                        block.get_by_position(arguments[0]).type->get_name(),
                                        block.get_by_position(arguments[1]).type->get_name());
        }
        block.replace_by_position(result,
                                  ColumnNullable::create(res_column, std::move(dst_null_column)));
        return Status::OK();
    }

private:
    //=========================== struct element===========================//
    // Resolve the 0-based field offset selected by a constant int/string index. Mirrors the logic
    // of the former struct_element function, which element_at now subsumes.
    Status get_struct_element_index(const DataTypeStruct& struct_type,
                                    const ColumnPtr& index_column, const DataTypePtr& index_type,
                                    size_t* result) const {
        if (!index_column) {
            return Status::RuntimeError("Function {}: second argument column is nullptr.",
                                        get_name());
        }
        size_t index = 0;
        if (is_int_or_bool(index_type->get_primitive_type())) {
            int64_t offset = index_column->get_int(0);
            size_t limit = struct_type.get_elements().size() + 1;
            if (offset < 1 || offset >= static_cast<int64_t>(limit)) {
                return Status::RuntimeError(
                        "Index out of bound for function {}: index {} should base from 1 and less "
                        "than {}.",
                        get_name(), offset, limit);
            }
            index = offset - 1; // the index starts from 1
        } else if (is_string_type(index_type->get_primitive_type())) {
            std::string field_name = index_column->get_data_at(0).to_string();
            std::optional<size_t> pos = struct_type.try_get_position_by_name(field_name);
            if (!pos.has_value()) {
                return Status::RuntimeError(
                        "Element not found for function {}: name {} not found in {}.", get_name(),
                        field_name, struct_type.get_name());
            }
            index = pos.value();
        } else {
            return Status::RuntimeError(
                    "Argument not supported for function {}: second arg type {} should be int or "
                    "string.",
                    get_name(), index_type->get_name());
        }
        *result = index;
        return Status::OK();
    }

    Status _execute_struct(Block& block, const ColumnNumbers& arguments, uint32_t result,
                           size_t input_rows_count) const {
        const auto& struct_arg = block.get_by_position(arguments[0]);
        ColumnPtr struct_col_ptr = struct_arg.column->convert_to_full_column_if_const();
        // element_at manages nulls itself (use_default_implementation_for_nulls() == false), so a
        // null struct row must be merged into the result null map manually.
        const ColumnUInt8* outer_null_map = nullptr;
        if (struct_col_ptr->is_nullable()) {
            const auto* nullable = assert_cast<const ColumnNullable*>(struct_col_ptr.get());
            outer_null_map = &nullable->get_null_map_column();
            struct_col_ptr = nullable->get_nested_column_ptr();
        }
        const auto* struct_type =
                check_and_get_data_type<DataTypeStruct>(remove_nullable(struct_arg.type).get());
        const auto* struct_col = check_and_get_column<ColumnStruct>(struct_col_ptr.get());
        if (!struct_col || !struct_type) {
            return Status::RuntimeError("unsupported types for function {}({}, {})", get_name(),
                                        struct_arg.type->get_name(),
                                        block.get_by_position(arguments[1]).type->get_name());
        }
        const auto& index_arg = block.get_by_position(arguments[1]);
        size_t index = 0;
        RETURN_IF_ERROR(
                get_struct_element_index(*struct_type, index_arg.column, index_arg.type, &index));

        ColumnPtr field_col = struct_col->get_column_ptr(index);
        auto res_null_column = ColumnUInt8::create(input_rows_count, 0);
        auto& res_null_map = res_null_column->get_data();
        ColumnPtr res_nested = field_col;
        if (field_col->is_nullable()) {
            const auto* field_nullable = assert_cast<const ColumnNullable*>(field_col.get());
            const auto& field_null_map = field_nullable->get_null_map_column().get_data();
            memcpy(res_null_map.data(), field_null_map.data(), input_rows_count);
            res_nested = field_nullable->get_nested_column_ptr();
        }
        if (outer_null_map) {
            const auto& outer = outer_null_map->get_data();
            for (size_t i = 0; i < input_rows_count; ++i) {
                res_null_map[i] |= outer[i];
            }
        }
        block.replace_by_position(
                result, ColumnNullable::create(res_nested->clone_resized(input_rows_count),
                                               std::move(res_null_column)));
        return Status::OK();
    }

    //=========================== map element===========================//
    ColumnPtr _get_mapped_idx(const ColumnArray& column,
                              const ColumnWithTypeAndName& argument) const {
        auto right_column = make_nullable(argument.column->convert_to_full_column_if_const());
        const ColumnArray::Offsets64& offsets = column.get_offsets();
        ColumnPtr nested_ptr = make_nullable(column.get_data_ptr());
        size_t rows = offsets.size();
        // prepare return data
        auto matched_indices = ColumnVector<MapIndiceDataType::PType>::create();
        matched_indices->reserve(rows);

        for (size_t i = 0; i < rows; i++) {
            bool matched = false;
            size_t begin = offsets[i - 1];
            size_t end = offsets[i];
            for (size_t j = begin; j < end; j++) {
                if (nested_ptr->compare_at(j, i, *right_column, -1) == 0) {
                    matched_indices->insert_value(
                            cast_set<MapIndiceDataType::FieldType, size_t, false>(j - begin + 1));
                    matched = true;
                    break;
                }
            }

            if (!matched) {
                matched_indices->insert_value(cast_set<MapIndiceDataType::FieldType, size_t, false>(
                        end - begin + 1)); // make indices for null
            }
        }

        return matched_indices;
    }

    template <typename ColumnType, typename IndexColumnType>
    ColumnPtr _execute_number(const ColumnArray::Offsets64& offsets, const IColumn& nested_column,
                              const UInt8* arr_null_map, const IColumn& indices,
                              const UInt8* nested_null_map, UInt8* dst_null_map,
                              const UInt8* idx_null_map, bool is_const_index, bool is_const_array,
                              size_t input_rows_count) const {
        const auto& nested_data = reinterpret_cast<const ColumnType&>(nested_column).get_data();
        const auto& index_data = assert_cast<const IndexColumnType&>(indices).get_data();

        auto dst_column = nested_column.clone_empty();
        auto& dst_data = reinterpret_cast<ColumnType&>(*dst_column).get_data();
        dst_data.resize(input_rows_count);

        for (size_t row = 0; row < input_rows_count; ++row) {
            size_t arr_row = index_check_const(row, is_const_array);
            size_t off = arr_row == 0 ? 0 : offsets[arr_row - 1];
            size_t len = offsets[arr_row] - off;
            size_t idx = index_check_const(row, is_const_index);
            auto index =
                    (idx_null_map && idx_null_map[idx]) ? 0 : static_cast<Int64>(index_data[idx]);
            bool null_flag = bool(arr_null_map && arr_null_map[arr_row]);
            if (!null_flag && index > 0 && index <= len) {
                index += off - 1;
            } else if (!null_flag && index < 0 && -index <= len) {
                index += off + len;
            } else {
                null_flag = true;
            }
            if (!null_flag && nested_null_map && nested_null_map[index]) {
                null_flag = true;
            }
            dst_null_map[row] = null_flag;
            dst_data[row] = !null_flag ? nested_data[index] : typename ColumnType::value_type();
        }
        return dst_column;
    }

    template <typename IndexColumnType>
    ColumnPtr _execute_string(const ColumnArray::Offsets64& offsets, const IColumn& nested_column,
                              const UInt8* arr_null_map, const IColumn& indices,
                              const UInt8* nested_null_map, UInt8* dst_null_map,
                              const UInt8* idx_null_map, bool is_const_index, bool is_const_array,
                              size_t input_rows_count) const {
        const auto& src_str_offs =
                reinterpret_cast<const ColumnString&>(nested_column).get_offsets();
        const auto& src_str_chars =
                reinterpret_cast<const ColumnString&>(nested_column).get_chars();
        const auto& index_data = assert_cast<const IndexColumnType&>(indices).get_data();

        // prepare return data
        auto dst_column = ColumnString::create();
        auto& dst_str_offs = dst_column->get_offsets();
        dst_str_offs.resize(input_rows_count);
        auto& dst_str_chars = dst_column->get_chars();
        dst_str_chars.reserve(src_str_chars.size());

        for (size_t row = 0; row < input_rows_count; ++row) {
            size_t arr_row = index_check_const(row, is_const_array);
            size_t off = arr_row == 0 ? 0 : offsets[arr_row - 1];
            size_t len = offsets[arr_row] - off;
            size_t idx = index_check_const(row, is_const_index);
            auto index =
                    (idx_null_map && idx_null_map[idx]) ? 0 : static_cast<Int64>(index_data[idx]);
            bool null_flag = bool(arr_null_map && arr_null_map[arr_row]);
            if (!null_flag && index > 0 && index <= len) {
                index += off - 1;
            } else if (!null_flag && index < 0 && -index <= len) {
                index += off + len;
            } else {
                null_flag = true;
            }
            if (!null_flag && nested_null_map && nested_null_map[index]) {
                null_flag = true;
            }
            if (!null_flag) {
                dst_null_map[row] = false;
                auto element_size = src_str_offs[index] - src_str_offs[index - 1];
                dst_str_offs[row] = (row == 0 ? 0 : dst_str_offs[row - 1]) + element_size;
                auto src_string_pos = src_str_offs[index - 1];
                auto dst_string_pos = row == 0 ? 0 : dst_str_offs[row - 1];
                dst_str_chars.resize(dst_string_pos + element_size);
                memcpy(&dst_str_chars[dst_string_pos], &src_str_chars[src_string_pos],
                       element_size);
            } else {
                dst_null_map[row] = true;
                dst_str_offs[row] = row == 0 ? 0 : dst_str_offs[row - 1];
            }
        }
        return dst_column;
    }

    ColumnPtr _execute_map(const ColumnsWithTypeAndName& arguments, size_t input_rows_count,
                           const UInt8* src_null_map, UInt8* dst_null_map) const {
        auto left_column = arguments[0].column->convert_to_full_column_if_const();
        DataTypePtr val_type =
                reinterpret_cast<const DataTypeMap&>(*arguments[0].type).get_value_type();
        const auto& map_column = reinterpret_cast<const ColumnMap&>(*left_column);

        // create column array to find keys
        auto key_arr = ColumnArray::create(map_column.get_keys_ptr(), map_column.get_offsets_ptr());
        auto val_arr =
                ColumnArray::create(map_column.get_values_ptr(), map_column.get_offsets_ptr());

        ColumnPtr matched_indices = _get_mapped_idx(*key_arr, arguments[1]);
        if (!matched_indices) {
            return nullptr;
        }
        DataTypePtr indices_type(std::make_shared<MapIndiceDataType>());
        ColumnWithTypeAndName indices(matched_indices, indices_type, "indices");
        ColumnWithTypeAndName data(std::move(val_arr), std::make_shared<DataTypeArray>(val_type),
                                   "value");
        ColumnsWithTypeAndName args = {data, indices};
        return _execute_nullable(args, input_rows_count, src_null_map, dst_null_map, false);
    }

    template <typename IndexColumnType>
    ColumnPtr _execute_common(const ColumnArray::Offsets64& offsets, const IColumn& nested_column,
                              const UInt8* arr_null_map, const IColumn& indices,
                              const UInt8* nested_null_map, UInt8* dst_null_map,
                              const UInt8* idx_null_map, bool is_const_index, bool is_const_array,
                              size_t input_rows_count) const {
        const auto& index_data = assert_cast<const IndexColumnType&>(indices).get_data();

        auto dst_column = nested_column.clone_empty();
        dst_column->reserve(input_rows_count);

        for (size_t row = 0; row < input_rows_count; ++row) {
            size_t arr_row = index_check_const(row, is_const_array);
            size_t off = arr_row == 0 ? 0 : offsets[arr_row - 1];
            size_t len = offsets[arr_row] - off;
            size_t idx = index_check_const(row, is_const_index);
            auto index =
                    (idx_null_map && idx_null_map[idx]) ? 0 : static_cast<Int64>(index_data[idx]);
            bool null_flag = bool(arr_null_map && arr_null_map[arr_row]);
            if (!null_flag && index > 0 && index <= len) {
                index += off - 1;
            } else if (!null_flag && index < 0 && -index <= len) {
                index += off + len;
            } else {
                null_flag = true;
            }
            if (!null_flag && nested_null_map && nested_null_map[index]) {
                null_flag = true;
            }
            if (!null_flag) {
                dst_null_map[row] = false;
                dst_column->insert_from(nested_column, index);
            } else {
                dst_null_map[row] = true;
                dst_column->insert_default();
            }
        }
        return dst_column;
    }

    ColumnPtr _execute_nullable(const ColumnsWithTypeAndName& arguments, size_t input_rows_count,
                                const UInt8* src_null_map, UInt8* dst_null_map,
                                bool is_const_array) const {
        // arguments[0].column is already the raw ColumnArray (possibly size-1 when is_const_array).
        // Do NOT call convert_to_full_column_if_const() here; const-awareness is handled below
        // via index_check_const(row, is_const_array).
        const auto& array_column = assert_cast<const ColumnArray&>(*arguments[0].column);
        const auto& offsets = array_column.get_offsets();
        DCHECK(is_const_array ? offsets.size() == 1 : offsets.size() == input_rows_count);
        const UInt8* nested_null_map = nullptr;
        ColumnPtr nested_column = nullptr;
        if (is_column_nullable(array_column.get_data())) {
            const auto& nested_null_column =
                    reinterpret_cast<const ColumnNullable&>(array_column.get_data());
            nested_null_map = nested_null_column.get_null_map_column().get_data().data();
            nested_column = nested_null_column.get_nested_column_ptr();
        } else {
            nested_column = array_column.get_data_ptr();
        }

        ColumnPtr res = nullptr;
        auto left_element_type = remove_nullable(
                assert_cast<const DataTypeArray&>(*remove_nullable(arguments[0].type))
                        .get_nested_type());
        const UInt8* idx_null_map = nullptr;
        auto idx_col_with_const = unpack_if_const(arguments[1].column);
        if (idx_col_with_const.first->is_nullable()) {
            const auto& idx_null_column =
                    reinterpret_cast<const ColumnNullable&>(*idx_col_with_const.first);
            idx_null_map = idx_null_column.get_null_map_column().get_data().data();
        }
        auto idx_col_raw = remove_nullable(idx_col_with_const.first);
        bool is_const_index = idx_col_with_const.second;

        PrimitiveType idx_ptype = remove_nullable(arguments[1].type)->get_primitive_type();

        // Outer dispatch on index column type (Int8/Int16/Int32/Int64),
        // inner dispatch on nested data column type.
        auto idx_dispatch = [&](const auto& idx_type) -> bool {
            using IdxDispatchType = std::decay_t<decltype(idx_type)>;
            using IndexColumnType = typename IdxDispatchType::ColumnType;

            auto data_call = [&](const auto& data_type) -> bool {
                using DataDispatchType = std::decay_t<decltype(data_type)>;
                res = _execute_number<typename DataDispatchType::ColumnType, IndexColumnType>(
                        offsets, *nested_column, src_null_map, *idx_col_raw, nested_null_map,
                        dst_null_map, idx_null_map, is_const_index, is_const_array,
                        input_rows_count);
                return true;
            };

            if (is_string_type(left_element_type->get_primitive_type())) {
                res = _execute_string<IndexColumnType>(offsets, *nested_column, src_null_map,
                                                       *idx_col_raw, nested_null_map, dst_null_map,
                                                       idx_null_map, is_const_index, is_const_array,
                                                       input_rows_count);
            } else if (!dispatch_switch_scalar(left_element_type->get_primitive_type(),
                                               data_call)) {
                res = _execute_common<IndexColumnType>(offsets, *nested_column, src_null_map,
                                                       *idx_col_raw, nested_null_map, dst_null_map,
                                                       idx_null_map, is_const_index, is_const_array,
                                                       input_rows_count);
            }
            return true;
        };

        bool dispatched = dispatch_switch_int(idx_ptype, idx_dispatch);
        DCHECK(dispatched) << "Unsupported index column type for element_at: " << idx_ptype;
        return res;
    }
};

} // namespace doris

Coverage Report

Created: 2026-06-10 13:17