diff --git a/be/src/exprs/function/function_map.cpp b/be/src/exprs/function/function_map.cpp index bd1b8281904607..ab4e1fba56acf0 100644 --- a/be/src/exprs/function/function_map.cpp +++ b/be/src/exprs/function/function_map.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -181,6 +182,133 @@ class FunctionMapFromArrays : public IFunction { } }; +// (MAP, ARRAY) -> MAP +class FunctionMapFilter : public IFunction { +public: + static constexpr auto name = "map_filter"; + static FunctionPtr create() { return std::make_shared(); } + + String get_name() const override { return name; } + size_t get_number_of_arguments() const override { return 2; } + bool use_default_implementation_for_nulls() const override { return false; } + + DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { + auto map_type = remove_nullable(arguments[0]); + return have_nullable(arguments) ? make_nullable(std::move(map_type)) : map_type; + } + + Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments, + uint32_t result, size_t input_rows_count) const override { + const auto& [unpacked_map_column, map_is_const] = + unpack_if_const(block.get_by_position(arguments[0]).column); + const auto& [unpacked_predicate_column, predicate_is_const] = + unpack_if_const(block.get_by_position(arguments[1]).column); + + auto result_null_map = ColumnUInt8::create(input_rows_count, 0); + auto& result_null_map_data = result_null_map->get_data(); + auto merge_null_map = [&](const ColumnPtr& column, bool is_const) -> const IColumn& { + if (const auto* nullable = check_and_get_column(column.get())) { + VectorizedUtils::update_null_map(result_null_map_data, + nullable->get_null_map_data(), is_const); + return nullable->get_nested_column(); + } + return *column; + }; + + const auto& map = + assert_cast(merge_null_map(unpacked_map_column, map_is_const)); + const auto& predicate = assert_cast( + merge_null_map(unpacked_predicate_column, predicate_is_const)); + RETURN_IF_ERROR(check_arguments(map, map_is_const, predicate, predicate_is_const, + result_null_map_data)); + + IColumn::Selector selector; + auto result_offsets = ColumnArray::ColumnOffsets::create(); + RETURN_IF_ERROR(build_selector_and_offsets(map, map_is_const, predicate, predicate_is_const, + result_null_map_data, selector, + *result_offsets)); + + auto result_keys = map.get_keys().clone_empty(); + auto result_values = map.get_values().clone_empty(); + if (!map_is_const) { + map.get_keys().append_data_by_selector(result_keys, selector); + map.get_values().append_data_by_selector(result_values, selector); + } else if (!selector.empty()) { + result_keys->insert_indices_from(map.get_keys(), selector.data(), + selector.data() + selector.size()); + result_values->insert_indices_from(map.get_values(), selector.data(), + selector.data() + selector.size()); + } + auto result_map = ColumnMap::create(std::move(result_keys), std::move(result_values), + std::move(result_offsets)); + if (block.get_by_position(result).type->is_nullable()) { + block.replace_by_position(result, ColumnNullable::create(std::move(result_map), + std::move(result_null_map))); + } else { + block.replace_by_position(result, std::move(result_map)); + } + return Status::OK(); + } + +private: + static Status build_selector_and_offsets(const ColumnMap& map, bool map_is_const, + const ColumnArray& predicate, bool predicate_is_const, + const NullMap& result_null_map, + IColumn::Selector& selector, + ColumnArray::ColumnOffsets& result_offsets) { + constexpr auto max_selector_position = std::numeric_limits::max(); + if (map.get_keys().size() > 0 && map.get_keys().size() - 1 > max_selector_position) { + return Status::InvalidArgument( + "Function {} cannot process {} nested entries because selector positions are " + "limited to {}", + name, map.get_keys().size(), max_selector_position); + } + const auto& nullable_predicate = assert_cast(predicate.get_data()); + const auto& predicate_null_map = nullable_predicate.get_null_map_data(); + const auto& predicate_values = + assert_cast(nullable_predicate.get_nested_column()).get_data(); + + if (!map_is_const) { + selector.reserve(map.get_keys().size()); + } + result_offsets.reserve(result_null_map.size()); + + for (size_t row = 0; row < result_null_map.size(); ++row) { + if (!result_null_map[row]) { + const size_t map_row = index_check_const(row, map_is_const); + const size_t map_begin = map.get_offsets()[map_row - 1]; + const size_t predicate_begin = + predicate.get_offsets()[index_check_const(row, predicate_is_const) - 1]; + const size_t entry_count = map.size_at(map_row); + for (size_t entry = 0; entry < entry_count; ++entry) { + const size_t predicate_entry = predicate_begin + entry; + const bool selected = predicate_null_map[predicate_entry] == 0 && + predicate_values[predicate_entry] != 0; + if (selected) { + selector.push_back(map_begin + entry); + } + } + } + result_offsets.insert_value(selector.size()); + } + return Status::OK(); + } + + static Status check_arguments(const ColumnMap& map, bool map_is_const, + const ColumnArray& predicate, bool predicate_is_const, + const NullMap& result_null_map) { + for (size_t row = 0; row < result_null_map.size(); ++row) { + if (!result_null_map[row] && + map.size_at(index_check_const(row, map_is_const)) != + predicate.size_at(index_check_const(row, predicate_is_const))) { + return Status::InvalidArgument( + "The map and lambda result offsets of function {} must be identical", name); + } + } + return Status::OK(); + } +}; + // construct a map // map(key1, value2, key2, value2) -> {key1: value2, key2: value2} class FunctionMap : public IFunction { @@ -457,9 +585,10 @@ class FunctionMapEntries : public IFunction { } }; +template class FunctionMapFromEntries : public IFunction { public: - static constexpr auto name = "map_from_entries"; + static constexpr auto name = keys_are_unique ? "%map_from_entries_unique%" : "map_from_entries"; static FunctionPtr create() { return std::make_shared(); } String get_name() const override { return name; } @@ -496,7 +625,9 @@ class FunctionMapFromEntries : public IFunction { auto result_map = ColumnMap::create(make_nullable(entry_struct.get_column_ptr(0)), make_nullable(entry_struct.get_column_ptr(1)), entries.get_offsets_ptr()); - RETURN_IF_ERROR(result_map->deduplicate_keys()); + if constexpr (!keys_are_unique) { + RETURN_IF_ERROR(result_map->deduplicate_keys()); + } if (nullable_array != nullptr) { block.replace_by_position( result, ColumnNullable::create(std::move(result_map), @@ -531,6 +662,108 @@ class FunctionMapFromEntries : public IFunction { } }; +/// Internal conversion used by map-filter rewrites. Each non-null input row is converted from +/// ARRAY>> to MAP by dropping null struct entries. Null outer arrays +/// produce null maps, while null key or value fields inside retained entries remain valid. +/// The input entries come from an existing map, so their keys are already unique and no key +/// deduplication is needed here. +class FunctionMapFromFilteredEntries : public IFunction { +public: + static constexpr auto name = "%map_from_filtered_entries_unique%"; + static FunctionPtr create() { return std::make_shared(); } + + String get_name() const override { return name; } + size_t get_number_of_arguments() const override { return 1; } + bool use_default_implementation_for_nulls() const override { return false; } + + DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { + const auto& array_type = assert_cast(*remove_nullable(arguments[0])); + const auto& struct_type = + assert_cast(*remove_nullable(array_type.get_nested_type())); + DCHECK_EQ(struct_type.get_elements().size(), 2); + auto map_type = std::make_shared(make_nullable(struct_type.get_element(0)), + make_nullable(struct_type.get_element(1))); + return arguments[0]->is_nullable() ? make_nullable(std::move(map_type)) : map_type; + } + + Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments, + uint32_t result, size_t input_rows_count) const override { + ColumnPtr entries_column = block.get_by_position(arguments[0]).column; + const auto* nullable_array = check_and_get_column(entries_column.get()); + if (nullable_array != nullptr) { + entries_column = nullable_array->get_nested_column_ptr(); + } + + const auto& entries = assert_cast(*entries_column); + const auto& nullable_entries = assert_cast(entries.get_data()); + const auto& entry_struct = + assert_cast(nullable_entries.get_nested_column()); + DCHECK_EQ(entry_struct.get_columns().size(), 2); + + ColumnPtr result_keys = entry_struct.get_column_ptr(0); + ColumnPtr result_values = entry_struct.get_column_ptr(1); + ColumnPtr result_offsets = entries.get_offsets_ptr(); + + if (nullable_entries.has_null()) { + IColumn::Selector selector; + auto filtered_offsets = ColumnArray::ColumnOffsets::create(); + RETURN_IF_ERROR(build_selector_and_offsets(entries, nullable_entries, nullable_array, + selector, *filtered_offsets)); + + auto filtered_keys = entry_struct.get_column(0).clone_empty(); + auto filtered_values = entry_struct.get_column(1).clone_empty(); + entry_struct.get_column(0).append_data_by_selector(filtered_keys, selector); + entry_struct.get_column(1).append_data_by_selector(filtered_values, selector); + result_keys = std::move(filtered_keys); + result_values = std::move(filtered_values); + result_offsets = std::move(filtered_offsets); + } + + auto result_map = + ColumnMap::create(make_nullable(result_keys), make_nullable(result_values), + std::move(result_offsets)); + if (nullable_array != nullptr) { + block.replace_by_position( + result, ColumnNullable::create(std::move(result_map), + nullable_array->get_null_map_column_ptr())); + } else { + block.replace_by_position(result, std::move(result_map)); + } + return Status::OK(); + } + +private: + static Status build_selector_and_offsets(const ColumnArray& entries, + const ColumnNullable& nullable_entries, + const ColumnNullable* nullable_array, + IColumn::Selector& selector, + ColumnArray::ColumnOffsets& filtered_offsets) { + constexpr auto max_selector_position = std::numeric_limits::max(); + if (nullable_entries.size() > 0 && nullable_entries.size() - 1 > max_selector_position) { + return Status::InvalidArgument( + "Function {} cannot process {} nested entries because selector positions are " + "limited to {}", + name, nullable_entries.size(), max_selector_position); + } + selector.reserve(nullable_entries.size()); + filtered_offsets.reserve(entries.size()); + + for (size_t row = 0; row < entries.size(); ++row) { + if (nullable_array == nullptr || !nullable_array->is_null_at(row)) { + const size_t begin = row == 0 ? 0 : entries.get_offsets()[row - 1]; + const size_t end = entries.get_offsets()[row]; + for (size_t entry = begin; entry < end; ++entry) { + if (!nullable_entries.is_null_at(entry)) { + selector.push_back(entry); + } + } + } + filtered_offsets.insert_value(selector.size()); + } + return Status::OK(); + } +}; + class FunctionStrToMap : public IFunction { public: static constexpr auto name = "str_to_map"; @@ -985,13 +1218,16 @@ class FunctionDeduplicateMap : public IFunction { void register_function_map(SimpleFunctionFactory& factory) { factory.register_function(); + factory.register_function(); factory.register_function(); factory.register_function>(); factory.register_function>(); factory.register_function>(); factory.register_function>(); factory.register_function(); - 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/lambda_function/varray_map_function.cpp b/be/src/exprs/lambda_function/varray_map_function.cpp index da0722ee0adf21..7538003d3ad2a6 100644 --- a/be/src/exprs/lambda_function/varray_map_function.cpp +++ b/be/src/exprs/lambda_function/varray_map_function.cpp @@ -265,11 +265,17 @@ class ArrayMapFunction : public LambdaFunction { const size_t lambda_batch_rows = _calculate_lambda_batch_size(children[0], lambda_datas, block, required_input_column_ids, has_row_dependent_captures); + // Reuse the nested input columns directly when they fit within eight regular lambda + // batches. Larger inputs use the base batch size, while a smaller byte-budget-derived + // batch remains authoritative. + const size_t lambda_fast_path_rows = lambda_batch_rows == _lambda_block_budget.max_rows + ? _lambda_block_budget.max_rows * 8 + : lambda_batch_rows; // Lambda arguments are already stored contiguously in the input arrays. When all nested - // rows fit in one lambda batch, reuse those columns directly and only materialize captured - // outer columns whose values depend on the outer row. - if (nested_array_column_rows > 0 && nested_array_column_rows <= lambda_batch_rows) { + // rows fit within the direct-execution limit, reuse those columns and only materialize + // captured outer columns whose values depend on the outer row. + if (nested_array_column_rows > 0 && nested_array_column_rows <= lambda_fast_path_rows) { Block lambda_block; PaddedPODArray captured_source_row_indices; MutableColumns captured_columns(lambda_argument_base); diff --git a/be/test/exprs/function/function_map_test.cpp b/be/test/exprs/function/function_map_test.cpp index 22cdd72256fc1e..d61cc4acc2d1b4 100644 --- a/be/test/exprs/function/function_map_test.cpp +++ b/be/test/exprs/function/function_map_test.cpp @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -26,6 +27,7 @@ #include "core/column/column_array.h" #include "core/column/column_const.h" #include "core/column/column_map.h" +#include "core/column/column_nothing.h" #include "core/column/column_nullable.h" #include "core/column/column_struct.h" #include "core/column/column_vector.h" @@ -51,6 +53,16 @@ MutableColumnPtr make_nullable_int_column(const std::vector>& values) { + auto nested = ColumnUInt8::create(); + auto null_map = ColumnUInt8::create(); + for (const auto& value : values) { + nested->insert_value(value.value_or(false)); + null_map->insert_value(value.has_value() ? 0 : 1); + } + return ColumnNullable::create(std::move(nested), std::move(null_map)); +} + MutableColumnPtr make_offsets(const std::vector& offsets) { auto result = ColumnArray::ColumnOffsets::create(); for (size_t offset : offsets) { @@ -59,11 +71,23 @@ MutableColumnPtr make_offsets(const std::vector& offsets) { return result; } +ColumnPtr make_int_map(const std::vector>& keys, + const std::vector>& values, + const std::vector& offsets) { + return ColumnMap::create(make_nullable_int_column(keys), make_nullable_int_column(values), + make_offsets(offsets)); +} + ColumnPtr make_int_array(const std::vector>& values, const std::vector& offsets) { return ColumnArray::create(make_nullable_int_column(values), make_offsets(offsets)); } +ColumnPtr make_bool_array(const std::vector>& values, + const std::vector& offsets) { + return ColumnArray::create(make_nullable_bool_column(values), make_offsets(offsets)); +} + ColumnPtr make_int_entry_array(const std::vector>& keys, const std::vector>& values, const std::vector& offsets, @@ -392,4 +416,237 @@ TEST(FunctionMapTest, map_from_entries_nullable) { EXPECT_FALSE(result.is_null_at(0)); EXPECT_TRUE(result.is_null_at(1)); } + +TEST(FunctionMapTest, map_filter) { + auto nullable_int = make_nullable(std::make_shared()); + auto map_type = std::make_shared(nullable_int, nullable_int); + auto nullable_map_type = make_nullable(map_type); + auto bool_array_type = + std::make_shared(make_nullable(std::make_shared())); + + { + Block block; + block.insert({make_int_map({1, 2, 3}, {10, 20, 30}, {2, 3}), map_type, "map"}); + block.insert({make_bool_array({true, std::nullopt, false}, {2, 3}), bool_array_type, + "predicate"}); + block.insert({nullptr, map_type, "result"}); + + ASSERT_TRUE(execute_map_function("map_filter", block, {0, 1}, 2, map_type).ok()); + const auto& result = assert_cast(*block.get_by_position(2).column); + ASSERT_EQ(result.get_offsets()[0], 1); + ASSERT_EQ(result.get_offsets()[1], 1); + ASSERT_EQ(get_nullable_int(result.get_keys(), 0), 1); + ASSERT_EQ(get_nullable_int(result.get_values(), 0), 10); + } + + { + Block block; + block.insert( + {ColumnConst::create(make_int_map({1, 2}, {10, 20}, {2}), 3), map_type, "map"}); + block.insert({make_bool_array({true, true, true, false, false, true}, {2, 4, 6}), + bool_array_type, "predicate"}); + block.insert({nullptr, map_type, "result"}); + + ASSERT_TRUE(execute_map_function("map_filter", block, {0, 1}, 2, map_type).ok()); + const auto& result = assert_cast(*block.get_by_position(2).column); + ASSERT_EQ(result.get_offsets()[0], 2); + ASSERT_EQ(result.get_offsets()[1], 3); + ASSERT_EQ(result.get_offsets()[2], 4); + EXPECT_EQ(get_nullable_int(result.get_keys(), 2), 1); + EXPECT_EQ(get_nullable_int(result.get_keys(), 3), 2); + } + + { + auto map_null_map = ColumnUInt8::create(); + map_null_map->insert_value(0); + map_null_map->insert_value(1); + Block block; + block.insert({ColumnNullable::create(make_int_map({1, 2}, {10, 20}, {1, 2}), + std::move(map_null_map)), + nullable_map_type, "map"}); + block.insert({make_bool_array({true, true}, {1, 2}), bool_array_type, "predicate"}); + block.insert({nullptr, nullable_map_type, "result"}); + + ASSERT_TRUE(execute_map_function("map_filter", block, {0, 1}, 2, nullable_map_type).ok()); + const auto& result = assert_cast(*block.get_by_position(2).column); + EXPECT_FALSE(result.is_null_at(0)); + EXPECT_TRUE(result.is_null_at(1)); + const auto& nested = assert_cast(result.get_nested_column()); + EXPECT_EQ(nested.get_offsets()[0], 1); + EXPECT_EQ(nested.get_offsets()[1], 1); + } + + { + Block block; + block.insert({make_int_map({}, {}, {0, 0}), map_type, "map"}); + block.insert({make_bool_array({}, {0, 0}), bool_array_type, "predicate"}); + block.insert({nullptr, map_type, "result"}); + + ASSERT_TRUE(execute_map_function("map_filter", block, {0, 1}, 2, map_type).ok()); + const auto& result = assert_cast(*block.get_by_position(2).column); + EXPECT_EQ(result.get_keys().size(), 0); + EXPECT_EQ(result.get_offsets()[0], 0); + EXPECT_EQ(result.get_offsets()[1], 0); + } + + { + Block block; + block.insert({make_int_map({1, 2, 3, 4}, {10, 20, 30, 40}, {2, 4}), map_type, "map"}); + block.insert({ColumnConst::create(make_bool_array({true, false}, {2}), 2), bool_array_type, + "predicate"}); + block.insert({nullptr, map_type, "result"}); + + ASSERT_TRUE(execute_map_function("map_filter", block, {0, 1}, 2, map_type).ok()); + const auto& result = assert_cast(*block.get_by_position(2).column); + ASSERT_EQ(result.get_offsets()[0], 1); + ASSERT_EQ(result.get_offsets()[1], 2); + EXPECT_EQ(get_nullable_int(result.get_keys(), 0), 1); + EXPECT_EQ(get_nullable_int(result.get_keys(), 1), 3); + } + + { + Block block; + block.insert({make_int_map({1, 2}, {10, 20}, {1, 2}), map_type, "map"}); + block.insert({make_bool_array({true}, {1, 1}), bool_array_type, "predicate"}); + block.insert({nullptr, map_type, "result"}); + + auto status = execute_map_function("map_filter", block, {0, 1}, 2, map_type); + ASSERT_TRUE(status.is()) << status.to_string(); + EXPECT_NE(status.to_string().find("The map and lambda result offsets of function " + "map_filter must be identical"), + std::string::npos); + } + + { + auto nullable_bool_array_type = make_nullable(bool_array_type); + auto predicate_null_map = ColumnUInt8::create(); + predicate_null_map->insert_value(0); + predicate_null_map->insert_value(1); + + Block block; + block.insert({make_int_map({1, 2, 3}, {10, 20, 30}, {1, 3}), map_type, "map"}); + block.insert({ColumnNullable::create(make_bool_array({true, true}, {1, 2}), + std::move(predicate_null_map)), + nullable_bool_array_type, "predicate"}); + block.insert({nullptr, nullable_map_type, "result"}); + + ASSERT_TRUE(execute_map_function("map_filter", block, {0, 1}, 2, nullable_map_type).ok()); + const auto& result = assert_cast(*block.get_by_position(2).column); + EXPECT_FALSE(result.is_null_at(0)); + EXPECT_TRUE(result.is_null_at(1)); + const auto& nested = assert_cast(result.get_nested_column()); + EXPECT_EQ(nested.get_offsets()[0], 1); + EXPECT_EQ(nested.get_offsets()[1], 1); + ASSERT_EQ(nested.get_keys().size(), 1); + EXPECT_EQ(get_nullable_int(nested.get_keys(), 0), 1); + } +} + +TEST(FunctionMapTest, map_from_entries_unique_skips_deduplication) { + auto nullable_int = make_nullable(std::make_shared()); + auto entry_struct_type = std::make_shared( + DataTypes {nullable_int, nullable_int}, Strings {"key", "value"}); + auto entry_array_type = std::make_shared(make_nullable(entry_struct_type)); + auto map_type = std::make_shared(nullable_int, nullable_int); + + Block block; + block.insert({make_int_entry_array({1, 1}, {10, 20}, {2}), entry_array_type, "entries"}); + block.insert({nullptr, map_type, "result"}); + + ASSERT_TRUE(execute_map_function("%map_from_entries_unique%", block, {0}, 1, map_type).ok()); + const auto& result = assert_cast(*block.get_by_position(1).column); + EXPECT_EQ(result.get_keys().size(), 2); + EXPECT_EQ(result.get_values().size(), 2); +} + +TEST(FunctionMapTest, map_filter_rejects_unrepresentable_selector_position) { + if (std::numeric_limits::max() <= std::numeric_limits::max()) { + GTEST_SKIP() << "size_t cannot represent positions beyond IColumn::Selector"; + } + + const size_t entry_count = static_cast( + static_cast(std::numeric_limits::max()) + 2); + auto offsets = ColumnArray::ColumnOffsets::create(); + offsets->insert_value(entry_count); + auto predicate_offsets = ColumnArray::ColumnOffsets::create(); + predicate_offsets->insert_value(entry_count); + + auto nullable_int = make_nullable(std::make_shared()); + auto map_type = std::make_shared(nullable_int, nullable_int); + auto bool_array_type = + std::make_shared(make_nullable(std::make_shared())); + + Block block; + block.insert({ColumnMap::create(ColumnNothing::create(entry_count), + ColumnNothing::create(entry_count), std::move(offsets)), + map_type, "map"}); + block.insert( + {ColumnArray::create(ColumnNothing::create(entry_count), std::move(predicate_offsets)), + bool_array_type, "predicate"}); + block.insert({nullptr, map_type, "result"}); + + auto status = execute_map_function("map_filter", block, {0, 1}, 2, map_type); + ASSERT_TRUE(status.is()) << status.to_string(); + EXPECT_NE(status.to_string().find("selector positions are limited to 4294967295"), + std::string::npos); +} + +TEST(FunctionMapTest, map_from_filtered_entries_unique) { + auto nullable_int = make_nullable(std::make_shared()); + auto entry_struct_type = std::make_shared( + DataTypes {nullable_int, nullable_int}, Strings {"key", "value"}); + auto entry_array_type = std::make_shared(make_nullable(entry_struct_type)); + auto map_type = std::make_shared(nullable_int, nullable_int); + + { + Block block; + block.insert({make_int_entry_array({1, 2, std::nullopt, 4}, {10, 20, 30, std::nullopt}, + {2, 4}, {false, true, false, false}), + entry_array_type, "entries"}); + block.insert({nullptr, map_type, "result"}); + + ASSERT_TRUE( + execute_map_function("%map_from_filtered_entries_unique%", block, {0}, 1, map_type) + .ok()); + const auto& result = assert_cast(*block.get_by_position(1).column); + ASSERT_EQ(result.get_offsets()[0], 1); + ASSERT_EQ(result.get_offsets()[1], 3); + ASSERT_EQ(result.get_keys().size(), 3); + EXPECT_EQ(get_nullable_int(result.get_keys(), 0), 1); + EXPECT_TRUE(assert_cast(result.get_keys()).is_null_at(1)); + EXPECT_EQ(get_nullable_int(result.get_values(), 1), 30); + EXPECT_TRUE(assert_cast(result.get_values()).is_null_at(2)); + } + + { + auto outer_null_map = ColumnUInt8::create(); + outer_null_map->insert_value(0); + outer_null_map->insert_value(1); + outer_null_map->insert_value(0); + auto nullable_entry_array_type = make_nullable(entry_array_type); + auto nullable_map_type = make_nullable(map_type); + + Block block; + block.insert({ColumnNullable::create( + make_int_entry_array({1, 98, 99, 4}, {10, 980, 990, 40}, {1, 3, 4}, + {false, false, true, false}), + std::move(outer_null_map)), + nullable_entry_array_type, "entries"}); + block.insert({nullptr, nullable_map_type, "result"}); + + ASSERT_TRUE(execute_map_function("%map_from_filtered_entries_unique%", block, {0}, 1, + nullable_map_type) + .ok()); + const auto& result = assert_cast(*block.get_by_position(1).column); + EXPECT_FALSE(result.is_null_at(0)); + EXPECT_TRUE(result.is_null_at(1)); + EXPECT_FALSE(result.is_null_at(2)); + const auto& result_map = assert_cast(result.get_nested_column()); + ASSERT_EQ(result_map.get_offsets()[0], 1); + ASSERT_EQ(result_map.get_offsets()[1], 1); + ASSERT_EQ(result_map.get_offsets()[2], 2); + EXPECT_EQ(get_nullable_int(result_map.get_keys(), 0), 1); + EXPECT_EQ(get_nullable_int(result_map.get_keys(), 1), 4); + } +} } // namespace doris diff --git a/be/test/exprs/lambda_function/array_map_function_test.cpp b/be/test/exprs/lambda_function/array_map_function_test.cpp index 5e0e5da07817fa..da4e345274f5a1 100644 --- a/be/test/exprs/lambda_function/array_map_function_test.cpp +++ b/be/test/exprs/lambda_function/array_map_function_test.cpp @@ -773,7 +773,7 @@ TEST(ArrayMapFunctionTest, LargeLambdaProducesCorrectResult) { EXPECT_EQ(values.get_element(99999), 100000); } -TEST(ArrayMapFunctionTest, VariableLengthCaptureUsesOuterBatchSize) { +TEST(ArrayMapFunctionTest, VariableLengthCaptureUsesLargerBatch) { auto int_type = std::make_shared(); auto string_type = std::make_shared(); auto array_int_type = std::make_shared(int_type); @@ -800,14 +800,12 @@ TEST(ArrayMapFunctionTest, VariableLengthCaptureUsesOuterBatchSize) { ColumnPtr result; auto status = root->execute_column(&context, &block, nullptr, block.rows(), result); ASSERT_TRUE(status.ok()) << status.to_string(); - ASSERT_EQ(observed_batch_sizes.size(), 3); - EXPECT_EQ(observed_batch_sizes[0], 2); - EXPECT_EQ(observed_batch_sizes[1], 2); - EXPECT_EQ(observed_batch_sizes[2], 1); + ASSERT_EQ(observed_batch_sizes.size(), 1); + EXPECT_EQ(observed_batch_sizes[0], 5); EXPECT_EQ(get_int_array_values(result).size(), 5); } -TEST(ArrayMapFunctionTest, ComplexLambdaUsesOuterBatchSize) { +TEST(ArrayMapFunctionTest, FixedLengthComplexLambdaUsesLargerBatch) { constexpr size_t intermediate_count = 100; constexpr size_t nested_count = 1000; constexpr int outer_batch_size = 256; @@ -840,14 +838,8 @@ TEST(ArrayMapFunctionTest, ComplexLambdaUsesOuterBatchSize) { auto status = root->execute_column(&context, &block, nullptr, 1, result); ASSERT_TRUE(status.ok()) << status.to_string(); - size_t observed_rows = 0; - for (size_t batch_rows : observed_batch_sizes) { - EXPECT_EQ(batch_rows, - std::min(static_cast(outer_batch_size), nested_count - observed_rows)); - observed_rows += batch_rows; - } - EXPECT_EQ(observed_rows, nested_count); - EXPECT_EQ(observed_batch_sizes.size(), 4); + ASSERT_EQ(observed_batch_sizes.size(), 1); + EXPECT_EQ(observed_batch_sizes[0], nested_count); const auto& values = get_int_array_values(result); ASSERT_EQ(values.size(), nested_count); @@ -855,7 +847,7 @@ TEST(ArrayMapFunctionTest, ComplexLambdaUsesOuterBatchSize) { EXPECT_EQ(values.get_element(nested_count - 1), nested_count); } -TEST(ArrayMapFunctionTest, FixedLengthInputsUseOuterBatchSize) { +TEST(ArrayMapFunctionTest, FixedLengthInputsUseLargerBatch) { constexpr size_t capture_count = 64; constexpr size_t nested_count = 1000; constexpr int outer_batch_size = 128; @@ -891,14 +883,8 @@ TEST(ArrayMapFunctionTest, FixedLengthInputsUseOuterBatchSize) { auto status = root->execute_column(&context, &block, nullptr, block.rows(), result); ASSERT_TRUE(status.ok()) << status.to_string(); - size_t observed_rows = 0; - for (size_t batch_rows : observed_batch_sizes) { - EXPECT_EQ(batch_rows, - std::min(static_cast(outer_batch_size), nested_count - observed_rows)); - observed_rows += batch_rows; - } - EXPECT_EQ(observed_rows, nested_count); - EXPECT_EQ(observed_batch_sizes.size(), 8); + ASSERT_EQ(observed_batch_sizes.size(), 1); + EXPECT_EQ(observed_batch_sizes[0], nested_count); EXPECT_EQ(get_int_array_values(result).size(), nested_count); } @@ -957,7 +943,8 @@ TEST(ArrayMapFunctionTest, FixedLengthInputsUsePreferredBlockSizeBudget) { TEST(ArrayMapFunctionTest, MultiBatchPreservesCaptureMappingAcrossSelectedArrayRows) { constexpr size_t selected_array_size = 300; - constexpr int outer_batch_size = 400; + constexpr int outer_batch_size = 64; + constexpr int lambda_batch_size = outer_batch_size; auto int_type = std::make_shared(); auto array_int_type = std::make_shared(int_type); std::vector observed_batch_sizes; @@ -1001,9 +988,15 @@ TEST(ArrayMapFunctionTest, MultiBatchPreservesCaptureMappingAcrossSelectedArrayR ASSERT_TRUE(status.ok()) << status.to_string(); const size_t total_nested_rows = 2 * selected_array_size; - ASSERT_EQ(observed_batch_sizes.size(), 2); - EXPECT_EQ(observed_batch_sizes[0], outer_batch_size); - EXPECT_EQ(observed_batch_sizes[1], total_nested_rows - outer_batch_size); + const size_t expected_full_batches = total_nested_rows / lambda_batch_size; + const size_t expected_last_batch_size = total_nested_rows % lambda_batch_size; + ASSERT_EQ(observed_batch_sizes.size(), expected_full_batches + (expected_last_batch_size > 0)); + for (size_t i = 0; i < expected_full_batches; ++i) { + EXPECT_EQ(observed_batch_sizes[i], lambda_batch_size); + } + if (expected_last_batch_size > 0) { + EXPECT_EQ(observed_batch_sizes.back(), expected_last_batch_size); + } const auto& result_array = assert_cast(*result); ASSERT_EQ(result_array.size(), 3); @@ -1016,10 +1009,9 @@ TEST(ArrayMapFunctionTest, MultiBatchPreservesCaptureMappingAcrossSelectedArrayR EXPECT_EQ(values.get_element(0), 10); EXPECT_EQ(values.get_element(selected_array_size - 1), 309); EXPECT_EQ(values.get_element(selected_array_size), 1050); - EXPECT_EQ(values.get_element(outer_batch_size - 1), - 50 + 1000 + static_cast(outer_batch_size - selected_array_size - 1)); - EXPECT_EQ(values.get_element(outer_batch_size), - 50 + 1000 + static_cast(outer_batch_size - selected_array_size)); + EXPECT_EQ(values.get_element(lambda_batch_size - 1), + 10 + static_cast(lambda_batch_size - 1)); + EXPECT_EQ(values.get_element(lambda_batch_size), 10 + lambda_batch_size); EXPECT_EQ(values.get_element(total_nested_rows - 1), 1349); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java index 62369888db9e23..9732cab755b6e8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java @@ -331,10 +331,14 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.MakeDate; import org.apache.doris.nereids.trees.expressions.functions.scalar.MakeSet; import org.apache.doris.nereids.trees.expressions.functions.scalar.MakeTime; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapAll; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapApply; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapContainsEntry; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapContainsKey; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapContainsValue; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapEntries; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapExists; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapFilter; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapFromArrays; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapFromEntries; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapKeys; @@ -542,6 +546,8 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.ToSeconds; import org.apache.doris.nereids.trees.expressions.functions.scalar.Tokenize; import org.apache.doris.nereids.trees.expressions.functions.scalar.TopLevelDomain; +import org.apache.doris.nereids.trees.expressions.functions.scalar.TransformKeys; +import org.apache.doris.nereids.trees.expressions.functions.scalar.TransformValues; import org.apache.doris.nereids.trees.expressions.functions.scalar.Translate; import org.apache.doris.nereids.trees.expressions.functions.scalar.Trim; import org.apache.doris.nereids.trees.expressions.functions.scalar.TrimIn; @@ -915,10 +921,14 @@ public class BuiltinScalarFunctions implements FunctionHelper { scalar(MakeDate.class, "makedate"), scalar(MakeSet.class, "make_set"), scalar(MakeTime.class, "maketime"), + scalar(MapAll.class, "map_all"), + scalar(MapApply.class, "map_apply"), scalar(MapContainsEntry.class, "map_contains_entry"), scalar(MapContainsKey.class, "map_contains_key"), scalar(MapContainsValue.class, "map_contains_value"), scalar(MapEntries.class, "map_entries"), + scalar(MapExists.class, "map_exists"), + scalar(MapFilter.class, "map_filter"), scalar(MapFromArrays.class, "map_from_arrays"), scalar(MapFromEntries.class, "map_from_entries"), scalar(MapKeys.class, "map_keys"), @@ -1134,6 +1144,8 @@ public class BuiltinScalarFunctions implements FunctionHelper { scalar(TopLevelDomain.class, "top_level_domain"), scalar(ToQuantileState.class, "to_quantile_state"), scalar(ToSeconds.class, "to_seconds"), + scalar(TransformKeys.class, "transform_keys"), + scalar(TransformValues.class, "transform_values"), scalar(Translate.class, "translate"), scalar(Trim.class, "trim"), scalar(TrimIn.class, "trim_in"), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java index 7a3d6495204297..b69dc292475547 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java @@ -763,8 +763,6 @@ public class Rewriter extends AbstractBatchJobExecutor { new MergePercentileToArray()) ), topic("add projection for volatile expression", - // separate AddProjectForVolatileExpression and MergeProjectable - // to avoid dead loop if code has bug topDown(new AddProjectForVolatileExpression()), topDown(new MergeProjectable()) ), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java index 1fc86341310153..7d812ee56a6ae9 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java @@ -3113,7 +3113,11 @@ public Expression visitLambdaExpression(LambdaExpressionContext ctx) { ImmutableList args = ctx.args.stream() .map(RuleContext::getText) .collect(ImmutableList.toImmutableList()); - Expression body = (Expression) visit(ctx.body); + // A tuple Lambda body such as (k, v) -> (k * 3, v + 1) is shorthand for + // (k, v) -> struct(k * 3, v + 1) + Expression body = ctx.body == null + ? new UnboundFunction("struct", visit(ctx.bodyItems, Expression.class)) + : (Expression) visit(ctx.body); return new Lambda(args, body); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java index a8f9d7bb27bf42..184e3991a85f7f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java @@ -41,6 +41,7 @@ import org.apache.doris.nereids.trees.expressions.Alias; import org.apache.doris.nereids.trees.expressions.And; import org.apache.doris.nereids.trees.expressions.ArrayItemReference; +import org.apache.doris.nereids.trees.expressions.ArrayItemReference.ArrayItemSlot; import org.apache.doris.nereids.trees.expressions.Between; import org.apache.doris.nereids.trees.expressions.BinaryArithmetic; import org.apache.doris.nereids.trees.expressions.BitNot; @@ -68,6 +69,7 @@ import org.apache.doris.nereids.trees.expressions.Placeholder; import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; import org.apache.doris.nereids.trees.expressions.Variable; import org.apache.doris.nereids.trees.expressions.WhenClause; import org.apache.doris.nereids.trees.expressions.WindowExpression; @@ -80,9 +82,11 @@ import org.apache.doris.nereids.trees.expressions.functions.agg.SupportMultiDistinct; import org.apache.doris.nereids.trees.expressions.functions.scalar.ElementAt; import org.apache.doris.nereids.trees.expressions.functions.scalar.Lambda; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapEntries; import org.apache.doris.nereids.trees.expressions.functions.udf.AliasUdfBuilder; import org.apache.doris.nereids.trees.expressions.functions.udf.UdfBuilder; import org.apache.doris.nereids.trees.expressions.literal.IntegerLikeLiteral; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; import org.apache.doris.nereids.trees.expressions.literal.Literal; import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; import org.apache.doris.nereids.trees.expressions.literal.StringLikeLiteral; @@ -96,6 +100,7 @@ import org.apache.doris.nereids.types.BigIntType; import org.apache.doris.nereids.types.BooleanType; import org.apache.doris.nereids.types.DataType; +import org.apache.doris.nereids.types.MapType; import org.apache.doris.nereids.types.NestedColumnPrunable; import org.apache.doris.nereids.types.StringType; import org.apache.doris.nereids.types.StructField; @@ -117,12 +122,16 @@ import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableList.Builder; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Lists; +import com.google.common.collect.Sets; import org.apache.commons.lang3.StringUtils; import java.util.ArrayList; import java.util.List; +import java.util.Locale; import java.util.Optional; +import java.util.Set; import java.util.stream.Collectors; import javax.annotation.Nullable; @@ -144,6 +153,14 @@ protected Expression processCompoundNewChildren(CompoundPredicate cp, List MAP_ENTRY_LAMBDA_FUNCTIONS = ImmutableSet.of( + "map_all", + "map_apply", + "map_exists", + "map_filter", + "transform_keys", + "transform_values"); + private final Plan currentPlan; /* bounded={table.a, a} @@ -438,12 +455,93 @@ private UnboundFunction processHighOrderFunction(UnboundFunction unboundFunction // bindLambdaFunction Lambda lambda = (Lambda) unboundFunction.children().get(0); Expression lambdaFunction = lambda.getLambdaFunction(); + if (MAP_ENTRY_LAMBDA_FUNCTIONS.contains(unboundFunction.getName().toLowerCase(Locale.ROOT))) { + return bindMapLambdaFunction(unboundFunction, lambda, lambdaFunction, subChildren, context); + } List arrayItemReferences = lambda.makeArguments(unboundFunction.getName(), subChildren); List boundedSlots = arrayItemReferences.stream() .map(ArrayItemReference::toSlot) .collect(ImmutableList.toImmutableList()); + lambdaFunction = analyzeLambdaFunction(lambda, lambdaFunction, boundedSlots, context); + + Lambda lambdaClosure = lambda.withLambdaFunctionArguments(lambdaFunction, arrayItemReferences); + + // We don't add the ArrayExpression in high order function at all + return unboundFunction.withChildren(ImmutableList.of(lambdaClosure)); + } + + private UnboundFunction bindMapLambdaFunction(UnboundFunction unboundFunction, Lambda lambda, + Expression lambdaFunction, List lambdaInputs, ExpressionRewriteContext context) { + String functionName = unboundFunction.getName(); + if (lambdaInputs.size() != 1) { + throw new AnalysisException(String.format( + "%s requires exactly one map argument but has %d", functionName, lambdaInputs.size())); + } + if (lambda.getLambdaArgumentNames().size() != 2) { + throw new AnalysisException(String.format( + "lambda of %s requires exactly two arguments but has %d", + functionName, lambda.getLambdaArgumentNames().size())); + } + + Expression mapExpression = lambdaInputs.get(0); + if (!(mapExpression.getDataType() instanceof MapType)) { + throw new AnalysisException(String.format( + "the non-lambda argument of %s must be map but is %s", + functionName, mapExpression.getDataType().toSql())); + } + + MapType mapType = (MapType) mapExpression.getDataType(); + ExprId keyExprId = StatementScopeIdGenerator.newExprId(); + ExprId valueExprId = StatementScopeIdGenerator.newExprId(); + ArrayItemSlot keySlot = new ArrayItemSlot( + keyExprId, lambda.getLambdaArgumentName(0), mapType.getKeyType(), true); + ArrayItemSlot valueSlot = new ArrayItemSlot( + valueExprId, lambda.getLambdaArgumentName(1), mapType.getValueType(), true); + lambdaFunction = analyzeLambdaFunction( + lambda, lambdaFunction, ImmutableList.of(keySlot, valueSlot), context); + + Set occupiedNames = Sets.newHashSet(lambda.getLambdaArgumentNames()); + for (Slot slot : lambdaFunction.collect(expression -> expression instanceof Slot)) { + occupiedNames.add(slot.getName()); + } + for (Lambda nestedLambda : lambdaFunction.collect(expression -> expression instanceof Lambda)) { + occupiedNames.addAll(nestedLambda.getLambdaArgumentNames()); + } + + ExprId entryExprId; + String entryName; + do { + entryExprId = StatementScopeIdGenerator.newExprId(); + entryName = "$_map_entry_" + entryExprId.asInt() + "_$"; + } while (occupiedNames.contains(entryName)); + + ArrayItemReference entryArgument = new ArrayItemReference( + entryExprId, entryName, new MapEntries(mapExpression)); + Slot entrySlot = entryArgument.toSlot(); + Expression key = new ElementAt(entrySlot, new IntegerLiteral(1)); + Expression value = new ElementAt(entrySlot, new IntegerLiteral(2)); + Expression rewrittenLambdaFunction = lambdaFunction.rewriteDownShortCircuit(expression -> { + if (expression instanceof ArrayItemSlot) { + ExprId exprId = ((ArrayItemSlot) expression).getExprId(); + if (exprId.equals(keyExprId)) { + return key; + } + if (exprId.equals(valueExprId)) { + return value; + } + } + return expression; + }); + + Lambda lambdaClosure = new Lambda( + ImmutableList.of(entryName), rewrittenLambdaFunction, ImmutableList.of(entryArgument)); + return unboundFunction.withChildren(ImmutableList.of(lambdaClosure)); + } + + private Expression analyzeLambdaFunction(Lambda lambda, Expression lambdaFunction, + List boundedSlots, ExpressionRewriteContext context) { ExpressionAnalyzer lambdaAnalyzer = new ExpressionAnalyzer(currentPlan, new Scope(Optional.of(getScope()), boundedSlots), context == null ? null : context.cascadesContext, true, true) { @@ -454,12 +552,7 @@ protected void couldNotFoundColumn(UnboundSlot unboundSlot, String tableName) { + " in lambda arguments" + lambda.getLambdaArgumentNames()); } }; - lambdaFunction = lambdaAnalyzer.analyze(lambdaFunction, context); - - Lambda lambdaClosure = lambda.withLambdaFunctionArguments(lambdaFunction, arrayItemReferences); - - // We don't add the ArrayExpression in high order function at all - return unboundFunction.withChildren(ImmutableList.of(lambdaClosure)); + return lambdaAnalyzer.analyze(lambdaFunction, context); } UnboundFunction preProcessUnboundFunction(UnboundFunction unboundFunction, ExpressionRewriteContext context) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AccessPathExpressionCollector.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AccessPathExpressionCollector.java index fdbc88de615072..316a0f76833d55 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AccessPathExpressionCollector.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AccessPathExpressionCollector.java @@ -52,6 +52,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.MapContainsEntry; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapContainsKey; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapContainsValue; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapEntries; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapKeys; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapSize; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapValues; @@ -405,6 +406,26 @@ public Void visitMapValues(MapValues mapValues, CollectorContext context) { return continueCollectAccessPath(mapValues.getArgument(0), context); } + @Override + public Void visitMapEntries(MapEntries mapEntries, CollectorContext context) { + LinkedList path = context.accessPathBuilder.accessPath; + if (path.size() >= 2 && AccessPathInfo.ACCESS_ALL.equals(path.get(0))) { + String entryField = path.get(1); + if ("key".equalsIgnoreCase(entryField) || "value".equalsIgnoreCase(entryField)) { + CollectorContext mapContext = new CollectorContext( + context.statementContext, context.bottomFilter); + mapContext.accessPathBuilder.accessPath.addAll(path.subList(2, path.size())); + mapContext.accessPathBuilder.addPrefix("key".equalsIgnoreCase(entryField) + ? AccessPathInfo.ACCESS_MAP_KEYS : AccessPathInfo.ACCESS_MAP_VALUES); + return continueCollectAccessPath(mapEntries.getArgument(0), mapContext); + } + } + if (path.isEmpty()) { + context.accessPathBuilder.addPrefix(AccessPathInfo.ACCESS_ALL); + } + return continueCollectAccessPath(mapEntries.getArgument(0), context); + } + private static boolean isUnderIsNull(List suffixPath) { return suffixPath.size() == 1 && AccessPathInfo.ACCESS_NULL.equals(suffixPath.get(0)); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Lambda.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Lambda.java index 192f04e563227b..09da4522b3cde6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Lambda.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Lambda.java @@ -38,7 +38,6 @@ * After bind, x -> x : arguments("x") -> children: Expression(x) ArrayItemReference(x) */ public class Lambda extends Expression { - private final List argumentNames; /** @@ -61,15 +60,15 @@ public Lambda(List argumentNames, List children) { /** * make slot according array expression * @param functionName function name - * @param arrays array expression + * @param lambdaArgs array or map expression * @return item slots of array expression */ - public ImmutableList makeArguments(String functionName, List arrays) { + public ImmutableList makeArguments(String functionName, List lambdaArgs) { Builder builder = new ImmutableList.Builder<>(); - if (arrays.size() != argumentNames.size()) { + if (lambdaArgs.size() != argumentNames.size()) { // In the lambda expression of array_sort, x and y point to the same slot. - if (functionName.equalsIgnoreCase("array_sort") && arrays.size() == 1 && argumentNames.size() == 2) { - Expression array = arrays.get(0); + if (functionName.equalsIgnoreCase("array_sort") && lambdaArgs.size() == 1 && argumentNames.size() == 2) { + Expression array = lambdaArgs.get(0); if (!(array.getDataType() instanceof ArrayType)) { throw new AnalysisException(String.format("lambda argument must be array but is %s", array)); } @@ -80,8 +79,8 @@ public ImmutableList makeArguments(String functionName, List throw new AnalysisException(String.format("lambda %s arguments' size is not equal parameters' size", toSql())); } - for (int i = 0; i < arrays.size(); i++) { - Expression array = arrays.get(i); + for (int i = 0; i < lambdaArgs.size(); i++) { + Expression array = lambdaArgs.get(i); if (!(array.getDataType() instanceof ArrayType)) { throw new AnalysisException(String.format("lambda argument must be array but is %s", array)); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapAll.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapAll.java new file mode 100644 index 00000000000000..c703fad6e7dddc --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapAll.java @@ -0,0 +1,90 @@ +// 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.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable; +import org.apache.doris.nereids.trees.expressions.functions.RewriteWhenAnalyze; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.BooleanType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * Scalar function map_all. + * + *

The Map lambda first produces one Boolean per entry, then ArrayMatchAll checks the result: + * + * 

+ * map_all((mapKey, mapValue) -> predicate, inputMap)
+ *   ->
+ * array_match_all(
+ *   array_map(
+ *     entry -> predicate(entry[1], entry[2]),
+ *     map_entries(inputMap)))
+ *
+ */ +public class MapAll extends ScalarFunction + implements HighOrderFunction, AlwaysNullable, RewriteWhenAnalyze { + + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.ret(BooleanType.INSTANCE).args(ArrayType.of(BooleanType.INSTANCE))); + + // The argument is a bound Lambda. + public MapAll(Expression arg) { + this(MapLambdaFunctionUtils.requireLambda("map_all", arg)); + } + + private MapAll(Lambda lambda) { + this(MapLambdaFunctionUtils.rewrite(lambda, (body, key, value, entry) -> body)); + } + + private MapAll(MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda) { + super("map_all", rewrittenLambda.toArrayMap()); + } + + private MapAll(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public MapAll withChildren(List children) { + Preconditions.checkArgument(children.size() == 1); + return new MapAll(getFunctionParams(children)); + } + + @Override + public List getImplSignature() { + return SIGNATURES; + } + + @Override + public Expression rewriteWhenAnalyze() { + return new ArrayMatchAll(getArgument(0)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitMapAll(this, context); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapApply.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapApply.java new file mode 100644 index 00000000000000..f0afeef273b209 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapApply.java @@ -0,0 +1,129 @@ +// 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.catalog.FunctionSignature; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.PreferPushDownProject; +import org.apache.doris.nereids.trees.expressions.functions.CustomSignature; +import org.apache.doris.nereids.trees.expressions.functions.PropagateNullable; +import org.apache.doris.nereids.trees.expressions.functions.RewriteWhenAnalyze; +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.MapType; +import org.apache.doris.nereids.types.StructField; +import org.apache.doris.nereids.types.StructType; + +import com.google.common.base.Preconditions; + +import java.util.List; + +/** + * Scalar function map_apply. + * + *

The lambda produces a two-field Struct for each Map entry. After analysis, the mapped entry + * array is converted directly to Map: + * + * 

+ * map_apply((mapKey, mapValue) -> struct(newKey, newValue), inputMap)
+ *   ->
+ * map_from_entries(array_map(
+ *   entry -> struct(newKey(entry[1], entry[2]), newValue(entry[1], entry[2])),
+ *   map_entries(inputMap)))
+ *
+ */ +public class MapApply extends ScalarFunction + implements CustomSignature, PropagateNullable, PreferPushDownProject, + RewriteWhenAnalyze { + + public MapApply(Expression arg) { + this(MapLambdaFunctionUtils.requireLambda("map_apply", arg)); + } + + private MapApply(Lambda lambda) { + this(validateAndRewrite(lambda)); + } + + private MapApply(MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda) { + super("map_apply", rewrittenLambda.getMapExpression(), rewrittenLambda.toArrayMap()); + } + + private MapApply(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public FunctionSignature customSignature() { + DataType mappedEntriesType = getArgument(1).getDataType(); + if (!(mappedEntriesType instanceof ArrayType) + || !(((ArrayType) mappedEntriesType).getItemType() instanceof StructType)) { + throw invalidReturnType(); + } + StructType structType = (StructType) ((ArrayType) mappedEntriesType).getItemType(); + if (structType.getFields().size() != 2) { + throw invalidReturnType(); + } + MapType inputMapType = (MapType) getArgument(0).getDataType(); + List fields = structType.getFields(); + MapType resultType = MapType.of(fields.get(0).getDataType(), fields.get(1).getDataType()); + resultType.validateDataType(); + return FunctionSignature.ret(resultType).args(inputMapType, mappedEntriesType); + } + + @Override + public MapApply withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new MapApply(getFunctionParams(children)); + } + + @Override + public Expression rewriteWhenAnalyze() { + return new MapFromEntries(getArgument(1)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitMapApply(this, context); + } + + private static MapLambdaFunctionUtils.RewrittenMapLambda validateAndRewrite(Lambda lambda) { + MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda = MapLambdaFunctionUtils.rewrite( + lambda, (body, key, value, entry) -> body); + validateLambdaReturn(lambda); + return rewrittenLambda; + } + + private static void validateLambdaReturn(Lambda lambda) { + Expression lambdaBody = lambda.getLambdaFunction(); + if (!(lambdaBody.getDataType() instanceof StructType) + || ((StructType) lambdaBody.getDataType()).getFields().size() != 2 + || lambdaBody.nullable()) { + throw invalidReturnType(); + } + StructType structType = (StructType) lambdaBody.getDataType(); + List fields = structType.getFields(); + MapType.of(fields.get(0).getDataType(), fields.get(1).getDataType()).validateDataType(); + } + + private static AnalysisException invalidReturnType() { + return new AnalysisException( + "Lambda of map_apply must return a non-nullable struct with exactly two fields"); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapExists.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapExists.java new file mode 100644 index 00000000000000..36687cf81be88d --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapExists.java @@ -0,0 +1,90 @@ +// 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.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable; +import org.apache.doris.nereids.trees.expressions.functions.RewriteWhenAnalyze; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.BooleanType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * Scalar function map_exists. + * + *

The Map lambda first produces one Boolean per entry, then ArrayMatchAny checks the result: + * + * 

+ * map_exists((mapKey, mapValue) -> predicate, inputMap)
+ *   ->
+ * array_match_any(
+ *   array_map(
+ *     entry -> predicate(entry[1], entry[2]),
+ *     map_entries(inputMap)))
+ *
+ */ +public class MapExists extends ScalarFunction + implements HighOrderFunction, AlwaysNullable, RewriteWhenAnalyze { + + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.ret(BooleanType.INSTANCE).args(ArrayType.of(BooleanType.INSTANCE))); + + /** Constructor with a bound Lambda argument. */ + public MapExists(Expression arg) { + this(MapLambdaFunctionUtils.requireLambda("map_exists", arg)); + } + + private MapExists(Lambda lambda) { + this(MapLambdaFunctionUtils.rewrite(lambda, (body, key, value, entry) -> body)); + } + + private MapExists(MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda) { + super("map_exists", rewrittenLambda.toArrayMap()); + } + + private MapExists(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public MapExists withChildren(List children) { + Preconditions.checkArgument(children.size() == 1); + return new MapExists(getFunctionParams(children)); + } + + @Override + public List getImplSignature() { + return SIGNATURES; + } + + @Override + public Expression rewriteWhenAnalyze() { + return new ArrayMatchAny(getArgument(0)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitMapExists(this, context); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFilter.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFilter.java new file mode 100644 index 00000000000000..b5613b0ebe7a24 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFilter.java @@ -0,0 +1,112 @@ +// 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.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.PropagateNullable; +import org.apache.doris.nereids.trees.expressions.functions.RewriteWhenAnalyze; +import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.BooleanType; +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; + +import java.util.List; + +/** + * Scalar function map_filter. + * + *

The Map lambda is evaluated by an ArrayMap over the Map's entry array: + * + * 

+ * map_filter((mapKey, mapValue) -> predicate, inputMap)
+ *   ->
+ * %map_from_filtered_entries_unique%(
+ *   array_map(
+ *     entry -> if(predicate(entry[1], entry[2]), entry, null),
+ *     map_entries(inputMap)))
+ *
+ */ +public class MapFilter extends ScalarFunction + implements HighOrderFunction, PropagateNullable, RewriteWhenAnalyze { + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.retArgType(0).args( + MapType.of(new AnyDataType(0), new AnyDataType(1)), + ArrayType.of(BooleanType.INSTANCE))); + private static final List MAP_LAMBDA_SIGNATURES = ImmutableList.of( + FunctionSignature.retArgType(0).args( + MapType.of(new AnyDataType(0), new AnyDataType(1)), + ArrayType.of(new AnyDataType(2)))); + + private final boolean validateMapLambdaInput; + + // The argument is a bound Lambda. + public MapFilter(Expression arg) { + this(MapLambdaFunctionUtils.requireLambda("map_filter", arg)); + } + + public MapFilter(Expression map, Expression filter) { + super("map_filter", map, filter); + validateMapLambdaInput = false; + } + + private MapFilter(Lambda lambda) { + this(MapLambdaFunctionUtils.rewrite(lambda, + (body, key, value, entry) -> new If( + body, entry, new NullLiteral(entry.getDataType())))); + } + + private MapFilter(MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda) { + super("map_filter", + rewrittenLambda.getMapExpression(), rewrittenLambda.toArrayMap()); + validateMapLambdaInput = true; + } + + private MapFilter(ScalarFunctionParams functionParams, boolean validateMapLambdaInput) { + super(functionParams); + this.validateMapLambdaInput = validateMapLambdaInput; + } + + @Override + public MapFilter withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new MapFilter(getFunctionParams(children), validateMapLambdaInput); + } + + @Override + public List getImplSignature() { + return validateMapLambdaInput ? MAP_LAMBDA_SIGNATURES : SIGNATURES; + } + + @Override + public Expression rewriteWhenAnalyze() { + return validateMapLambdaInput + ? new MapFromFilteredEntriesUnique(getArgument(1)) + : this; + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitMapFilter(this, context); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromEntries.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromEntries.java index e2f6d1ec01cf8f..c79552776d6d5f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromEntries.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromEntries.java @@ -48,7 +48,11 @@ public MapFromEntries(Expression entries) { super("map_from_entries", entries); } - private MapFromEntries(ScalarFunctionParams functionParams) { + protected MapFromEntries(String name, Expression entries) { + super(name, entries); + } + + protected MapFromEntries(ScalarFunctionParams functionParams) { super(functionParams); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromEntriesUnique.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromEntriesUnique.java new file mode 100644 index 00000000000000..24176e0b6a28e8 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromEntriesUnique.java @@ -0,0 +1,42 @@ +// 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.trees.expressions.Expression; + +import com.google.common.base.Preconditions; + +import java.util.List; + +/** Internal Map constructor used when entry keys are known to be unique. */ +public class MapFromEntriesUnique extends MapFromEntries { + + public MapFromEntriesUnique(Expression entries) { + super("%map_from_entries_unique%", entries); + } + + private MapFromEntriesUnique(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public MapFromEntriesUnique withChildren(List children) { + Preconditions.checkArgument(children.size() == 1); + return new MapFromEntriesUnique(getFunctionParams(children)); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromFilteredEntriesUnique.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromFilteredEntriesUnique.java new file mode 100644 index 00000000000000..1f7cd306299f78 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromFilteredEntriesUnique.java @@ -0,0 +1,42 @@ +// 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.trees.expressions.Expression; + +import com.google.common.base.Preconditions; + +import java.util.List; + +/** Internal Map constructor that drops null entries produced by a Map-filter Lambda. */ +public class MapFromFilteredEntriesUnique extends MapFromEntries { + + public MapFromFilteredEntriesUnique(Expression entries) { + super("%map_from_filtered_entries_unique%", entries); + } + + private MapFromFilteredEntriesUnique(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public MapFromFilteredEntriesUnique withChildren(List children) { + Preconditions.checkArgument(children.size() == 1); + return new MapFromFilteredEntriesUnique(getFunctionParams(children)); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaFunctionUtils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaFunctionUtils.java new file mode 100644 index 00000000000000..9bd8224c903f0e --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaFunctionUtils.java @@ -0,0 +1,94 @@ +// 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.trees.expressions.ArrayItemReference; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** Utilities for building Map Lambda functions from one bound Map-entry array driver. */ +final class MapLambdaFunctionUtils { + + private MapLambdaFunctionUtils() { + } + + /** Build the function-specific body around a bound Map-entry Lambda. */ + static RewrittenMapLambda rewrite(Lambda lambda, EntryBodyBuilder bodyBuilder) { + Expression mapExpression = extractMapExpression(lambda); + List arguments = lambda.getLambdaArguments(); + ArrayItemReference entryArgument = arguments.get(0); + Slot entrySlot = entryArgument.toSlot(); + Expression key = new ElementAt(entrySlot, new IntegerLiteral(1)); + Expression value = new ElementAt(entrySlot, new IntegerLiteral(2)); + + Expression rewrittenBody = bodyBuilder.build(lambda.getLambdaFunction(), key, value, entrySlot); + Lambda rewrittenLambda = new Lambda( + ImmutableList.of(entryArgument.getName()), rewrittenBody, ImmutableList.of(entryArgument)); + return new RewrittenMapLambda(mapExpression, rewrittenLambda); + } + + /** Require a bound Lambda argument. */ + static Lambda requireLambda(String functionName, Expression expression) { + if (!(expression instanceof Lambda)) { + throw new AnalysisException(String.format( + "The 1st arg of %s must be lambda but is %s", functionName, expression)); + } + return (Lambda) expression; + } + + private static Expression extractMapExpression(Lambda lambda) { + List arguments = lambda.getLambdaArguments(); + Preconditions.checkArgument(arguments.size() == 1, + "A bound Map Lambda must have one entry argument"); + Expression entries = arguments.get(0).getArrayExpression(); + Preconditions.checkArgument(entries instanceof MapEntries, + "A bound Map Lambda must use a map_entries argument"); + return entries.child(0); + } + + @FunctionalInterface + interface EntryBodyBuilder { + Expression build(Expression body, Expression key, Expression value, Slot entry); + } + + /** One original Map and its one-driver entry Lambda. */ + static final class RewrittenMapLambda { + private final Expression mapExpression; + private final Lambda lambda; + + private RewrittenMapLambda(Expression mapExpression, Lambda lambda) { + this.mapExpression = mapExpression; + this.lambda = lambda; + } + + Expression getMapExpression() { + return mapExpression; + } + + ArrayMap toArrayMap() { + return new ArrayMap(lambda); + } + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TransformKeys.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TransformKeys.java new file mode 100644 index 00000000000000..d5b832075bc743 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TransformKeys.java @@ -0,0 +1,86 @@ +// 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.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.PreferPushDownProject; +import org.apache.doris.nereids.trees.expressions.functions.CustomSignature; +import org.apache.doris.nereids.trees.expressions.functions.PropagateNullable; +import org.apache.doris.nereids.trees.expressions.functions.RewriteWhenAnalyze; +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.MapType; +import org.apache.doris.nereids.types.StructField; +import org.apache.doris.nereids.types.StructType; + +import com.google.common.base.Preconditions; + +import java.util.List; + +/** Scalar function transform_keys. */ +public class TransformKeys extends ScalarFunction + implements CustomSignature, PropagateNullable, PreferPushDownProject, RewriteWhenAnalyze { + + public TransformKeys(Expression arg) { + this(MapLambdaFunctionUtils.requireLambda("transform_keys", arg)); + } + + private TransformKeys(Lambda lambda) { + this(MapLambdaFunctionUtils.rewrite(lambda, + (body, key, value, entry) -> new CreateStruct(body, value))); + } + + private TransformKeys(MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda) { + super("transform_keys", + rewrittenLambda.getMapExpression(), rewrittenLambda.toArrayMap()); + } + + private TransformKeys(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public FunctionSignature customSignature() { + MapType inputMapType = (MapType) getArgument(0).getDataType(); + ArrayType mappedEntriesType = (ArrayType) getArgument(1).getDataType(); + StructType entryType = (StructType) mappedEntriesType.getItemType(); + List fields = entryType.getFields(); + DataType resultKeyType = fields.get(0).getDataType(); + MapType resultType = MapType.of(resultKeyType, inputMapType.getValueType()); + resultType.validateDataType(); + return FunctionSignature.ret(resultType).args(inputMapType, mappedEntriesType); + } + + @Override + public TransformKeys withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new TransformKeys(getFunctionParams(children)); + } + + @Override + public Expression rewriteWhenAnalyze() { + return new MapFromEntries(getArgument(1)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitTransformKeys(this, context); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TransformValues.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TransformValues.java new file mode 100644 index 00000000000000..47f549d09177b5 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TransformValues.java @@ -0,0 +1,99 @@ +// 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.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.PreferPushDownProject; +import org.apache.doris.nereids.trees.expressions.functions.CustomSignature; +import org.apache.doris.nereids.trees.expressions.functions.PropagateNullable; +import org.apache.doris.nereids.trees.expressions.functions.RewriteWhenAnalyze; +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.MapType; +import org.apache.doris.nereids.types.StructField; +import org.apache.doris.nereids.types.StructType; + +import com.google.common.base.Preconditions; + +import java.util.List; + +/** + * Scalar function transform_values. + * + *

The original keys are retained while the Map lambda produces the new value array: + * + * 

+ * transform_values((mapKey, mapValue) -> newValue, inputMap)
+ *   ->
+ * %map_from_entries_unique%(
+ *   array_map(
+ *     entry -> struct(entry[1], newValue(entry[1], entry[2])),
+ *     map_entries(inputMap)))
+ *
+ */ +public class TransformValues extends ScalarFunction + implements CustomSignature, PropagateNullable, PreferPushDownProject, RewriteWhenAnalyze { + + public TransformValues(Expression arg) { + this(MapLambdaFunctionUtils.requireLambda("transform_values", arg)); + } + + private TransformValues(Lambda lambda) { + this(MapLambdaFunctionUtils.rewrite(lambda, + (body, key, value, entry) -> new CreateStruct(key, body))); + } + + private TransformValues(MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda) { + super("transform_values", + rewrittenLambda.getMapExpression(), rewrittenLambda.toArrayMap()); + } + + private TransformValues(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public FunctionSignature customSignature() { + MapType inputMapType = (MapType) getArgument(0).getDataType(); + ArrayType mappedEntriesType = (ArrayType) getArgument(1).getDataType(); + StructType entryType = (StructType) mappedEntriesType.getItemType(); + List fields = entryType.getFields(); + DataType resultValueType = fields.get(1).getDataType(); + MapType resultType = MapType.of(inputMapType.getKeyType(), resultValueType); + resultType.validateDataType(); + return FunctionSignature.ret(resultType).args(inputMapType, mappedEntriesType); + } + + @Override + public TransformValues withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new TransformValues(getFunctionParams(children)); + } + + @Override + public Expression rewriteWhenAnalyze() { + return new MapFromEntriesUnique(getArgument(1)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitTransformValues(this, context); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java index c0f03da1789b48..ec86766df02f2e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java @@ -350,10 +350,14 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.MakeDate; import org.apache.doris.nereids.trees.expressions.functions.scalar.MakeSet; import org.apache.doris.nereids.trees.expressions.functions.scalar.MakeTime; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapAll; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapApply; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapContainsEntry; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapContainsKey; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapContainsValue; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapEntries; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapExists; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapFilter; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapFromArrays; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapFromEntries; import org.apache.doris.nereids.trees.expressions.functions.scalar.MapKeys; @@ -560,6 +564,8 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.ToSeconds; import org.apache.doris.nereids.trees.expressions.functions.scalar.Tokenize; import org.apache.doris.nereids.trees.expressions.functions.scalar.TopLevelDomain; +import org.apache.doris.nereids.trees.expressions.functions.scalar.TransformKeys; +import org.apache.doris.nereids.trees.expressions.functions.scalar.TransformValues; import org.apache.doris.nereids.trees.expressions.functions.scalar.Translate; import org.apache.doris.nereids.trees.expressions.functions.scalar.Trim; import org.apache.doris.nereids.trees.expressions.functions.scalar.TrimIn; @@ -2828,6 +2834,14 @@ default R visitCreateMap(CreateMap createMap, C context) { return visitScalarFunction(createMap, context); } + default R visitMapAll(MapAll mapAll, C context) { + return visitScalarFunction(mapAll, context); + } + + default R visitMapApply(MapApply mapApply, C context) { + return visitScalarFunction(mapApply, context); + } + default R visitMapContainsKey(MapContainsKey mapContainsKey, C context) { return visitScalarFunction(mapContainsKey, context); } @@ -2844,6 +2858,14 @@ default R visitMapEntries(MapEntries mapEntries, C context) { return visitScalarFunction(mapEntries, context); } + default R visitMapExists(MapExists mapExists, C context) { + return visitScalarFunction(mapExists, context); + } + + default R visitMapFilter(MapFilter mapFilter, C context) { + return visitScalarFunction(mapFilter, context); + } + default R visitMapFromArrays(MapFromArrays mapFromArrays, C context) { return visitScalarFunction(mapFromArrays, context); } @@ -2864,6 +2886,14 @@ default R visitMapValues(MapValues mapValues, C context) { return visitScalarFunction(mapValues, context); } + default R visitTransformKeys(TransformKeys transformKeys, C context) { + return visitScalarFunction(transformKeys, context); + } + + default R visitTransformValues(TransformValues transformValues, C context) { + return visitScalarFunction(transformValues, context); + } + default R visitXor(Xor xor, C context) { return visitScalarFunction(xor, context); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java index 3f636840d97731..16f524f4037dcd 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java @@ -486,10 +486,33 @@ public void testPruneArrayLambda() throws Exception { assertColumn("select array_map(m -> array_map(x -> element_at(map_values(m)[0], 'a'), [1]), " + "element_at(s, 'data')) from tbl", + "struct>>>", + ImmutableList.of(path("s", "data", "*", "VALUES", "a")), + ImmutableList.of() + ); + + assertColumn("select array_map(m -> array_map(x -> element_at(map_values(m)[0], 'a'), [1]), " + + "element_at(s, 'data')) from tbl", + "struct>>>", + ImmutableList.of(path("s", "data", "*", "VALUES", "a")), + ImmutableList.of() + ); + } + + @Test + public void testPruneMapEntryLambda() throws Exception { + assertColumn("select map_exists((k, v) -> element_at(v, 'a') > 0, " + + "element_at(s, 'data')[1]) from tbl", "struct>>>", ImmutableList.of(path("s", "data", "*", "VALUES", "a")), ImmutableList.of() ); + + assertColumn("select map_all((k, v) -> k > 0, element_at(s, 'data')[1]) from tbl", + "struct>>>", + ImmutableList.of(path("s", "data", "*", "KEYS")), + ImmutableList.of() + ); } @Test diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaFunctionsTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaFunctionsTest.java new file mode 100644 index 00000000000000..b86920ebe1dfe2 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaFunctionsTest.java @@ -0,0 +1,281 @@ +// 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.NereidsPlanner; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.glue.translator.PhysicalPlanTranslator; +import org.apache.doris.nereids.glue.translator.PlanTranslatorContext; +import org.apache.doris.nereids.parser.NereidsParser; +import org.apache.doris.nereids.properties.PhysicalProperties; +import org.apache.doris.nereids.rules.expression.ExpressionRewriteTestHelper; +import org.apache.doris.nereids.trees.expressions.ArrayItemReference; +import org.apache.doris.nereids.trees.expressions.Cast; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.physical.PhysicalPlan; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.types.MapType; +import org.apache.doris.nereids.types.SmallIntType; +import org.apache.doris.nereids.types.TinyIntType; +import org.apache.doris.nereids.util.MemoTestUtils; +import org.apache.doris.utframe.TestWithFeService; + +import com.google.common.collect.ImmutableList; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.List; + +public class MapLambdaFunctionsTest extends TestWithFeService { + + private static final NereidsParser PARSER = new NereidsParser(); + + @Override + protected void runBeforeAll() throws Exception { + createDatabase("map_lambda_function_test"); + useDatabase("map_lambda_function_test"); + createTables( + "CREATE TABLE map_lambda_left (id INT) DUPLICATE KEY(id) " + + "DISTRIBUTED BY HASH(id) BUCKETS 1 " + + "PROPERTIES ('replication_num' = '1')", + "CREATE TABLE map_lambda_right (id INT) DUPLICATE KEY(id) " + + "DISTRIBUTED BY HASH(id) BUCKETS 1 " + + "PROPERTIES ('replication_num' = '1')", + "CREATE TABLE map_lambda_filter (id INT, numeric_map MAP) DUPLICATE KEY(id) " + + "DISTRIBUTED BY HASH(id) BUCKETS 1 " + + "PROPERTIES ('replication_num' = '1')"); + } + + @Override + protected void runBeforeEach() throws Exception { + StatementScopeIdGenerator.clear(); + } + + @Test + public void testMapLambdaWrappersAndTypes() { + Expression mapApply = analyze("map_apply((k, v) -> struct(k + 1, v * 2), map(1, 10, 2, 20))"); + Assertions.assertTrue(mapApply instanceof MapFromEntries); + assertMapEntryArray(mapApply.child(0)); + Assertions.assertEquals(MapType.of(SmallIntType.INSTANCE, SmallIntType.INSTANCE), + mapApply.getDataType()); + + Expression tupleMapApply = analyze( + "map_apply((k, v) -> (k + 1, v * 2), map(1, 10, 2, 20))"); + Assertions.assertTrue(tupleMapApply instanceof MapFromEntries); + assertMapEntryArray(tupleMapApply.child(0)); + Assertions.assertEquals(MapType.of(SmallIntType.INSTANCE, SmallIntType.INSTANCE), + tupleMapApply.getDataType()); + + Expression mapFilter = analyze("map_filter((k, v) -> v > 10, map(1, 10, 2, 20))"); + Assertions.assertTrue(mapFilter instanceof MapFromFilteredEntriesUnique); + assertMapEntryArray(mapFilter.child(0)); + + Expression mapFilterWithMask = analyze("map_filter(map(1, 10, 2, 20), " + + "array_map((k, v) -> v > k, [1, 2], [10, 20]))"); + Assertions.assertTrue(mapFilterWithMask instanceof MapFilter); + Assertions.assertTrue(mapFilterWithMask.child(1).getDataType() instanceof ArrayType); + + Expression transformKeys = analyze("transform_keys((k, v) -> k + 1, map(1, 10, 2, 20))"); + Assertions.assertTrue(transformKeys instanceof MapFromEntries); + assertMapEntryArray(transformKeys.child(0)); + Assertions.assertEquals(MapType.of(SmallIntType.INSTANCE, TinyIntType.INSTANCE), + transformKeys.getDataType()); + + Expression transformValues = analyze( + "transform_values((k, v) -> v + 1, map(1, 10, 2, 20))"); + Assertions.assertTrue(transformValues instanceof MapFromEntriesUnique); + assertMapEntryArray(transformValues.child(0)); + Assertions.assertEquals(MapType.of(TinyIntType.INSTANCE, SmallIntType.INSTANCE), + transformValues.getDataType()); + } + + @Test + public void testMapExistsAndAllRewriteToArrayMatch() { + Expression mapExists = analyze("map_exists((k, v) -> v > 10, map(1, 10, 2, 20))"); + Assertions.assertTrue(mapExists instanceof ArrayMatchAny); + assertMapEntryArray(mapExists.child(0)); + + Expression mapAll = analyze("map_all((k, v) -> v > 10, map(1, 10, 2, 20))"); + Assertions.assertTrue(mapAll instanceof ArrayMatchAll); + assertMapEntryArray(mapAll.child(0)); + } + + @Test + public void testNestedLambdaCanCaptureImmediateOuterScope() { + Expression nested = analyze("map_exists((x, v) -> " + + "array_match_any(x -> x > v, [1]), map(1, 10))"); + Assertions.assertTrue(nested instanceof ArrayMatchAny); + } + + @Test + public void testFreshEntryNameAvoidsUserLambdaArgumentCollision() { + Expression mapExists = analyze("map_exists(($_map_entry_2_$, v) -> " + + "$_map_entry_2_$ > 0, map(1, 10))"); + Assertions.assertTrue(mapExists instanceof ArrayMatchAny); + ArrayMap arrayMap = (ArrayMap) mapExists.child(0); + Lambda lambda = (Lambda) arrayMap.child(0); + Assertions.assertNotEquals("$_map_entry_2_$", lambda.getLambdaArgumentName(0)); + assertMapEntryArray(arrayMap); + } + + @Test + public void testComputedMapIsAccepted() { + SlotReference value = new SlotReference("value", IntegerType.INSTANCE); + CreateMap computedMap = new CreateMap(Literal.of(1), value); + ArrayItemReference entryArgument = new ArrayItemReference("entry", new MapEntries(computedMap)); + Lambda boundLambda = new Lambda( + ImmutableList.of("entry"), + new ElementAt(entryArgument.toSlot(), new IntegerLiteral(2)), + ImmutableList.of(entryArgument)); + TransformValues transformValues = new TransformValues(boundLambda); + + Assertions.assertSame(computedMap, transformValues.child(0)); + + Expression nondeterministicMap = analyze( + "transform_values((k, v) -> v, map(cast(random() as int), 10))"); + Assertions.assertTrue(nondeterministicMap instanceof MapFromEntriesUnique); + assertMapEntryArray(nondeterministicMap.child(0)); + } + + @Test + public void testMapDependingOnBothJoinSidesCanBeTranslated() { + translate("SELECT l.id, r.id FROM map_lambda_left l JOIN map_lambda_right r " + + "ON map_exists((k, v) -> v > 0, map(l.id, r.id))"); + translate("SELECT l.id, r.id FROM map_lambda_left l JOIN map_lambda_right r " + + "ON map_contains_key(transform_values((k, v) -> v, map(l.id, r.id)), l.id)"); + } + + @Test + public void testComputedMapInNestedMapLambdaCanBeTranslated() { + translate("SELECT transform_values((ok, ov) -> transform_values(" + + "(ik, iv) -> ik, map(ok + random(), ov)), map(1, 10))"); + } + + @Test + public void testComputedMapInArrayMapLambdaCanBeTranslated() { + translate("SELECT array_map(x -> map_keys(" + + "transform_values((k, v) -> k, map(uuid(), x)))[1], [1, 2])"); + translate("SELECT array_map(x -> map_keys(" + + "map_apply((k, v) -> struct(k, k), map(uuid(), x)))[1], [1, 2])"); + } + + @Test + public void testPartiallyMaterializedMapLambdaInputsCanBeTranslated() { + translate("SELECT transform_keys((k, v) -> 1, map(1, 10, 2, 20))"); + translate("SELECT transform_values((k, v) -> 1, map(1, 10, 2, 20))"); + translate("SELECT map_filter((k, v) -> k > 0, map(1, 10, 2, 20))"); + translate("SELECT map_filter(map(1, 10, 2, 20), " + + "array_map((k, v) -> v > k, [1, 2], [10, 20]))"); + translate("SELECT count(map_filter(numeric_map, array_map(" + + "(k, v) -> v > k + id, map_keys(numeric_map), map_values(numeric_map)))) " + + "FROM map_lambda_filter"); + } + + @Test + public void testMapFromArraysCanBeAnalyzed() { + Expression map = analyze("map_from_arrays([1, 2], [10, 20])"); + Assertions.assertTrue(map instanceof MapFromArrays); + Assertions.assertEquals( + MapType.of(TinyIntType.INSTANCE, TinyIntType.INSTANCE), map.getDataType()); + + Assertions.assertThrows(RuntimeException.class, + () -> analyze("map_from_arrays([1, 2], 10)")); + + AnalysisException complexKeyException = Assertions.assertThrows(AnalysisException.class, + () -> analyze("map_from_arrays([[1]], [10])")); + Assertions.assertTrue(complexKeyException.getMessage().contains( + "MAP key type must be a primitive type"), complexKeyException::getMessage); + + Expression nestedNullValueMap = analyze("map_from_arrays([1], [[]])"); + Assertions.assertEquals( + MapType.of(TinyIntType.INSTANCE, ArrayType.of(TinyIntType.INSTANCE)), + nestedNullValueMap.getDataType()); + } + + @Test + public void testMapFromEntriesCanBeAnalyzed() { + Expression map = analyze("map_from_entries(array(struct(1, 10), struct(2, 20)))"); + Assertions.assertTrue(map instanceof MapFromEntries); + Assertions.assertEquals( + MapType.of(TinyIntType.INSTANCE, TinyIntType.INSTANCE), map.getDataType()); + + Assertions.assertThrows(AnalysisException.class, + () -> analyze("map_from_entries(1)")); + Assertions.assertThrows(AnalysisException.class, + () -> analyze("map_from_entries(array(struct(1)))")); + + AnalysisException complexKeyException = Assertions.assertThrows(AnalysisException.class, + () -> analyze("map_from_entries(array(struct([1], 10)))")); + Assertions.assertTrue(complexKeyException.getMessage().contains( + "MAP key type must be a primitive type"), complexKeyException::getMessage); + } + + @Test + public void testMapLambdaRejectsInvalidArguments() { + Assertions.assertThrows(RuntimeException.class, + () -> analyze("map_filter(k -> k > 0, map(1, 10))")); + Assertions.assertThrows(RuntimeException.class, + () -> analyze("map_filter((k, v) -> v > 0, [1, 2])")); + Assertions.assertThrows(RuntimeException.class, + () -> analyze("map_filter((k, v) -> v > 0, map(1, 10), map(2, 20))")); + Assertions.assertThrows(RuntimeException.class, + () -> analyze("map_apply((k, v) -> if(k > 0, struct(k, v), null), map(1, 10))")); + Assertions.assertThrows(RuntimeException.class, + () -> analyze("map_apply((k, v) -> (k, v, k + v), map(1, 10))")); + } + + private Expression analyze(String sql) { + return ExpressionRewriteTestHelper.typeCoercion(PARSER.parseExpression(sql)); + } + + private void translate(String sql) { + StatementContext statementContext = MemoTestUtils.createStatementContext(connectContext, sql); + NereidsPlanner planner = new NereidsPlanner(statementContext); + PhysicalPlan plan = planner.planWithLock(PARSER.parseSingle(sql), PhysicalProperties.ANY); + new PhysicalPlanTranslator(new PlanTranslatorContext(planner.getCascadesContext())) + .translatePlan(plan); + } + + private void assertMapEntryArray(Expression expression) { + while (expression instanceof Cast) { + expression = expression.child(0); + } + Assertions.assertTrue(expression instanceof ArrayMap); + ArrayMap arrayMap = (ArrayMap) expression; + Assertions.assertTrue(arrayMap.child(0) instanceof Lambda); + + Lambda lambda = (Lambda) arrayMap.child(0); + List arguments = lambda.getLambdaArguments(); + Assertions.assertEquals(1, arguments.size()); + Assertions.assertTrue(arguments.get(0).getArrayExpression() instanceof MapEntries); + Assertions.assertTrue(arguments.get(0).getName().startsWith("$_map_entry_")); + Assertions.assertEquals(1, + arrayMap.collect(child -> child instanceof MapEntries).size()); + Assertions.assertTrue(arrayMap.collect(child -> child instanceof MapKeys).isEmpty()); + Assertions.assertTrue(arrayMap.collect(child -> child instanceof MapValues).isEmpty()); + + Assertions.assertTrue(arrayMap.withChildren(arrayMap.children()) instanceof ArrayMap); + } + +} diff --git a/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 b/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 index 531310b611a727..53fde2c2663766 100644 --- a/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 +++ b/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 @@ -23,6 +23,40 @@ options { tokenVocab = DorisLexer; } @members { public boolean ansiSQLSyntax = false; + + private boolean isTupleLambdaBody() { + if (_input.LA(1) != LEFT_PAREN) { + return false; + } + int parenthesisDepth = 0; + int bracketDepth = 0; + int braceDepth = 0; + for (int offset = 1; ; offset++) { + int tokenType = _input.LA(offset); + if (tokenType == Token.EOF) { + return false; + } + if (tokenType == LEFT_PAREN) { + parenthesisDepth++; + } else if (tokenType == RIGHT_PAREN) { + parenthesisDepth--; + if (parenthesisDepth == 0) { + return false; + } + } else if (tokenType == LEFT_BRACKET) { + bracketDepth++; + } else if (tokenType == RIGHT_BRACKET) { + bracketDepth--; + } else if (tokenType == LEFT_BRACE) { + braceDepth++; + } else if (tokenType == RIGHT_BRACE) { + braceDepth--; + } else if (tokenType == COMMA && parenthesisDepth == 1 + && bracketDepth == 0 && braceDepth == 0) { + return true; + } + } + } } multiStatements @@ -1648,7 +1682,9 @@ lambdaExpression | LEFT_PAREN args+=errorCapturingIdentifier (COMMA args+=errorCapturingIdentifier)+ RIGHT_PAREN - ARROW body=booleanExpression + ARROW ({isTupleLambdaBody()}? + LEFT_PAREN bodyItems+=expression (COMMA bodyItems+=expression)+ RIGHT_PAREN + | body=booleanExpression) ; booleanExpression diff --git a/regression-test/data/query_p0/sql_functions/map_functions/test_map_lambda.out b/regression-test/data/query_p0/sql_functions/map_functions/test_map_lambda.out new file mode 100644 index 00000000000000..8bdd1713035e4d --- /dev/null +++ b/regression-test/data/query_p0/sql_functions/map_functions/test_map_lambda.out @@ -0,0 +1,122 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !map_apply -- +2 20 40 + +-- !map_apply_tuple_lambda -- +2 11 21 2 11 21 + +-- !map_from_arrays -- +2 10 20 + +-- !map_from_arrays_last_win -- +1 20 + +-- !map_from_entries -- +2 10 20 + +-- !map_from_entries_last_win -- +1 20 + +-- !map_from_entries_column -- +1 2 10 20 +2 0 \N \N +3 \N \N \N + +-- !map_from_entries_null -- +\N + +-- !map_filter_null_is_false -- +1 \N 20 + +-- !map_filter_two_argument_nullable_map -- +1 {1:10, 2:20} +2 {} +3 \N + +-- !map_filter_constant_map_captures_column -- +1 2 10 20 +2 2 10 20 +3 2 10 20 + +-- !transform_keys_last_win -- +1 20 + +-- !map_apply_last_win -- +1 40 + +-- !transform_values -- +11 21 + +-- !transform_values_string -- +x:a y:b + +-- !transform_values_array -- +[10, 1] [20, 21, 2] + +-- !nested_array_lambda -- +{1:[21], 2:[32, 33]} + +-- !map_quantifiers -- +1 0 + +-- !empty_map_quantifiers -- +0 1 + +-- !null_predicate_quantifiers -- +\N \N + +-- !empty_map_all_functions -- +0 0 0 0 0 1 + +-- !null_map_all_functions -- +1 1 1 1 1 1 + +-- !constant_map -- +11 21 + +-- !materialized_computed_map -- +11 31 + +-- !direct_computed_map -- +11 31 + +-- !all_functions_direct_computed_map -- +11 30 10 31 1 1 + +-- !map_lambda_aggregate_output -- +1 1 12 +2 1 13 +3 1 14 + +-- !map_lambda_having -- +1 1 +2 1 +3 1 + +-- !map_lambda_join_on -- +1 2 +2 3 + +-- !nondeterministic_map_input -- +1 10 + +-- !nondeterministic_map_single_evaluation -- +1 + +-- !map_lambda_generate_volatile -- +1 + +-- !nested_nondeterministic_map_single_evaluation -- +1 1 + +-- !lambda_nested_lambda_map_single_evaluation -- +1 + +-- !array_map_nested_map_lambda_single_evaluation -- +1 + +-- !array_map_nested_map_apply_single_evaluation -- +1 + +-- !map_from_arrays_nested_empty_array -- +{1:[]} diff --git a/regression-test/suites/query_p0/sql_functions/map_functions/test_map_lambda.groovy b/regression-test/suites/query_p0/sql_functions/map_functions/test_map_lambda.groovy new file mode 100644 index 00000000000000..9d0266da3dfb91 --- /dev/null +++ b/regression-test/suites/query_p0/sql_functions/map_functions/test_map_lambda.groovy @@ -0,0 +1,400 @@ +// 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_lambda", "p0") { + sql "set enable_nereids_planner = true" + sql "set enable_fallback_to_original_planner = false" + sql "drop table if exists test_map_lambda" + sql """ + create table test_map_lambda ( + id int, + bias int, + mii map, + mss map, + mia map>, + mim map> + ) + duplicate key(id) + distributed by hash(id) buckets 1 + properties("replication_num" = "1") + """ + sql """ + insert into test_map_lambda values + (1, 10, + map(1, 10, 2, 20), + map('a', 'x', 'b', 'y'), + map(1, [10], 2, [20, 21]), + map(1, map(2, 20), 3, map(4, 40))), + (2, 0, + cast(map() as map), + cast(map() as map), + cast(map() as map>), + cast(map() as map>)), + (3, 0, null, null, null, null) + """ + + qt_map_apply """ + select map_size(r), r[2], r[3] + from ( + select map_apply((k, v) -> struct(k + 1, v * 2), mii) r + from test_map_lambda where id = 1 + ) t + """ + + qt_map_apply_tuple_lambda """ + select map_size(tuple_result), tuple_result[3], tuple_result[6], + map_size(struct_result), struct_result[3], struct_result[6] + from ( + select map_apply((k, v) -> (k * 3, v + 1), mii) tuple_result, + map_apply((k, v) -> struct(k * 3, v + 1), mii) struct_result + from test_map_lambda where id = 1 + ) t + """ + + qt_map_from_arrays """ + select map_size(r), r[1], r[2] + from ( + select map_from_arrays([1, 2], [10, 20]) r + ) t + """ + + qt_map_from_arrays_last_win """ + select map_size(r), r[1] + from ( + select map_from_arrays([1, 1], [10, 20]) r + ) t + """ + + qt_map_from_entries """ + select map_size(r), r[1], r[2] + from ( + select map_from_entries(array(struct(1, 10), struct(2, 20))) r + ) t + """ + + qt_map_from_entries_last_win """ + select map_size(r), r[1] + from ( + select map_from_entries(array(struct(1, 10), struct(1, 20))) r + ) t + """ + + order_qt_map_from_entries_column """ + select id, map_size(r), r[1], r[2] + from ( + select id, map_from_entries(map_entries(mii)) r + from test_map_lambda + ) t + order by id + """ + + qt_map_from_entries_null """ + select map_from_entries(cast(null as array>)) + """ + + testFoldConst("select map_from_entries(array(struct(1, 'a'), struct(2, 'b')))") + + qt_map_filter_null_is_false """ + select map_size(r), r[1], r[2] + from ( + select map_filter( + (k, v) -> if(k = 1, cast(null as boolean), true), mii) r + from test_map_lambda where id = 1 + ) t + """ + + order_qt_map_filter_two_argument_nullable_map """ + select id, map_filter( + mii, if(id = 1, [true, true], if(id = 2, [], [true]))) + from test_map_lambda + order by id + """ + + order_qt_map_filter_constant_map_captures_column """ + select id, map_size(r), r[1], r[2] + from ( + select id, map_filter((k, v) -> v > id, map(1, 10, 2, 20)) r + from test_map_lambda + ) t + order by id + """ + + // Both transformations intentionally produce duplicate keys. ColumnMap uses last-win. + qt_transform_keys_last_win """ + select map_size(r), r[1] + from ( + select transform_keys((k, v) -> 1, mii) r + from test_map_lambda where id = 1 + ) t + """ + qt_map_apply_last_win """ + select map_size(r), r[1] + from ( + select map_apply((k, v) -> struct(1, v * 2), mii) r + from test_map_lambda where id = 1 + ) t + """ + + qt_transform_values """ + select r[1], r[2] + from ( + select transform_values((k, v) -> v + 1, mii) r + from test_map_lambda where id = 1 + ) t + """ + qt_transform_values_string """ + select r['a'], r['b'] + from ( + select transform_values((k, v) -> concat(v, ':', k), mss) r + from test_map_lambda where id = 1 + ) t + """ + qt_transform_values_array """ + select r[1], r[2] + from ( + select transform_values((k, v) -> array_pushback(v, k), mia) r + from test_map_lambda where id = 1 + ) t + """ + qt_nested_array_lambda """ + select map_apply( + (k, vals) -> struct( + k, array_map(x -> x + k + 10, vals)), + map(1, [10], 2, [20, 21])) + """ + + qt_map_quantifiers """ + select cast(map_exists((k, v) -> v = 20, mii) as int), + cast(map_all((k, v) -> v > 10, mii) as int) + from test_map_lambda where id = 1 + """ + qt_empty_map_quantifiers """ + select cast(map_exists((k, v) -> true, mii) as int), + cast(map_all((k, v) -> false, mii) as int) + from test_map_lambda where id = 2 + """ + qt_null_predicate_quantifiers """ + select map_exists((k, v) -> cast(null as boolean), mii), + map_all((k, v) -> cast(null as boolean), mii) + from test_map_lambda where id = 1 + """ + + qt_empty_map_all_functions """ + select map_size(map_apply((k, v) -> struct(k, v), mii)), + map_size(map_filter((k, v) -> true, mii)), + map_size(transform_keys((k, v) -> k + 1, mii)), + map_size(transform_values((k, v) -> v + 1, mii)), + cast(map_exists((k, v) -> true, mii) as int), + cast(map_all((k, v) -> false, mii) as int) + from test_map_lambda where id = 2 + """ + + qt_null_map_all_functions """ + select cast(map_apply((k, v) -> struct(k, v), mii) is null as int), + cast(map_filter((k, v) -> true, mii) is null as int), + cast(transform_keys((k, v) -> k + 1, mii) is null as int), + cast(transform_values((k, v) -> v + 1, mii) is null as int), + cast(map_exists((k, v) -> true, mii) is null as int), + cast(map_all((k, v) -> false, mii) is null as int) + from test_map_lambda where id = 3 + """ + + // A deterministic constant Map is an explicitly supported stable source. + qt_constant_map """ + select r[1], r[2] + from ( + select transform_values( + (k, v) -> v + 1, map(1, 10, 2, 20)) r + ) t + """ + + qt_materialized_computed_map """ + select r[1], r[3] + from ( + select transform_values((k, v) -> v + 1, computed_map) r + from ( + select map(1, mii[1], 3, 30) computed_map + from test_map_lambda where id = 1 + ) producer + ) consumer + """ + + // Deterministic computed Maps compose directly with all Map Lambda functions. + qt_direct_computed_map """ + select r[1], r[3] + from ( + select transform_values( + (k, v) -> v + 1, map(1, mii[1], 3, 30)) r + from test_map_lambda where id = 1 + ) t + """ + + qt_all_functions_direct_computed_map """ + select map_apply( + (k, v) -> struct(k, v + 1), + map(1, mii[1], 3, 30))[1], + map_filter( + (k, v) -> k = 3, + map(1, mii[1], 3, 30))[3], + transform_keys( + (k, v) -> k + 1, + map(1, mii[1], 3, 30))[2], + transform_values( + (k, v) -> v + 1, + map(1, mii[1], 3, 30))[3], + cast(map_exists( + (k, v) -> v = 30, + map(1, mii[1], 3, 30)) as int), + cast(map_all( + (k, v) -> v >= 10, + map(1, mii[1], 3, 30)) as int) + from test_map_lambda where id = 1 + """ + + order_qt_map_lambda_aggregate_output """ + select id, count(*) as row_count, + transform_values( + (k, v) -> v + 1, + map(1, id, 2, id + 10))[2] as transformed_value + from test_map_lambda + group by id + order by id + """ + + order_qt_map_lambda_having """ + select id, count(*) as row_count + from test_map_lambda + group by id + having map_exists( + (k, v) -> true, + map(1, id, 2, id + 10)) + order by id + """ + + order_qt_map_lambda_join_on """ + select l.id, r.id + from test_map_lambda l + join test_map_lambda r + on transform_values( + (k, v) -> v, + map(1, l.id, 2, l.id + 1))[2] = r.id + order by l.id, r.id + """ + + qt_nondeterministic_map_input """ + select map_size(r), map_values(r)[1] + from ( + select transform_values( + (k, v) -> v, map(cast(random() as int), 10)) r + ) t + """ + + qt_nondeterministic_map_single_evaluation """ + select cast(map_keys(r)[1] = map_values(r)[1] as int) + from ( + select transform_values( + (k, v) -> k, map(random(1, 1000000000), 10)) r + ) t + """ + + qt_map_lambda_generate_volatile """ + select cast(k = v as int) + from (select 1 as seed) t + lateral view explode_map( + transform_values( + (mk, mv) -> mk, + map(random(1, 1000000000), 0))) tmp as k, v + """ + + qt_nested_nondeterministic_map_single_evaluation """ + select map_size(r), cast(map_keys(r)[1] = map_values(r)[1] as int) + from ( + select transform_values( + (k, v) -> v, + transform_values( + (ik, iv) -> ik, + map(random(1, 1000000000), 0))) r + ) t + """ + + qt_lambda_nested_lambda_map_single_evaluation """ + select cast(map_keys(r[1])[1] = map_values(r[1])[1] as int) + from ( + select transform_values( + (ok, ov) -> transform_values( + (ik, iv) -> ik, + map(ok + random(1, 1000000000), ov)), + map(1, 10)) r + ) t + """ + + qt_array_map_nested_map_lambda_single_evaluation """ + select cast( + map_keys(r[1])[1] = map_values(r[1])[1] + and map_keys(r[2])[1] = map_values(r[2])[1] + and map_keys(r[1])[1] != map_keys(r[2])[1] + as int) + from ( + select array_map( + x -> transform_values((k, v) -> k, map(uuid(), x)), + [1, 2]) r + ) t + """ + + qt_array_map_nested_map_apply_single_evaluation """ + select cast( + map_keys(r[1])[1] = map_values(r[1])[1] + and map_keys(r[2])[1] = map_values(r[2])[1] + and map_keys(r[1])[1] != map_keys(r[2])[1] + as int) + from ( + select array_map( + x -> map_apply((k, v) -> struct(k, k), map(uuid(), x)), + [1, 2]) r + ) t + """ + + test { + sql "select map_filter(k -> k > 0, map(1, 10))" + exception "requires exactly two arguments" + } + test { + sql "select map_from_arrays([[1]], [10])" + exception "MAP key type must be a primitive type" + } + qt_map_from_arrays_nested_empty_array "select map_from_arrays([1], [[]])" + test { + sql "select map_from_entries(1)" + exception "requires an array of structs with exactly two fields" + } + test { + sql "select map_from_entries(array(struct(1)))" + exception "requires an array of structs with exactly two fields" + } + test { + sql "select map_from_entries(array(cast(null as struct)))" + exception "Map entry of function map_from_entries cannot be null" + } + test { + sql """ + select map_apply( + (k, v) -> if(k > 0, struct(k, v), cast(null as struct)), + map(1, 10)) + """ + exception "must return a non-nullable struct with exactly two fields" + } +}