diff --git a/be/src/exprs/function/array/function_array_distance.cpp b/be/src/exprs/function/array/function_array_distance.cpp index 3f37775d6beedf..103f737edbfccf 100644 --- a/be/src/exprs/function/array/function_array_distance.cpp +++ b/be/src/exprs/function/array/function_array_distance.cpp @@ -19,6 +19,7 @@ #include +#include "exprs/function/function_inner_product.h" #include "exprs/function/simple_function_factory.h" namespace doris { @@ -90,7 +91,7 @@ void register_function_array_distance(SimpleFunctionFactory& factory) { factory.register_function>(); factory.register_function>(); factory.register_function>(); - factory.register_function>(); + factory.register_function(); factory.register_function>(); factory.register_function>(); } diff --git a/be/src/exprs/function/function_inner_product.h b/be/src/exprs/function/function_inner_product.h new file mode 100644 index 00000000000000..3485db95c6d12e --- /dev/null +++ b/be/src/exprs/function/function_inner_product.h @@ -0,0 +1,479 @@ +// 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 +struct InnerProductMapKeyTraits { + using ColumnType = PrimitiveTypeTraits::ColumnType; + using Key = PrimitiveTypeTraits::CppType; + using KeyAccessor = const Key*; + using Hash = HashCRC32; + + 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 { + 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 { +public: + static FunctionPtr create() { return std::make_shared(); } + + 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::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(*remove_nullable(arguments[0])); + const auto& right_type = assert_cast(*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(); + } + + 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::execute_impl(context, block, arguments, result, + input_rows_count); + } + +private: + using ColumnType = PrimitiveTypeTraits::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, + bool& is_const) { + const IColumn* raw_column = column.get(); + is_const = is_column_const(*raw_column); + if (is_const) { + raw_column = assert_cast(raw_column)->get_data_column_ptr().get(); + } + + if (const auto* nullable = check_and_get_column(raw_column)) { + if (raw_column->has_null()) { + throw doris::Exception(ErrorCode::INVALID_ARGUMENT, + "{} for function {} cannot be null", argument_name, + InnerProduct::name); + } + raw_column = nullable->get_nested_column_ptr().get(); + } + + return assert_cast(*raw_column); + } + + static const IColumn& _get_key_column(const IColumn& column, const UInt8*& null_map) { + null_map = nullptr; + if (const auto* nullable = check_and_get_column(&column)) { + null_map = nullable->get_null_map_data().data(); + return nullable->get_nested_column(); + } + return column; + } + + static const IColumn& _get_value_column(const IColumn& column, + const ColumnNullable*& nullable_with_null) { + nullable_with_null = nullptr; + if (const auto* nullable = check_and_get_column(&column)) { + if (nullable->has_null()) { + nullable_with_null = nullable; + } + return nullable->get_nested_column(); + } + return column; + } + + template + static void _validate_retained_values(typename KeyTraits::KeyAccessor keys, + const UInt8* key_null_map, + const ColumnNullable& nullable_values, MapRange range, + const char* argument_name) { + if (!nullable_values.has_null(range.begin, range.begin + range.size)) { + return; + } + + using Key = typename KeyTraits::Key; + doris::flat_hash_set seen_keys; + seen_keys.reserve(range.size); + const auto& value_null_map = nullable_values.get_null_map_data(); + bool has_null_key = false; + + // Only the last value for each key is visible. Ignore NULL values shadowed by a later + // duplicate, matching ColumnMap::deduplicate_keys() semantics. + for (size_t i = range.begin + range.size; i > range.begin; --i) { + const size_t index = i - 1; + if (key_null_map != nullptr && key_null_map[index]) { + if (!has_null_key) { + has_null_key = true; + if (value_null_map[index]) { + throw doris::Exception(ErrorCode::INVALID_ARGUMENT, + "{} for function {} cannot have null", argument_name, + InnerProduct::name); + } + } + continue; + } + + if (seen_keys.emplace(KeyTraits::get_key(keys, index)).second && + value_null_map[index]) { + throw doris::Exception(ErrorCode::INVALID_ARGUMENT, + "{} for function {} cannot have null", argument_name, + InnerProduct::name); + } + } + } + + template + struct MapData { + typename KeyTraits::KeyAccessor keys; + const UInt8* key_null_map; + const float* values; + const ColumnNullable* nullable_values; + }; + + template + static void _execute_map_with_cached_constant(const MapData& constant, + MapRange constant_range, + const MapData& varying, + const ColumnMap& varying_map, + const char* varying_argument_name, + ColumnType::Container& destination_data, + size_t input_rows_count) { + using Key = typename KeyTraits::Key; + struct CachedValue { + float value; + size_t last_matched_row; + }; + + doris::flat_hash_map constant_values_by_key; + constant_values_by_key.reserve(constant_range.size); + bool has_null_key = false; + float null_key_value = 0.0F; + size_t null_key_last_matched_row = input_rows_count; + + for (size_t i = constant_range.begin + constant_range.size; i > constant_range.begin; --i) { + const size_t index = i - 1; + if (constant.key_null_map != nullptr && constant.key_null_map[index]) { + if (!has_null_key) { + has_null_key = true; + null_key_value = constant.values[index]; + } + } else { + constant_values_by_key.emplace( + KeyTraits::get_key(constant.keys, index), + CachedValue {constant.values[index], input_rows_count}); + } + } + + for (size_t row = 0; row < input_rows_count; ++row) { + const MapRange varying_range = _get_map_range(varying_map, false, row); + if (varying.nullable_values != nullptr) { + _validate_retained_values(varying.keys, varying.key_null_map, + *varying.nullable_values, varying_range, + varying_argument_name); + } + if (constant_range.size == 0 || varying_range.size == 0) { + destination_data[row] = 0.0F; + continue; + } + + float inner_product = 0.0F; + for (size_t i = varying_range.begin + varying_range.size; i > varying_range.begin; + --i) { + const size_t index = i - 1; + if (varying.key_null_map != nullptr && varying.key_null_map[index]) { + if (has_null_key && null_key_last_matched_row != row) { + inner_product += null_key_value * varying.values[index]; + null_key_last_matched_row = row; + } + continue; + } + + const auto it = + constant_values_by_key.find(KeyTraits::get_key(varying.keys, index)); + if (it != constant_values_by_key.end() && it->second.last_matched_row != row) { + inner_product += it->second.value * varying.values[index]; + it->second.last_matched_row = row; + } + } + destination_data[row] = inner_product; + } + } + + template + 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; + using Key = typename KeyTraits::Key; + using KeyColumn = typename KeyTraits::ColumnType; + using TypedMapData = MapData; + + const UInt8* left_key_null_map = nullptr; + const UInt8* right_key_null_map = nullptr; + const auto& left_keys = + assert_cast(_get_key_column(left.get_keys(), left_key_null_map)); + const auto& right_keys = assert_cast( + _get_key_column(right.get_keys(), right_key_null_map)); + const ColumnNullable* left_nullable_values = nullptr; + const ColumnNullable* right_nullable_values = nullptr; + const auto& left_values = + assert_cast( + _get_value_column(left.get_values(), left_nullable_values)) + .get_data(); + const auto& right_values = + assert_cast( + _get_value_column(right.get_values(), right_nullable_values)) + .get_data(); + + const TypedMapData left_data {KeyTraits::get_key_accessor(left_keys), left_key_null_map, + left_values.data(), left_nullable_values}; + const TypedMapData right_data {KeyTraits::get_key_accessor(right_keys), right_key_null_map, + right_values.data(), right_nullable_values}; + + if (left_is_const && left_data.nullable_values != nullptr) { + _validate_retained_values(left_data.keys, left_data.key_null_map, + *left_data.nullable_values, + _get_map_range(left, true, 0), "First argument"); + } + if (right_is_const && right_data.nullable_values != nullptr) { + _validate_retained_values(right_data.keys, right_data.key_null_map, + *right_data.nullable_values, + _get_map_range(right, true, 0), "Second argument"); + } + + if (left_is_const != right_is_const && input_rows_count > 1) { + const bool constant_is_left = left_is_const; + const TypedMapData constant = constant_is_left ? left_data : right_data; + const TypedMapData varying = constant_is_left ? right_data : left_data; + const auto& varying_map = constant_is_left ? right : left; + const MapRange constant_range = + _get_map_range(constant_is_left ? left : right, true, 0); + + // Reuse the constant side only when its scratch space does not exceed the total + // varying input. Otherwise the per-row path below keeps memory bounded by the smaller + // map in each row. + if (constant_range.size <= varying_map.get_keys().size()) { + _execute_map_with_cached_constant( + constant, constant_range, varying, varying_map, + constant_is_left ? "Second argument" : "First argument", destination_data, + input_rows_count); + return; + } + } + + // Build the hash table from the smaller map row to minimize temporary memory. + doris::flat_hash_map values_by_key; + for (size_t row = 0; row < input_rows_count; ++row) { + const MapRange left_range = _get_map_range(left, left_is_const, row); + const MapRange right_range = _get_map_range(right, right_is_const, row); + if (!left_is_const && left_data.nullable_values != nullptr) { + _validate_retained_values(left_data.keys, left_data.key_null_map, + *left_data.nullable_values, left_range, + "First argument"); + } + if (!right_is_const && right_data.nullable_values != nullptr) { + _validate_retained_values(right_data.keys, right_data.key_null_map, + *right_data.nullable_values, right_range, + "Second argument"); + } + if (left_range.size == 0 || right_range.size == 0) { + destination_data[row] = 0.0F; + continue; + } + const bool build_left = left_range.size <= right_range.size; + const TypedMapData build = build_left ? left_data : right_data; + const TypedMapData probe = build_left ? right_data : left_data; + const MapRange build_range = build_left ? left_range : right_range; + const MapRange probe_range = build_left ? right_range : left_range; + + values_by_key.clear(); + values_by_key.reserve(build_range.size); + bool has_null_key = false; + float null_key_value = 0.0F; + // Scan backwards so emplace keeps the last value for duplicate keys. + for (size_t i = build_range.begin + build_range.size; i > build_range.begin; --i) { + const size_t index = i - 1; + if (build.key_null_map != nullptr && build.key_null_map[index]) { + if (!has_null_key) { + has_null_key = true; + null_key_value = build.values[index]; + } + } else { + values_by_key.emplace(KeyTraits::get_key(build.keys, index), + build.values[index]); + } + } + + float inner_product = 0.0F; + // Erase matches while scanning backwards so probe duplicates also use the last value. + for (size_t i = probe_range.begin + probe_range.size; i > probe_range.begin; --i) { + const size_t index = i - 1; + if (probe.key_null_map != nullptr && probe.key_null_map[index]) { + if (has_null_key) { + inner_product += null_key_value * probe.values[index]; + has_null_key = false; + } + continue; + } + const auto it = values_by_key.find(KeyTraits::get_key(probe.keys, index)); + if (it != values_by_key.end()) { + inner_product += it->second * probe.values[index]; + values_by_key.erase(it); + } + } + destination_data[row] = inner_product; + } + } + + Status _execute_map(Block& block, const ColumnNumbers& arguments, uint32_t result, + size_t input_rows_count) const { + bool left_is_const = false; + bool right_is_const = false; + const auto& left = _get_map_column(block.get_by_position(arguments[0]).column, + "First argument", left_is_const); + const auto& right = _get_map_column(block.get_by_position(arguments[1]).column, + "Second argument", right_is_const); + + auto destination = ColumnType::create(input_rows_count); + auto& destination_data = destination->get_data(); + const auto& map_type = assert_cast( + *remove_nullable(block.get_by_position(arguments[0]).type)); + switch (remove_nullable(map_type.get_key_type())->get_primitive_type()) { + case TYPE_TINYINT: + _execute_map_typed(left, left_is_const, right, right_is_const, + destination_data, input_rows_count); + break; + case TYPE_SMALLINT: + _execute_map_typed(left, left_is_const, right, right_is_const, + destination_data, input_rows_count); + break; + case TYPE_INT: + _execute_map_typed(left, left_is_const, right, right_is_const, + destination_data, input_rows_count); + break; + case TYPE_BIGINT: + _execute_map_typed(left, left_is_const, right, right_is_const, + destination_data, input_rows_count); + break; + case TYPE_LARGEINT: + _execute_map_typed(left, left_is_const, right, right_is_const, + destination_data, input_rows_count); + break; + case TYPE_CHAR: + case TYPE_VARCHAR: + case TYPE_STRING: + _execute_map_typed(left, left_is_const, right, right_is_const, + destination_data, input_rows_count); + break; + default: + return Status::InvalidArgument("Function {} only supports integer or string map keys", + get_name()); + } + + block.replace_by_position(result, std::move(destination)); + return Status::OK(); + } +}; + +} // namespace doris diff --git a/be/test/exprs/function/function_map_inner_product_test.cpp b/be/test/exprs/function/function_map_inner_product_test.cpp new file mode 100644 index 00000000000000..d2f1c55e1ec68a --- /dev/null +++ b/be/test/exprs/function/function_map_inner_product_test.cpp @@ -0,0 +1,429 @@ +// 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 + +#include +#include +#include + +#include "common/exception.h" +#include "common/status.h" +#include "core/assert_cast.h" +#include "core/block/block.h" +#include "core/column/column_array.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/column/column_vector.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_string.h" +#include "exprs/function/function_inner_product.h" +#include "exprs/function/simple_function_factory.h" + +namespace doris { +namespace { + +MutableColumnPtr make_offsets(const std::vector& offsets) { + auto result = ColumnArray::ColumnOffsets::create(); + for (size_t offset : offsets) { + result->insert_value(offset); + } + return result; +} + +template +MutableColumnPtr make_nullable_vector(const std::vector>& values) { + auto nested = ColumnType::create(); + auto null_map = ColumnUInt8::create(); + for (const auto& value : values) { + nested->insert_value(value.value_or(ValueType {})); + null_map->insert_value(value.has_value() ? 0 : 1); + } + return ColumnNullable::create(std::move(nested), std::move(null_map)); +} + +MutableColumnPtr make_nullable_string(const std::vector>& values) { + auto nested = ColumnString::create(); + auto null_map = ColumnUInt8::create(); + for (const auto& value : values) { + const std::string data = value.value_or(""); + nested->insert_data(data.data(), data.size()); + null_map->insert_value(value.has_value() ? 0 : 1); + } + return ColumnNullable::create(std::move(nested), std::move(null_map)); +} + +ColumnPtr make_int_float_map(const std::vector>& keys, + const std::vector>& values, + const std::vector& offsets) { + return ColumnMap::create(make_nullable_vector(keys), + make_nullable_vector(values), make_offsets(offsets)); +} + +ColumnPtr make_high_cardinality_map_with_shadowed_nulls(size_t key_count) { + std::vector> keys; + std::vector> values; + keys.reserve(key_count + 3); + values.reserve(key_count + 3); + for (size_t key = 0; key < key_count; ++key) { + keys.emplace_back(static_cast(key)); + if (key == 5) { + values.emplace_back(std::nullopt); + } else { + values.emplace_back(1.0F); + } + } + keys.emplace_back(5); + values.emplace_back(2.0F); + keys.emplace_back(std::nullopt); + values.emplace_back(std::nullopt); + keys.emplace_back(std::nullopt); + values.emplace_back(4.0F); + return make_int_float_map(keys, values, {key_count + 3}); +} + +ColumnPtr make_largeint_float_map(const std::vector>& keys, + const std::vector>& values, + const std::vector& offsets) { + return ColumnMap::create(make_nullable_vector(keys), + make_nullable_vector(values), make_offsets(offsets)); +} + +ColumnPtr make_string_float_map(const std::vector>& keys, + const std::vector>& values, + const std::vector& offsets) { + return ColumnMap::create(make_nullable_string(keys), + make_nullable_vector(values), make_offsets(offsets)); +} + +Status execute_inner_product(Block& block, const DataTypePtr& return_type) { + ColumnsWithTypeAndName arguments {block.get_by_position(0), block.get_by_position(1)}; + auto function = + SimpleFunctionFactory::instance().get_function("inner_product", arguments, return_type); + if (function == nullptr) { + return Status::InternalError("function inner_product is not registered"); + } + return function->execute(nullptr, block, {0, 1}, 2, block.rows()); +} + +DataTypePtr nullable_int_type() { + return make_nullable(std::make_shared()); +} + +DataTypePtr nullable_float_type() { + return make_nullable(std::make_shared()); +} + +} // namespace + +TEST(FunctionMapInnerProductTest, numeric_keys) { + auto map_type = std::make_shared(nullable_int_type(), nullable_float_type()); + auto return_type = std::make_shared(); + Block block; + block.insert({make_int_float_map({1, 2, 1, 3}, {1.0F, 2.0F, -2.0F, 0.5F}, {2, 4, 4}), map_type, + "left"}); + block.insert({make_int_float_map({2, 1, 3, 2, 1}, {3.0F, 4.0F, 8.0F, 99.0F, 4.0F}, {2, 5, 5}), + map_type, "right"}); + block.insert({nullptr, return_type, "result"}); + + ASSERT_TRUE(execute_inner_product(block, return_type).ok()); + const auto& result = + assert_cast(*block.get_by_position(2).column).get_data(); + ASSERT_EQ(result.size(), 3); + EXPECT_FLOAT_EQ(result[0], 10.0F); + EXPECT_FLOAT_EQ(result[1], -4.0F); + EXPECT_FLOAT_EQ(result[2], 0.0F); +} + +TEST(FunctionMapInnerProductTest, const_map) { + auto map_type = std::make_shared(nullable_int_type(), nullable_float_type()); + auto return_type = std::make_shared(); + Block block; + block.insert({ColumnConst::create(make_int_float_map({1, 2}, {2.0F, 3.0F}, {2}), 2), map_type, + "left"}); + block.insert({make_int_float_map({2, 1}, {4.0F, 5.0F}, {1, 2}), map_type, "right"}); + block.insert({nullptr, return_type, "result"}); + + ASSERT_TRUE(execute_inner_product(block, return_type).ok()); + const auto& result = + assert_cast(*block.get_by_position(2).column).get_data(); + ASSERT_EQ(result.size(), 2); + EXPECT_FLOAT_EQ(result[0], 12.0F); + EXPECT_FLOAT_EQ(result[1], 10.0F); +} + +TEST(FunctionMapInnerProductTest, reuses_high_cardinality_const_map) { + constexpr size_t key_count = 1024; + constexpr size_t row_count = 3; + auto map_type = std::make_shared(nullable_int_type(), nullable_float_type()); + auto return_type = std::make_shared(); + auto constant_map = make_high_cardinality_map_with_shadowed_nulls(key_count); + + std::vector> varying_keys; + std::vector> varying_values; + std::vector varying_offsets; + varying_keys.reserve(row_count * (key_count + 2)); + varying_values.reserve(row_count * (key_count + 2)); + varying_offsets.reserve(row_count); + for (size_t row = 0; row < row_count; ++row) { + for (size_t key = 0; key < key_count; ++key) { + varying_keys.emplace_back(static_cast(key)); + varying_values.emplace_back(1.0F); + } + varying_keys.emplace_back(5); + varying_values.emplace_back(3.0F); + varying_keys.emplace_back(std::nullopt); + varying_values.emplace_back(static_cast(row + 1)); + varying_offsets.emplace_back(varying_keys.size()); + } + auto varying_map = make_int_float_map(varying_keys, varying_values, varying_offsets); + + auto expect_results = [&](bool constant_is_left) { + Block block; + ColumnPtr constant = ColumnConst::create(constant_map, row_count); + block.insert({constant_is_left ? constant : varying_map, map_type, "left"}); + block.insert({constant_is_left ? varying_map : constant, map_type, "right"}); + block.insert({nullptr, return_type, "result"}); + + ASSERT_TRUE(execute_inner_product(block, return_type).ok()); + const auto& result = + assert_cast(*block.get_by_position(2).column).get_data(); + ASSERT_EQ(result.size(), row_count); + EXPECT_FLOAT_EQ(result[0], 1033.0F); + EXPECT_FLOAT_EQ(result[1], 1037.0F); + EXPECT_FLOAT_EQ(result[2], 1041.0F); + }; + + expect_results(true); + expect_results(false); +} + +TEST(FunctionMapInnerProductTest, avoids_caching_oversized_const_map) { + constexpr size_t key_count = 1024; + constexpr size_t row_count = 3; + auto map_type = std::make_shared(nullable_int_type(), nullable_float_type()); + auto return_type = std::make_shared(); + std::vector> constant_keys; + std::vector> constant_values; + constant_keys.reserve(key_count); + constant_values.reserve(key_count); + for (size_t key = 0; key < key_count; ++key) { + constant_keys.emplace_back(static_cast(key)); + constant_values.emplace_back(1.0F); + } + auto constant_map = make_int_float_map(constant_keys, constant_values, {key_count}); + auto varying_map = make_int_float_map({5, 7, -1}, {3.0F, 2.0F, 7.0F}, {1, 2, 3}); + + auto expect_results = [&](bool constant_is_left) { + Block block; + ColumnPtr constant = ColumnConst::create(constant_map, row_count); + block.insert({constant_is_left ? constant : varying_map, map_type, "left"}); + block.insert({constant_is_left ? varying_map : constant, map_type, "right"}); + block.insert({nullptr, return_type, "result"}); + + ASSERT_TRUE(execute_inner_product(block, return_type).ok()); + const auto& result = + assert_cast(*block.get_by_position(2).column).get_data(); + ASSERT_EQ(result.size(), row_count); + EXPECT_FLOAT_EQ(result[0], 3.0F); + EXPECT_FLOAT_EQ(result[1], 2.0F); + EXPECT_FLOAT_EQ(result[2], 0.0F); + }; + + expect_results(true); + expect_results(false); +} + +TEST(FunctionMapInnerProductTest, empty_map_short_circuits_high_cardinality_row) { + constexpr size_t key_count = 4096; + auto map_type = std::make_shared(nullable_int_type(), nullable_float_type()); + auto return_type = std::make_shared(); + std::vector> keys; + std::vector> values; + keys.reserve(key_count); + values.reserve(key_count); + for (size_t key = 0; key < key_count; ++key) { + keys.emplace_back(static_cast(key)); + values.emplace_back(1.0F); + } + + Block block; + block.insert({make_int_float_map(keys, values, {0, key_count}), map_type, "left"}); + block.insert({make_int_float_map(keys, values, {key_count, key_count}), map_type, "right"}); + block.insert({nullptr, return_type, "result"}); + + ASSERT_TRUE(execute_inner_product(block, return_type).ok()); + const auto& result = + assert_cast(*block.get_by_position(2).column).get_data(); + ASSERT_EQ(result.size(), 2); + EXPECT_FLOAT_EQ(result[0], 0.0F); + EXPECT_FLOAT_EQ(result[1], 0.0F); +} + +TEST(FunctionMapInnerProductTest, largeint_keys) { + auto largeint_type = make_nullable(std::make_shared()); + auto map_type = std::make_shared(largeint_type, nullable_float_type()); + auto return_type = std::make_shared(); + Block block; + block.insert({make_largeint_float_map({Int128 {1}, Int128 {2}}, {2.0F, 3.0F}, {2}), map_type, + "left"}); + block.insert({make_largeint_float_map({Int128 {2}, Int128 {3}}, {4.0F, 5.0F}, {2}), map_type, + "right"}); + block.insert({nullptr, return_type, "result"}); + + ASSERT_TRUE(execute_inner_product(block, return_type).ok()); + const auto& result = + assert_cast(*block.get_by_position(2).column).get_data(); + ASSERT_EQ(result.size(), 1); + EXPECT_FLOAT_EQ(result[0], 12.0F); +} + +TEST(FunctionMapInnerProductTest, string_and_null_keys) { + auto nullable_string = make_nullable(std::make_shared()); + auto map_type = std::make_shared(nullable_string, nullable_float_type()); + auto return_type = std::make_shared(); + Block block; + block.insert({make_string_float_map({"a", std::nullopt}, {2.0F, 3.0F}, {2}), map_type, "left"}); + block.insert( + {make_string_float_map({std::nullopt, "a"}, {4.0F, 5.0F}, {2}), map_type, "right"}); + block.insert({nullptr, return_type, "result"}); + + ASSERT_TRUE(execute_inner_product(block, return_type).ok()); + const auto& result = + assert_cast(*block.get_by_position(2).column).get_data(); + ASSERT_EQ(result.size(), 1); + EXPECT_FLOAT_EQ(result[0], 22.0F); +} + +TEST(FunctionMapInnerProductTest, duplicate_keys_use_last_value) { + auto map_type = std::make_shared(nullable_int_type(), nullable_float_type()); + auto return_type = std::make_shared(); + Block block; + // Rows 0 and 2 build the left map; rows 1 and 3 build the right map. + // Rows 0 and 1 use ordinary duplicate keys; rows 2 and 3 use duplicate NULL keys. + block.insert( + {make_int_float_map( + {1, 1, 1, 1, 2, std::nullopt, std::nullopt, std::nullopt, std::nullopt, 2}, + {2.0F, 3.0F, 4.0F, 5.0F, 7.0F, 2.0F, 3.0F, 4.0F, 5.0F, 7.0F}, {2, 5, 7, 10}), + map_type, "left"}); + block.insert( + {make_int_float_map( + {1, 1, 2, 1, 1, std::nullopt, std::nullopt, 2, std::nullopt, std::nullopt}, + {4.0F, 5.0F, 7.0F, 2.0F, 3.0F, 4.0F, 5.0F, 7.0F, 2.0F, 3.0F}, {3, 5, 8, 10}), + map_type, "right"}); + block.insert({nullptr, return_type, "result"}); + + ASSERT_TRUE(execute_inner_product(block, return_type).ok()); + const auto& result = + assert_cast(*block.get_by_position(2).column).get_data(); + ASSERT_EQ(result.size(), 4); + EXPECT_FLOAT_EQ(result[0], 15.0F); + EXPECT_FLOAT_EQ(result[1], 15.0F); + EXPECT_FLOAT_EQ(result[2], 15.0F); + EXPECT_FLOAT_EQ(result[3], 15.0F); +} + +TEST(FunctionMapInnerProductTest, shadowed_null_values_are_ignored) { + auto map_type = std::make_shared(nullable_int_type(), nullable_float_type()); + auto return_type = std::make_shared(); + Block block; + // Cover a shadowed NULL value on either side and in both build/probe roles. Rows 1 and 3 + // additionally exercise the separate NULL-key bucket. + block.insert( + {make_int_float_map( + {1, 1, std::nullopt, std::nullopt, 2, 1, 2, std::nullopt, 2, 3}, + {std::nullopt, 2.0F, std::nullopt, 2.0F, 5.0F, 3.0F, 4.0F, 3.0F, 4.0F, 5.0F}, + {2, 5, 7, 10}), + map_type, "left"}); + block.insert( + {make_int_float_map( + {1, 2, 3, std::nullopt, 3, 1, 1, 3, std::nullopt, std::nullopt}, + {3.0F, 4.0F, 5.0F, 3.0F, 4.0F, std::nullopt, 2.0F, 5.0F, std::nullopt, 2.0F}, + {3, 5, 8, 10}), + map_type, "right"}); + block.insert({nullptr, return_type, "result"}); + + ASSERT_TRUE(execute_inner_product(block, return_type).ok()); + const auto& result = + assert_cast(*block.get_by_position(2).column).get_data(); + ASSERT_EQ(result.size(), 4); + EXPECT_FLOAT_EQ(result[0], 6.0F); + EXPECT_FLOAT_EQ(result[1], 6.0F); + EXPECT_FLOAT_EQ(result[2], 6.0F); + EXPECT_FLOAT_EQ(result[3], 6.0F); +} + +TEST(FunctionMapInnerProductTest, rejects_retained_null_values) { + auto map_type = std::make_shared(nullable_int_type(), nullable_float_type()); + auto return_type = std::make_shared(); + + auto expect_rejected = [&](ColumnPtr left, ColumnPtr right, const std::string& message) { + Block block; + block.insert({std::move(left), map_type, "left"}); + block.insert({std::move(right), map_type, "right"}); + block.insert({nullptr, return_type, "result"}); + + const auto status = execute_inner_product(block, return_type); + ASSERT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find(message), std::string::npos); + }; + + const std::string first_argument_error = + "First argument for function inner_product cannot have null"; + const std::string second_argument_error = + "Second argument for function inner_product cannot have null"; + + // Retained NULL on the left while the left and right maps are selected for build. + expect_rejected(make_int_float_map({1}, {std::nullopt}, {1}), + make_int_float_map({2, 3}, {2.0F, 3.0F}, {2}), first_argument_error); + expect_rejected(make_int_float_map({std::nullopt, std::nullopt}, {1.0F, std::nullopt}, {2}), + make_int_float_map({1}, {3.0F}, {1}), first_argument_error); + + // Retained NULL on the right while the left and right maps are selected for build. + expect_rejected(make_int_float_map({1}, {3.0F}, {1}), + make_int_float_map({std::nullopt, std::nullopt}, {1.0F, std::nullopt}, {2}), + second_argument_error); + expect_rejected(make_int_float_map({2, 3}, {2.0F, 3.0F}, {2}), + make_int_float_map({1}, {std::nullopt}, {1}), second_argument_error); + + // Empty-map short-circuiting must not hide a retained NULL in the nonempty map. + expect_rejected(make_int_float_map({1}, {std::nullopt}, {1}), make_int_float_map({}, {}, {0}), + first_argument_error); + expect_rejected(make_int_float_map({}, {}, {0}), make_int_float_map({1}, {std::nullopt}, {1}), + second_argument_error); +} + +TEST(FunctionMapInnerProductTest, rejects_unsupported_key_type) { + auto double_type = make_nullable(std::make_shared()); + auto map_type = std::make_shared(double_type, nullable_float_type()); + DataTypes arguments {map_type, map_type}; + + try { + FunctionInnerProduct::create()->get_return_type_impl(arguments); + FAIL() << "Expected unsupported map key type to be rejected"; + } catch (const doris::Exception& exception) { + EXPECT_NE(std::string(exception.what()) + .find("inner_product only supports integer or string map keys"), + std::string::npos); + } +} + +} // namespace doris diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/InnerProduct.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/InnerProduct.java index 78780bf931fa94..a5e73abb7db5a0 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/InnerProduct.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/InnerProduct.java @@ -18,13 +18,17 @@ package org.apache.doris.nereids.trees.expressions.functions.scalar; import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.exceptions.AnalysisException; import org.apache.doris.nereids.trees.expressions.Expression; import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable; import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression; import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.DataType; import org.apache.doris.nereids.types.FloatType; +import org.apache.doris.nereids.types.MapType; +import org.apache.doris.nereids.types.coercion.AnyDataType; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; @@ -39,7 +43,10 @@ public class InnerProduct extends ScalarFunction implements ExplicitlyCastableSi public static final List SIGNATURES = ImmutableList.of( FunctionSignature.ret(FloatType.INSTANCE) - .args(ArrayType.of(FloatType.INSTANCE), ArrayType.of(FloatType.INSTANCE)) + .args(ArrayType.of(FloatType.INSTANCE), ArrayType.of(FloatType.INSTANCE)), + FunctionSignature.ret(FloatType.INSTANCE) + .args(MapType.of(new AnyDataType(0), FloatType.INSTANCE), + MapType.of(new AnyDataType(0), FloatType.INSTANCE)) ); /** @@ -54,6 +61,39 @@ private InnerProduct(ScalarFunctionParams functionParams) { super(functionParams); } + @Override + public void checkLegalityBeforeTypeCoercion() { + checkMapKeyTypes(true); + } + + @Override + public void checkLegalityAfterRewrite() { + checkMapKeyTypes(false); + } + + private void checkMapKeyTypes(boolean allowNullKeyType) { + DataType firstKeyType = null; + for (Expression argument : getArguments()) { + DataType argumentType = argument.getDataType(); + if (argumentType.isMapType()) { + DataType keyType = ((MapType) argumentType).getKeyType(); + if (allowNullKeyType && keyType.isNullType()) { + continue; + } + if (!keyType.isIntegralType() && !keyType.isStringLikeType()) { + throw new AnalysisException("inner_product only supports integer or string map keys," + + " but got " + keyType.toSql() + " in expression " + toSql()); + } + if (firstKeyType != null && firstKeyType.isIntegralType() != keyType.isIntegralType()) { + throw new AnalysisException("inner_product requires map keys from the same type family," + + " but got " + firstKeyType.toSql() + " and " + keyType.toSql() + + " in expression " + toSql()); + } + firstKeyType = keyType; + } + } + } + /** * withChildren. */ diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/scalar/InnerProductTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/scalar/InnerProductTest.java new file mode 100644 index 00000000000000..e53c100b295627 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/scalar/InnerProductTest.java @@ -0,0 +1,82 @@ +// 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. + +package org.apache.doris.nereids.trees.expressions.functions.scalar; + +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.parser.NereidsParser; +import org.apache.doris.nereids.rules.expression.ExpressionRewriteTestHelper; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.types.MapType; +import org.apache.doris.qe.GlobalVariable; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class InnerProductTest { + + private static final NereidsParser PARSER = new NereidsParser(); + + @Test + public void testInferNullMapKeyType() { + Expression expression = analyze( + "inner_product(map(null, cast(2 as float))," + + " map(cast(null as int), cast(3 as float)))"); + + Assertions.assertTrue(expression instanceof InnerProduct); + for (Expression argument : ((InnerProduct) expression).getArguments()) { + Assertions.assertTrue(argument.getDataType().isMapType()); + Assertions.assertEquals(IntegerType.INSTANCE, + ((MapType) argument.getDataType()).getKeyType()); + } + Assertions.assertDoesNotThrow(expression::checkLegalityAfterRewrite); + } + + @Test + public void testRejectMixedMapKeyFamiliesInBothCoercionModes() { + boolean originalBehavior = GlobalVariable.enableNewTypeCoercionBehavior; + try { + for (boolean enableNewBehavior : new boolean[] {true, false}) { + GlobalVariable.enableNewTypeCoercionBehavior = enableNewBehavior; + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> analyze("inner_product(map(1, cast(2 as float))," + + " map('1', cast(3 as float)))")); + Assertions.assertTrue(exception.getMessage().contains("same type family"), + exception::getMessage); + } + } finally { + GlobalVariable.enableNewTypeCoercionBehavior = originalBehavior; + } + } + + @Test + public void testRejectAllNullMapKeyTypesAfterCoercion() { + Expression expression = analyze( + "inner_product(map(null, cast(1 as float))," + + " map(null, cast(2 as float)))"); + + AnalysisException exception = Assertions.assertThrows( + AnalysisException.class, expression::checkLegalityAfterRewrite); + Assertions.assertTrue(exception.getMessage().contains( + "only supports integer or string map keys"), exception::getMessage); + } + + private Expression analyze(String sql) { + return ExpressionRewriteTestHelper.typeCoercion(PARSER.parseExpression(sql)); + } +} diff --git a/regression-test/data/query_p0/sql_functions/map_functions/test_map_inner_product.out b/regression-test/data/query_p0/sql_functions/map_functions/test_map_inner_product.out new file mode 100644 index 00000000000000..df73024c626ee7 --- /dev/null +++ b/regression-test/data/query_p0/sql_functions/map_functions/test_map_inner_product.out @@ -0,0 +1,20 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !map_inner_product_rows -- +1 10.0 +2 -4.0 +3 0.0 + +-- !map_inner_product_string_keys -- +22.0 + +-- !map_inner_product_null_key -- +38.0 + +-- !map_inner_product_inferred_null_key -- +6.0 + +-- !map_inner_product_disjoint -- +0.0 + +-- !map_inner_product_dense_compatibility -- +11.0 diff --git a/regression-test/suites/query_p0/sql_functions/map_functions/test_map_inner_product.groovy b/regression-test/suites/query_p0/sql_functions/map_functions/test_map_inner_product.groovy new file mode 100644 index 00000000000000..a7c11284b6c67a --- /dev/null +++ b/regression-test/suites/query_p0/sql_functions/map_functions/test_map_inner_product.groovy @@ -0,0 +1,124 @@ +// 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. + +suite("test_map_inner_product", "p0,nonConcurrent") { + sql "drop table if exists test_map_inner_product" + sql """ + create table test_map_inner_product ( + id int, + lhs map, + rhs map + ) + duplicate key(id) + distributed by hash(id) buckets 1 + properties("replication_num" = "1") + """ + sql """ + insert into test_map_inner_product values + (1, map(1, 1.0, 2, 2.0), map(2, 3.0, 1, 4.0)), + (2, map(1, -2.0, 3, 0.5), map(1, 4.0, 2, 99.0, 3, 8.0)), + (3, cast(map() as map), map(1, 10.0)), + (4, map(1, cast(null as float)), map(1, 2.0)), + (5, cast(null as map), map(1, 2.0)) + """ + + order_qt_map_inner_product_rows """ + select id, inner_product(lhs, rhs) + from test_map_inner_product + where id <= 3 + order by id + """ + + qt_map_inner_product_string_keys """ + select inner_product( + map('a', 2.0, 'b', 3.0), + map('b', 4.0, 'c', 100.0, 'a', 5.0)) + """ + + qt_map_inner_product_null_key """ + select inner_product( + map(cast(null as int), 2.0, 1, 3.0), + map(1, 10.0, cast(null as int), 4.0)) + """ + + qt_map_inner_product_inferred_null_key """ + select inner_product( + map(null, cast(2 as float)), + map(cast(null as int), cast(3 as float))) + """ + + qt_map_inner_product_disjoint """ + select inner_product(map(1, 2.0), map(2, 3.0)) + """ + + qt_map_inner_product_dense_compatibility """ + select inner_product([1.0, 2.0], [3.0, 4.0]) + """ + + test { + sql """ + select inner_product(lhs, rhs) + from test_map_inner_product + where id = 4 + """ + exception "First argument for function inner_product cannot have null" + } + + test { + sql """ + select inner_product(lhs, rhs) + from test_map_inner_product + where id = 5 + """ + exception "First argument for function inner_product cannot be null" + } + + test { + sql """ + select inner_product( + map(cast('2024-01-01' as date), cast(1 as float)), + map(cast('2024-01-01' as date), cast(2 as float))) + """ + exception "inner_product only supports integer or string map keys" + } + + test { + sql """ + select inner_product( + map(null, cast(1 as float)), + map(null, cast(2 as float))) + """ + exception "inner_product only supports integer or string map keys" + } + + def originalTypeCoercionBehavior = sql """ + show global variables like 'enable_new_type_coercion_behavior' + """ + try { + sql "set global enable_new_type_coercion_behavior = false" + test { + sql """ + select inner_product( + map(1, cast(2 as float)), + map('1', cast(3 as float))) + """ + exception "inner_product requires map keys from the same type family" + } + } finally { + sql "set global enable_new_type_coercion_behavior = ${originalTypeCoercionBehavior[0][1]}" + } +}