BiteTheDDDDt commented on code in PR #67311: URL: https://github.com/apache/doris/pull/67311#discussion_r3904839613
########## be/src/exprs/function/function_inner_product.h: ########## @@ -0,0 +1,317 @@ +// 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 "core/assert_cast.h" +#include "core/column/column_const.h" +#include "core/column/column_map.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" +#include "core/data_type/data_type_map.h" +#include "core/string_ref.h" +#include "exec/common/hash_table/hash.h" +#include "exec/common/hash_table/phmap_fwd_decl.h" +#include "exec/common/util.hpp" +#include "exprs/function/array/function_array_distance.h" + +namespace doris { + +namespace detail { + +template <PrimitiveType KeyType> +struct InnerProductMapKeyTraits { + using ColumnType = PrimitiveTypeTraits<KeyType>::ColumnType; + using Key = PrimitiveTypeTraits<KeyType>::CppType; + using KeyAccessor = const Key*; + using Hash = HashCRC32<Key>; + + static KeyAccessor get_key_accessor(const ColumnType& column) { + return column.get_data().data(); + } + + static Key get_key(KeyAccessor keys, size_t index) { return keys[index]; } +}; + +template <> +struct InnerProductMapKeyTraits<TYPE_STRING> { + using ColumnType = ColumnString; + using Key = StringRef; + using KeyAccessor = const ColumnType*; + using Hash = StringRefHash; + + static KeyAccessor get_key_accessor(const ColumnType& column) { return &column; } + + static Key get_key(KeyAccessor keys, size_t index) { return keys->get_data_at(index); } +}; + +} // namespace detail + +class FunctionInnerProduct final : public FunctionArrayDistance<InnerProduct> { +public: + static FunctionPtr create() { return std::make_shared<FunctionInnerProduct>(); } + + DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { + if (arguments.size() != 2) { + throw doris::Exception(ErrorCode::INVALID_ARGUMENT, "Invalid number of arguments"); + } + + const bool both_arrays = arguments[0]->get_primitive_type() == TYPE_ARRAY && + arguments[1]->get_primitive_type() == TYPE_ARRAY; + if (both_arrays) { + return FunctionArrayDistance<InnerProduct>::get_return_type_impl(arguments); + } + + const bool both_maps = arguments[0]->get_primitive_type() == TYPE_MAP && + arguments[1]->get_primitive_type() == TYPE_MAP; + if (!both_maps) { + throw doris::Exception(ErrorCode::INVALID_ARGUMENT, + "Arguments for function {} must be arrays or maps", get_name()); + } + + const auto& left_type = assert_cast<const DataTypeMap&>(*remove_nullable(arguments[0])); + const auto& right_type = assert_cast<const DataTypeMap&>(*remove_nullable(arguments[1])); + if (!left_type.get_key_type()->equals(*right_type.get_key_type())) { + throw doris::Exception(ErrorCode::INVALID_ARGUMENT, + "Map keys for function {} must have the same type", get_name()); + } + const auto key_type = remove_nullable(left_type.get_key_type())->get_primitive_type(); + if (!_is_supported_map_key_type(key_type)) { + throw doris::Exception(ErrorCode::INVALID_ARGUMENT, + "Function {} only supports integer or string map keys", + get_name()); + } + if (remove_nullable(left_type.get_value_type())->get_primitive_type() != TYPE_FLOAT || + remove_nullable(right_type.get_value_type())->get_primitive_type() != TYPE_FLOAT) { + throw doris::Exception(ErrorCode::INVALID_ARGUMENT, + "Map values for function {} must be FLOAT", get_name()); + } + return std::make_shared<DataTypeFloat32>(); + } + + Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments, + uint32_t result, size_t input_rows_count) const override { + if (block.get_by_position(arguments[0]).type->get_primitive_type() == TYPE_MAP) { + return _execute_map(block, arguments, result, input_rows_count); + } + return FunctionArrayDistance<InnerProduct>::execute_impl(context, block, arguments, result, + input_rows_count); + } + +private: + using ColumnType = PrimitiveTypeTraits<TYPE_FLOAT>::ColumnType; + + struct MapRange { + size_t begin; + size_t size; + }; + + static ALWAYS_INLINE MapRange _get_map_range(const ColumnMap& map, bool is_const, size_t row) { + const size_t actual_row = index_check_const(row, is_const); + return {map.offset_at(actual_row), map.size_at(actual_row)}; + } + + static bool _is_supported_map_key_type(PrimitiveType type) { + switch (type) { + case TYPE_TINYINT: + case TYPE_SMALLINT: + case TYPE_INT: + case TYPE_BIGINT: + case TYPE_LARGEINT: + case TYPE_CHAR: + case TYPE_VARCHAR: + case TYPE_STRING: + return true; + default: + return false; + } + } + + static const ColumnMap& _get_map_column(const ColumnPtr& column, const char* argument_name, + const String& function_name, bool& is_const) { + const IColumn* raw_column = column.get(); + is_const = is_column_const(*raw_column); + if (is_const) { + raw_column = assert_cast<const ColumnConst*>(raw_column)->get_data_column_ptr().get(); + } + + if (const auto* nullable = check_and_get_column<ColumnNullable>(raw_column)) { + if (raw_column->has_null()) { + throw doris::Exception(ErrorCode::INVALID_ARGUMENT, + "{} for function {} cannot be null", argument_name, + function_name); + } + raw_column = nullable->get_nested_column_ptr().get(); + } + + const auto& map = assert_cast<const ColumnMap&>(*raw_column); + if (map.get_values().has_null()) { + throw doris::Exception(ErrorCode::INVALID_ARGUMENT, + "{} for function {} cannot have null", argument_name, + function_name); + } + return map; + } + + static const IColumn& _get_key_column(const IColumn& column, const UInt8*& null_map) { + null_map = nullptr; + if (const auto* nullable = check_and_get_column<ColumnNullable>(&column)) { + null_map = nullable->get_null_map_data().data(); + return nullable->get_nested_column(); + } + return column; + } + + template <PrimitiveType KeyType> + static void _execute_map_typed(const ColumnMap& left, bool left_is_const, + const ColumnMap& right, bool right_is_const, + ColumnType::Container& destination_data, + size_t input_rows_count) { + using KeyTraits = detail::InnerProductMapKeyTraits<KeyType>; + using Key = typename KeyTraits::Key; + using KeyAccessor = typename KeyTraits::KeyAccessor; + using KeyColumn = typename KeyTraits::ColumnType; + + const UInt8* left_key_null_map = nullptr; + const UInt8* right_key_null_map = nullptr; + const auto& left_keys = + assert_cast<const KeyColumn&>(_get_key_column(left.get_keys(), left_key_null_map)); Review Comment: Thanks for pointing this out. I checked the execution path more closely. Although DataTypeString::check_column() accepts ColumnString64 and ColumnMap::convert_column_if_overflow() recursively converts its key column, ColumnString64 is an internal representation for stateful operators and is recursively rejected at every operator boundary by Block::check_no_column_string64(). Existing map-key comparison functions do not support this representation either. map_contains_key dispatches strings through StringElementView, which casts to ColumnString; map_contains_entry dispatches TYPE_STRING to PrimitiveTypeTraits<TYPE_STRING>::ColumnType, also ColumnString; and element_at(map, key) requires matching ColumnStr<T> representations in compare_at. Could you point to a production operator/query path where a ColumnMap containing ColumnString64 can reach scalar-function execution? Without such a path, a direct ColumnString64 function UT would construct a block that violates the current operator-boundary invariant, and supporting it only in inner_product would be inconsistent with the other map functions. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
