From 1314dbe65916ff6dc0b73366dbe0bad156d384bd Mon Sep 17 00:00:00 2001 From: linzhenqi Date: Wed, 19 Aug 2026 17:15:31 +0800 Subject: [PATCH 1/4] [tmp] --- be/src/exprs/function/function_map.cpp | 113 +++ .../lambda_function/varray_map_function.cpp | 12 +- .../array_map_function_test.cpp | 54 +- .../doris/catalog/BuiltinScalarFunctions.java | 12 + .../glue/translator/ExpressionTranslator.java | 10 + .../doris/nereids/jobs/executor/Rewriter.java | 9 +- .../nereids/parser/LogicalPlanBuilder.java | 6 +- .../apache/doris/nereids/rules/RuleType.java | 1 + .../rewrite/AddProjectForMapLambdaInput.java | 786 ++++++++++++++++++ .../expressions/functions/scalar/Lambda.java | 48 +- .../expressions/functions/scalar/MapAll.java | 86 ++ .../functions/scalar/MapApply.java | 154 ++++ .../functions/scalar/MapEntryArrayMap.java | 38 + .../functions/scalar/MapExists.java | 86 ++ .../functions/scalar/MapFilter.java | 99 +++ .../functions/scalar/MapFromArraysUnique.java | 42 + .../functions/scalar/MapLambdaValidator.java | 168 ++++ .../functions/scalar/TransformKeys.java | 82 ++ .../functions/scalar/TransformValues.java | 94 +++ .../visitor/ScalarFunctionVisitor.java | 30 + .../apache/doris/nereids/util/PlanUtils.java | 19 +- .../AddProjectForMapLambdaInputTest.java | 427 ++++++++++ .../scalar/MapLambdaFunctionsTest.java | 344 ++++++++ .../org/apache/doris/nereids/DorisParser.g4 | 38 +- .../map_functions/test_map_lambda.out | 141 ++++ .../map_functions/test_map_lambda.groovy | 449 ++++++++++ 26 files changed, 3297 insertions(+), 51 deletions(-) create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInput.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapAll.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapApply.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapEntryArrayMap.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapExists.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFilter.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromArraysUnique.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaValidator.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TransformKeys.java create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TransformValues.java create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInputTest.java create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaFunctionsTest.java create mode 100644 regression-test/data/query_p0/sql_functions/map_functions/test_map_lambda.out create mode 100644 regression-test/suites/query_p0/sql_functions/map_functions/test_map_lambda.groovy diff --git a/be/src/exprs/function/function_map.cpp b/be/src/exprs/function/function_map.cpp index bd1b8281904607..9615a5523f8214 100644 --- a/be/src/exprs/function/function_map.cpp +++ b/be/src/exprs/function/function_map.cpp @@ -181,6 +181,118 @@ 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); + ColumnPtr map_column = unpacked_map_column; + ColumnPtr predicate_column = unpacked_predicate_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 = [&](ColumnPtr& column, bool is_const) { + 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); + column = nullable->get_nested_column_ptr(); + } + }; + merge_null_map(map_column, map_is_const); + merge_null_map(predicate_column, predicate_is_const); + + const auto& map = assert_cast(*map_column); + const auto& predicate = assert_cast(*predicate_column); + if ((!map_is_const && map.size() != input_rows_count) || + (!predicate_is_const && predicate.size() != input_rows_count)) { + return Status::InvalidArgument( + "The map and lambda result offsets of function {} must be identical", name); + } + for (size_t row = 0; row < input_rows_count; ++row) { + if (!result_null_map_data[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); + } + } + + const IColumn* predicate_data = &predicate.get_data(); + const UInt8* predicate_null_map = nullptr; + if (const auto* nullable_predicate = check_and_get_column(predicate_data)) { + predicate_null_map = nullable_predicate->get_null_map_data().data(); + predicate_data = &nullable_predicate->get_nested_column(); + } + const auto& predicate_values = assert_cast(*predicate_data).get_data(); + + IColumn::Selector selector; + if (!map_is_const) { + selector.reserve(map.get_keys().size()); + } + auto result_keys = map.get_keys().clone_empty(); + auto result_values = map.get_values().clone_empty(); + auto result_offsets = ColumnArray::ColumnOffsets::create(); + result_offsets->reserve(input_rows_count); + + size_t output_offset = 0; + for (size_t row = 0; row < input_rows_count; ++row) { + if (!result_null_map_data[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 == nullptr || + predicate_null_map[predicate_entry] == 0) && + predicate_values[predicate_entry] != 0; + if (selected) { + const size_t map_entry = map_begin + entry; + selector.push_back(map_entry); + ++output_offset; + } + } + } + result_offsets->insert_value(output_offset); + } + if (map_is_const && !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()); + } else if (!map_is_const) { + map.get_keys().append_data_by_selector(result_keys, selector); + map.get_values().append_data_by_selector(result_values, selector); + } + 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(); + } +}; + // construct a map // map(key1, value2, key2, value2) -> {key1: value2, key2: value2} class FunctionMap : public IFunction { @@ -985,6 +1097,7 @@ 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>(); 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/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/glue/translator/ExpressionTranslator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java index 84aea35f2c5001..5dd5fc7d4b5fa5 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java @@ -100,6 +100,9 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.DictGetMany; 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.MapEntryArrayMap; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapFilter; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapLambdaValidator; import org.apache.doris.nereids.trees.expressions.functions.scalar.ScalarFunction; import org.apache.doris.nereids.trees.expressions.functions.udf.JavaUdaf; import org.apache.doris.nereids.trees.expressions.functions.udf.JavaUdf; @@ -532,6 +535,9 @@ public Expr visitLambda(Lambda lambda, PlanTranslatorContext context) { @Override public Expr visitArrayMap(ArrayMap arrayMap, PlanTranslatorContext context) { Lambda lambda = (Lambda) arrayMap.child(0); + if (arrayMap instanceof MapEntryArrayMap) { + MapLambdaValidator.validateStablePhysicalInputs(arrayMap.getName(), lambda); + } List arguments = new ArrayList<>(arrayMap.children().size()); arguments.add(null); int columnId = 0; @@ -721,6 +727,10 @@ public Expr visitSearchExpression(SearchExpression searchExpression, @Override public Expr visitScalarFunction(ScalarFunction function, PlanTranslatorContext context) { + if (function instanceof MapFilter && ((MapFilter) function).shouldValidateMapLambdaInput()) { + MapLambdaValidator.validateOuterMapConsumer( + function.getName(), function.getArgument(1)); + } List arguments = function.getArguments().stream() .map(arg -> arg.accept(this, context)) .collect(Collectors.toList()); 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..0dc1ca7755e2a9 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 @@ -34,6 +34,7 @@ import org.apache.doris.nereids.rules.expression.QueryColumnCollector; import org.apache.doris.nereids.rules.rewrite.AddDefaultLimit; import org.apache.doris.nereids.rules.rewrite.AddProjectForJoin; +import org.apache.doris.nereids.rules.rewrite.AddProjectForMapLambdaInput; import org.apache.doris.nereids.rules.rewrite.AddProjectForVolatileExpression; import org.apache.doris.nereids.rules.rewrite.AdjustConjunctsReturnType; import org.apache.doris.nereids.rules.rewrite.AdjustNullable; @@ -762,9 +763,11 @@ public class Rewriter extends AbstractBatchJobExecutor { topDown(new SumLiteralRewrite(), new MergePercentileToArray()) ), - topic("add projection for volatile expression", - // separate AddProjectForVolatileExpression and MergeProjectable - // to avoid dead loop if code has bug + topic("add projection for volatile and lambda expression", + // Materialize Map Lambda inputs and volatile expressions before merging the + // generated Projects. Separate passes avoid repeatedly adding and merging the + // same Project in one top-down rewrite. + topDown(new AddProjectForMapLambdaInput()), 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/RuleType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java index a981dc496813b6..13da88ef23a299 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java @@ -228,6 +228,7 @@ public enum RuleType { PUSH_DOWN_MAX_MIN_FILTER(RuleTypeClass.REWRITE), ADD_PROJECT_FOR_JOIN(RuleTypeClass.REWRITE), + ADD_PROJECT_FOR_MAP_LAMBDA_INPUT(RuleTypeClass.REWRITE), ADD_PROJECT_FOR_VOLATILE_EXPRESSION(RuleTypeClass.REWRITE), VARIANT_SUB_PATH_PRUNING(RuleTypeClass.REWRITE), NESTED_COLUMN_PRUNING(RuleTypeClass.REWRITE), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInput.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInput.java new file mode 100644 index 00000000000000..6bc611215aab97 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInput.java @@ -0,0 +1,786 @@ +// 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.rules.rewrite; + +import org.apache.doris.common.Pair; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.ArrayItemReference; +import org.apache.doris.nereids.trees.expressions.ArrayItemReference.ArrayItemSlot; +import org.apache.doris.nereids.trees.expressions.Cast; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; +import org.apache.doris.nereids.trees.expressions.functions.Function; +import org.apache.doris.nereids.trees.expressions.functions.scalar.ArrayMap; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Lambda; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapEntryArrayMap; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapLambdaValidator; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalGenerate; +import org.apache.doris.nereids.trees.plans.logical.LogicalHaving; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalOneRowRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.util.ExpressionUtils; +import org.apache.doris.nereids.util.JoinUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; +import com.google.common.collect.Sets; + +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.Optional; +import java.util.Set; + +/** + * Materialize computed Map inputs used by {@link MapEntryArrayMap}. + * + *

A Map entry lambda takes {@code map_keys(computedMap)} and + * {@code map_values(computedMap)} as its two input arrays. rule evaThisluates + * {@code computedMap} in a child Project and replaces all its occurrences with the same Slot: + * + * 

+ * before:
+ *   Project[map_from_arrays(
+ *     map_keys(computedMap),
+ *     MapEntryArrayMap(
+ *       (mapKey, mapValue) -> valueExpression,
+ *       map_keys(computedMap), map_values(computedMap)))]
+ *     child
+ *
+ * after:
+ *   Project[map_from_arrays(
+ *     map_keys(materializedMapSlot),
+ *     MapEntryArrayMap(
+ *       (mapKey, mapValue) -> valueExpression,
+ *       map_keys(materializedMapSlot), map_values(materializedMapSlot)))]
+ *     Project[child.*, computedMap AS materializedMapSlot]
+ *       child
+ *
+ * + *

Besides the basic rewrite above, this rule handles + * repeated entry arrays, nested lambdas, and Join children through dedicated helper methods below. + */ +public class AddProjectForMapLambdaInput implements RewriteRuleFactory { + + @Override + public List buildRules() { + return ImmutableList.of( + new GenerateRewrite().build(), + new OneRowRelationRewrite().build(), + new ProjectRewrite().build(), + new FilterRewrite().build(), + new HavingRewrite().build(), + new AggregateRewrite().build(), + new JoinRewrite().build() + ); + } + + private class GenerateRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalGenerate().thenApply(ctx -> { + LogicalGenerate generate = ctx.root; + List generators = materializeNestedMapInputs(generate.getGenerators()); + Optional, LogicalProject>> + rewrittenOpt = rewriteExpressions(generate, generators); + if (rewrittenOpt.isPresent()) { + return generate.withGenerators(rewrittenOpt.get().first) + .withChildren(rewrittenOpt.get().second); + } else if (!generators.equals(generate.getGenerators())) { + return generate.withGenerators(generators); + } else { + return generate; + } + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class OneRowRelationRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalOneRowRelation().thenApply(ctx -> { + LogicalOneRowRelation oneRowRelation = ctx.root; + List projects = materializeNestedMapInputs(oneRowRelation.getProjects()); + List mapInputAliases = tryGenMapInputAliases(projects); + List rewrittenProjects = replaceExpressions(projects, mapInputAliases); + List entryArrayAliases = tryGenSharedEntryArrayAliases(rewrittenProjects); + if (mapInputAliases.isEmpty() && entryArrayAliases.isEmpty()) { + return projects.equals(oneRowRelation.getProjects()) + ? oneRowRelation : oneRowRelation.withProjects(projects); + } + + // A OneRowRelation has no child on which to install the usual materialization + // Project. Use the relation itself as the lowest projection, then stack the shared + // entry-array Project and the original output Project above it. + Plan child; + if (mapInputAliases.isEmpty()) { + child = oneRowRelation.withProjects(entryArrayAliases); + } else { + child = oneRowRelation.withProjects(mapInputAliases); + if (!entryArrayAliases.isEmpty()) { + child = appendProject(child, entryArrayAliases); + } + } + rewrittenProjects = replaceExpressions(rewrittenProjects, entryArrayAliases); + return new LogicalProject<>(rewrittenProjects, child); + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class ProjectRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalProject().thenApply(ctx -> { + LogicalProject project = ctx.root; + List projects = materializeNestedMapInputs(project.getProjects()); + Optional, LogicalProject>> + rewrittenOpt = rewriteExpressions(project, projects); + if (rewrittenOpt.isPresent()) { + return project.withProjectsAndChild(rewrittenOpt.get().first, rewrittenOpt.get().second); + } else if (!projects.equals(project.getProjects())) { + return project.withProjects(projects); + } else { + return project; + } + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class FilterRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalFilter().thenApply(ctx -> { + LogicalFilter filter = ctx.root; + List conjuncts = materializeNestedMapInputs(filter.getConjuncts()); + Optional, LogicalProject>> + rewrittenOpt = rewriteExpressions(filter, conjuncts); + if (rewrittenOpt.isPresent()) { + return filter.withConjunctsAndChild( + ImmutableSet.copyOf(rewrittenOpt.get().first), + rewrittenOpt.get().second); + } else if (!ImmutableSet.copyOf(conjuncts).equals(filter.getConjuncts())) { + return filter.withConjuncts(ImmutableSet.copyOf(conjuncts)); + } else { + return filter; + } + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class HavingRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalHaving().thenApply(ctx -> { + LogicalHaving having = ctx.root; + List conjuncts = materializeNestedMapInputs(having.getConjuncts()); + Optional, LogicalProject>> + rewrittenOpt = rewriteExpressions(having, conjuncts); + if (rewrittenOpt.isPresent()) { + return having.withConjuncts(ImmutableSet.copyOf(rewrittenOpt.get().first)) + .withChildren(rewrittenOpt.get().second); + } else if (!ImmutableSet.copyOf(conjuncts).equals(having.getConjuncts())) { + return having.withConjuncts(ImmutableSet.copyOf(conjuncts)); + } else { + return having; + } + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class AggregateRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalAggregate().thenApply(ctx -> { + LogicalAggregate aggregate = ctx.root; + List originalTargets = Lists.newArrayList(); + originalTargets.addAll(aggregate.getGroupByExpressions()); + originalTargets.addAll(aggregate.getOutputExpressions()); + List targets = materializeNestedMapInputs(originalTargets); + Optional, LogicalProject>> rewrittenOpt + = rewriteExpressions(aggregate, targets); + Plan newChild = rewrittenOpt.isPresent() + ? rewrittenOpt.get().second : aggregate.child(); + List newTargets = rewrittenOpt.isPresent() + ? rewrittenOpt.get().first : targets; + if (!rewrittenOpt.isPresent() && newTargets.equals(originalTargets)) { + return aggregate; + } + // rewriteExpressions treats group-by expressions and outputs as one ordered list + // so a common Map input is materialized only once. Restore the two original lists + // after replacement. + int groupBySize = aggregate.getGroupByExpressions().size(); + ImmutableList newGroupBy = ImmutableList.copyOf( + newTargets.subList(0, groupBySize)); + ImmutableList.Builder newOutputBuilder + = ImmutableList.builderWithExpectedSize(aggregate.getOutputExpressions().size()); + for (int i = groupBySize; i < newTargets.size(); i++) { + newOutputBuilder.add((NamedExpression) newTargets.get(i)); + } + return aggregate.withChildGroupByAndOutput(newGroupBy, newOutputBuilder.build(), newChild); + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class JoinRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalJoin().thenApply(ctx -> { + LogicalJoin join = ctx.root; + int hashOtherConjunctsSize = join.getHashJoinConjuncts().size() + + join.getOtherJoinConjuncts().size(); + int totalConjunctsSize = hashOtherConjunctsSize + join.getMarkJoinConjuncts().size(); + List allConjuncts = Lists.newArrayListWithExpectedSize(totalConjunctsSize); + allConjuncts.addAll(join.getHashJoinConjuncts()); + allConjuncts.addAll(join.getOtherJoinConjuncts()); + allConjuncts.addAll(join.getMarkJoinConjuncts()); + List originalAllConjuncts = ImmutableList.copyOf(allConjuncts); + allConjuncts = materializeNestedMapInputs(allConjuncts); + Optional rewrittenOpt = rewriteJoinExpressions(join, allConjuncts); + if (!rewrittenOpt.isPresent() && allConjuncts.equals(originalAllConjuncts)) { + return join; + } + + Plan newLeftChild = rewrittenOpt.map(result -> result.left).orElse(join.left()); + Plan newRightChild = rewrittenOpt.map(result -> result.right).orElse(join.right()); + List newAllConjuncts = rewrittenOpt + .map(result -> result.newConjuncts).orElse(allConjuncts); + List newHashOtherConjuncts = newAllConjuncts.subList(0, hashOtherConjunctsSize); + List newMarkJoinConjuncts = ImmutableList.copyOf( + newAllConjuncts.subList(hashOtherConjunctsSize, totalConjunctsSize)); + + Pair, List> pair = JoinUtils.extractExpressionForHashTable( + newLeftChild.getOutput(), newRightChild.getOutput(), newHashOtherConjuncts); + List newHashJoinConjuncts = pair.first; + List newOtherJoinConjuncts = pair.second; + JoinType joinType = join.getJoinType(); + if (joinType == JoinType.CROSS_JOIN && !newHashJoinConjuncts.isEmpty()) { + joinType = JoinType.INNER_JOIN; + } + return new LogicalJoin<>(joinType, + newHashJoinConjuncts, + newOtherJoinConjuncts, + newMarkJoinConjuncts, + join.getDistributeHint(), + join.getMarkJoinSlotReference(), + ImmutableList.of(newLeftChild, newRightChild), + join.getJoinReorderContext()); + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + /** + * Rewrite expressions owned by a single-child plan and install their materialization Projects. + * + *

It first materializes computed Map inputs and replaces them in {@code targets}. It then + * materializes any {@link MapEntryArrayMap} still used more than once. These are separate + * Project layers because the second expression can depend on a Map Slot created by the first. + * The returned pair contains the rewritten targets and the top materialization Project. + */ + private Optional, LogicalProject>> rewriteExpressions( + LogicalPlan plan, Collection targets) { + // computed map materialized + List mapInputAliases = tryGenMapInputAliases(targets); + List rewrittenTargets = replaceExpressions(targets, mapInputAliases); + // MapEntryArrayMap merteialized + List entryArrayAliases = tryGenSharedEntryArrayAliases(rewrittenTargets); + if (mapInputAliases.isEmpty() && entryArrayAliases.isEmpty()) { + return Optional.empty(); + } + + Plan child = plan.child(0); + if (!mapInputAliases.isEmpty()) { + child = appendProject(child, mapInputAliases); + } + if (!entryArrayAliases.isEmpty()) { + child = appendProject(child, entryArrayAliases); + rewrittenTargets = replaceExpressions(rewrittenTargets, entryArrayAliases); + } + + return Optional.of(Pair.of(rewrittenTargets, (LogicalProject) child)); + } + + /** Add aliases without hiding any output already produced by {@code child}. */ + private LogicalProject appendProject(Plan child, List aliases) { + List projects = ImmutableList.builder() + .addAll(child.getOutput()) + .addAll(aliases) + .build(); + return new LogicalProject<>(projects, child); + } + + /** Replace each aliased expression by its Slot in all target expression trees. */ + private List replaceExpressions( + Collection expressions, List aliases) { + if (aliases.isEmpty()) { + return ImmutableList.copyOf(expressions); + } + Map replaceMap = Maps.newHashMap(); + for (NamedExpression alias : aliases) { + replaceMap.put(alias.child(0), alias.toSlot()); + } + ImmutableList.Builder builder = ImmutableList.builderWithExpectedSize(expressions.size()); + for (T expression : expressions) { + builder.add((T) ExpressionUtils.replace(expression, replaceMap)); + } + return builder.build(); + } + + /** + * Rewrite Join conjuncts using the same two materialization stages as + * {@link #rewriteExpressions(LogicalPlan, Collection)}. + * + *

Unlike a single-child plan, each generated alias must be attached to the Join child that + * contains all its input Slots. An expression referencing both children cannot be evaluated in + * either child Project, so a deterministic expression is left unchanged and a volatile one is + * rejected. Entry-array aliases are assigned after Map aliases because they may use new Slots. + */ + private Optional rewriteJoinExpressions(LogicalJoin join, + Collection targets) { + List rewrittenTargets = ImmutableList.copyOf(targets); + Plan left = join.left(); + Plan right = join.right(); + + Map> mapInputSlots = Maps.newLinkedHashMap(); + for (Expression target : rewrittenTargets) { + Set mapInputs = Sets.newLinkedHashSet(); + collectMapInputs(target, mapInputs); + for (Expression mapInput : mapInputs) { + Set inputSlots = mapInput.getInputSlots(); + mapInputSlots.computeIfAbsent(mapInput, ignored -> Sets.newLinkedHashSet()) + .addAll(inputSlots.isEmpty() ? target.getInputSlots() : inputSlots); + } + } + + ImmutableList.Builder leftAliases = ImmutableList.builder(); + ImmutableList.Builder rightAliases = ImmutableList.builder(); + Map replaceMap = Maps.newHashMap(); + Set leftOutputSet = left.getOutputSet(); + Set rightOutputSet = right.getOutputSet(); + for (Entry> entry : mapInputSlots.entrySet()) { + Set inputSlots = entry.getValue(); + Set mapInputExpressionSlots = entry.getKey().getInputSlots(); + if (!mapInputExpressionSlots.isEmpty() + && !leftOutputSet.containsAll(inputSlots) + && !rightOutputSet.containsAll(inputSlots)) { + // No child Project can reference Slots from both sides. Recalculation is safe for + // a deterministic expression, but a volatile Map would no longer have one stable + // value shared by map_keys and map_values. + if (entry.getKey().containsVolatileExpression()) { + throw new AnalysisException( + "A computed Map input containing a volatile expression cannot " + + "reference both sides of a join"); + } + continue; + } + ExprId exprId = StatementScopeIdGenerator.newExprId(); + Alias alias = new Alias( + exprId, entry.getKey(), "$_map_input_" + exprId.asInt() + "_$"); + replaceMap.put(alias.child(0), alias.toSlot()); + if (!inputSlots.isEmpty() && rightOutputSet.containsAll(inputSlots)) { + rightAliases.add(alias); + } else { + leftAliases.add(alias); + } + } + if (!replaceMap.isEmpty()) { + List leftAliasList = leftAliases.build(); + List rightAliasList = rightAliases.build(); + left = appendProjectIfNeeded(left, leftAliasList); + right = appendProjectIfNeeded(right, rightAliasList); + rewrittenTargets = replaceExpressions(rewrittenTargets, + ImmutableList.builder() + .addAll(leftAliasList) + .addAll(rightAliasList) + .build()); + } + + List entryArrayAliases = tryGenSharedEntryArrayAliases(rewrittenTargets); + ImmutableList.Builder leftEntryAliases = ImmutableList.builder(); + ImmutableList.Builder rightEntryAliases = ImmutableList.builder(); + leftOutputSet = left.getOutputSet(); + rightOutputSet = right.getOutputSet(); + for (NamedExpression alias : entryArrayAliases) { + Expression entryArray = alias.child(0); + Set inputSlots = Sets.newLinkedHashSet(entryArray.getInputSlots()); + if (inputSlots.isEmpty()) { + // As with a slot-free Map, inherit the containing conjunct's scope only to choose + // a child. The expression itself remains valid on either side. + for (Expression target : rewrittenTargets) { + if (target.anyMatch(entryArray::equals)) { + inputSlots.addAll(target.getInputSlots()); + } + } + } + Set expressionSlots = entryArray.getInputSlots(); + if (!expressionSlots.isEmpty() + && !leftOutputSet.containsAll(inputSlots) + && !rightOutputSet.containsAll(inputSlots)) { + if (entryArray.containsVolatileExpression()) { + throw new AnalysisException( + "A shared Map entry array containing a volatile expression cannot " + + "reference both sides of a join"); + } + continue; + } + if (!inputSlots.isEmpty() && rightOutputSet.containsAll(inputSlots)) { + rightEntryAliases.add(alias); + } else { + leftEntryAliases.add(alias); + } + } + List leftEntryAliasList = leftEntryAliases.build(); + List rightEntryAliasList = rightEntryAliases.build(); + if (!leftEntryAliasList.isEmpty() || !rightEntryAliasList.isEmpty()) { + left = appendProjectIfNeeded(left, leftEntryAliasList); + right = appendProjectIfNeeded(right, rightEntryAliasList); + rewrittenTargets = replaceExpressions(rewrittenTargets, + ImmutableList.builder() + .addAll(leftEntryAliasList) + .addAll(rightEntryAliasList) + .build()); + } + + if (replaceMap.isEmpty() && leftEntryAliasList.isEmpty() && rightEntryAliasList.isEmpty()) { + return Optional.empty(); + } + return Optional.of(new JoinRewriteResult(rewrittenTargets, left, right)); + } + + /** Avoid creating an identity Project when one side of a Join has no aliases. */ + private Plan appendProjectIfNeeded(Plan child, List aliases) { + if (aliases.isEmpty()) { + return child; + } + List projects = ImmutableList.builder() + .addAll(child.getOutput()) + .addAll(aliases) + .build(); + return new LogicalProject<>(projects, child); + } + + /** + * Find and alias each distinct computed Map consumed by a {@link MapEntryArrayMap}. + * + *

This method turns the expressions found by {@link #collectMapInputs(Expression, Set)} into + * aliases. Slots and Map literals are excluded because they need no materialization. + */ + private List tryGenMapInputAliases( + Collection targets) { + Set mapInputs = Sets.newLinkedHashSet(); + for (Expression target : targets) { + collectMapInputs(target, mapInputs); + } + + ImmutableList.Builder aliases + = ImmutableList.builderWithExpectedSize(mapInputs.size()); + for (Expression mapInput : mapInputs) { + ExprId exprId = StatementScopeIdGenerator.newExprId(); + aliases.add(new Alias(exprId, mapInput, "$_map_input_" + exprId.asInt() + "_$")); + } + return aliases.build(); + } + + /** + * Find repeated {@link MapEntryArrayMap} expressions and create one shared alias for each. + * + *

This is used by the current safe lowering of {@code map_apply}; the implementation does + * not use the optional fast lowering into two independent two-parameter ArrayMaps. The original + * two-parameter lambda is evaluated first and produces {@code ARRAY<STRUCT>}: + * + * 

+     * mappedEntries = MapEntryArrayMap(
+     *   (mapKey, mapValue) -> struct(newKey, newValue),
+     *   map_keys(inputMap), map_values(inputMap))
+     * map_from_arrays(
+     *   array_map(mappedEntry -> mappedEntry[1], mappedEntries),
+     *   array_map(mappedEntry -> mappedEntry[2], mappedEntries))
+     *
+ * + *

The two extraction ArrayMaps have one parameter because they iterate the resulting Struct + * array, not the original Map. They do not copy or reevaluate the original lambda body. + */ + private List tryGenSharedEntryArrayAliases( + Collection targets) { + Map entryArrayCounts = Maps.newLinkedHashMap(); + for (Expression target : targets) { + collectEntryArrayCounts(target, entryArrayCounts); + } + + ImmutableList.Builder aliases = ImmutableList.builder(); + for (Entry entry : entryArrayCounts.entrySet()) { + if (entry.getValue() > 1) { + ExprId exprId = StatementScopeIdGenerator.newExprId(); + aliases.add(new Alias( + exprId, entry.getKey(), "$_map_entries_" + exprId.asInt() + "_$")); + } + } + return aliases.build(); + } + + /** Apply nested-lambda materialization independently to every target expression. */ + private List materializeNestedMapInputs(Collection expressions) { + ImmutableList.Builder builder = ImmutableList.builderWithExpectedSize(expressions.size()); + for (T expression : expressions) { + builder.add((T) materializeNestedMapInputs(expression)); + } + return builder.build(); + } + + /** + * Materialize computed Maps that depend on lambda item Slots inside the owning ArrayMap. + * + *

Consider: + * + * 

+     * select transform_values(
+     *   (outer_k, outer_v) -> transform_values((inner_k, inner_v) -> inner_k, map(outer_k + random(), outer_v)),
+     *   map(1, 10));
+     *
+ * + * A relation Project cannot evaluate {@code map(outer_k + random(), outer_v)} because {@code outer_k} and + * {@code outer_v} exist only while the outer lambda is running. The outer ArrayMap is rewritten to + * carry a hidden array whose item is that Map: + * + * 
+     * outer inputs before:
+     *   outer_k <- map_keys(outerMap)
+     *   outer_v <- map_values(outerMap)
+     *
+     * outer inputs after:
+     *   outer_k <- map_keys(outerMap)
+     *   outer_v <- map_values(outerMap)
+     *   materializedInnerMap
+     *      - array_map((outerKey, outerValue) -> map(outerKey + random(), outerValue),
+     *                   map_keys(outerMap), map_values(outerMap))
+     *
+     * outer body after:
+     *   transform_values((innerKey, innerValue) -> innerKey, materializedInnerMap)
+     *
+ * + *

Traversal is bottom-up. For each ArrayMap, computed Maps in its body become hidden input + * arrays; repeated entry arrays are handled afterward because they may use those hidden inputs. + */ + private Expression materializeNestedMapInputs(Expression expression) { + ImmutableList.Builder children + = ImmutableList.builderWithExpectedSize(expression.arity()); + boolean changed = false; + for (Expression child : expression.children()) { + Expression rewrittenChild = materializeNestedMapInputs(child); + children.add(rewrittenChild); + changed |= rewrittenChild != child; + } + Expression rewritten = changed ? expression.withChildren(children.build()) : expression; + if (!(rewritten instanceof ArrayMap)) { + return rewritten; + } + + Lambda lambda = (Lambda) rewritten.child(0); + Set mapInputs = Sets.newLinkedHashSet(); + collectMapInputs(lambda.getLambdaFunction(), mapInputs); + + List sourceArguments = lambda.getLambdaArguments(); + List argumentNames = Lists.newArrayList(lambda.getLambdaArgumentNames()); + List arguments = Lists.newArrayList(sourceArguments); + Expression lambdaBody = lambda.getLambdaFunction(); + for (Expression mapInput : mapInputs) { + Pair materialized = buildLambdaMaterializer(mapInput, sourceArguments); + ArrayItemReference hiddenArgument = new ArrayItemReference(materialized.second, materialized.first); + argumentNames.add(materialized.second); + arguments.add(hiddenArgument); + Map replaceMap = Maps.newHashMap(); + replaceMap.put(mapInput, hiddenArgument.toSlot()); + lambdaBody = ExpressionUtils.replace(lambdaBody, replaceMap); + } + + List entryArrayAliases = tryGenSharedEntryArrayAliases( + ImmutableList.of(lambdaBody)); + for (NamedExpression entryArrayAlias : entryArrayAliases) { + Expression entryArray = entryArrayAlias.child(0); + Pair materialized = buildLambdaMaterializer(entryArray, arguments); + ArrayItemReference hiddenArgument = new ArrayItemReference(materialized.second, materialized.first); + argumentNames.add(materialized.second); + arguments.add(hiddenArgument); + Map replaceMap = Maps.newHashMap(); + replaceMap.put(entryArray, hiddenArgument.toSlot()); + lambdaBody = ExpressionUtils.replace(lambdaBody, replaceMap); + } + if (arguments.size() == sourceArguments.size()) { + return rewritten; + } + + // A shared entry-array materializer can embed an earlier Map materializer. In that case the + // final body references only the entry-array argument. Keep all user arguments, but remove + // optimizer-added arguments no longer referenced by the final body to avoid evaluating the + // embedded Map expression a second time. + Set referencedArgumentIds = collectReferencedArgumentIds(lambdaBody); + ImmutableList.Builder retainedNames = ImmutableList.builder(); + ImmutableList.Builder retainedArguments = ImmutableList.builder(); + for (int i = 0; i < arguments.size(); i++) { + ArrayItemReference argument = arguments.get(i); + if (i < sourceArguments.size() + || referencedArgumentIds.contains(argument.getExprId())) { + retainedNames.add(argumentNames.get(i)); + retainedArguments.add(argument); + } + } + return rewritten.withChildren(ImmutableList.of( + new Lambda(retainedNames.build(), lambdaBody, retainedArguments.build()))); + } + + /** + * Build an ArrayMap that evaluates {@code expression} once per entry of the enclosing lambda. + * + *

Only enclosing arguments referenced by the expression are forwarded. For + * {@code map(ok + random(), ov)}, the generated lambda receives copies of {@code ok} and + * {@code ov}, with fresh ExprIds, and its body is rebound to those copies. If the expression + * only captures relation Slots, one enclosing array is still forwarded as a row-count and + * offset driver; all arrays of one ArrayMap have identical entry offsets. + * + * @return the materializing ArrayMap and the name of the hidden item argument that will expose + * each materialized result to the original lambda body + */ + private Pair buildLambdaMaterializer( + Expression expression, List sourceArguments) { + Set referencedArgumentIds = collectReferencedArgumentIds(expression); + + List selectedArguments = sourceArguments.stream() + .filter(argument -> referencedArgumentIds.contains(argument.getExprId())) + .collect(ImmutableList.toImmutableList()); + if (selectedArguments.isEmpty()) { + // ArrayMap needs an array to define the entry count even when the expression only + // captures relation slots. Any current lambda input has the same entry offsets. + selectedArguments = ImmutableList.of(sourceArguments.get(0)); + } + + Map replaceMap = Maps.newHashMap(); + ImmutableList.Builder materializerNames + = ImmutableList.builderWithExpectedSize(selectedArguments.size()); + ImmutableList.Builder materializerArguments + = ImmutableList.builderWithExpectedSize(selectedArguments.size()); + for (ArrayItemReference sourceArgument : selectedArguments) { + ExprId exprId = StatementScopeIdGenerator.newExprId(); + String name = "$_map_materialize_arg_" + exprId.asInt() + "_$"; + ArrayItemReference materializerArgument = new ArrayItemReference( + exprId, name, sourceArgument.getArrayExpression()); + materializerNames.add(name); + materializerArguments.add(materializerArgument); + replaceMap.put(sourceArgument.toSlot(), materializerArgument.toSlot()); + } + + Expression materializerBody = ExpressionUtils.replace(expression, replaceMap); + Lambda materializerLambda = new Lambda( + materializerNames.build(), materializerBody, materializerArguments.build()); + ExprId hiddenExprId = StatementScopeIdGenerator.newExprId(); + String hiddenName = "$_map_input_" + hiddenExprId.asInt() + "_$"; + return Pair.of(new ArrayMap(materializerLambda), hiddenName); + } + + /** Return ExprIds of lambda item Slots referenced by an expression. */ + private Set collectReferencedArgumentIds(Expression expression) { + Set referencedArgumentIds = Sets.newHashSet(); + expression.foreach(node -> { + if (node instanceof ArrayItemSlot) { + referencedArgumentIds.add(((ArrayItemSlot) node).getExprId()); + } + }); + return referencedArgumentIds; + } + + /** Traverse an expression and collect the Map input of every {@link MapEntryArrayMap} marker. */ + private void collectMapInputs(Expression expression, Set mapInputs) { + MapEntryArrayMap marker = unwrapMarker(expression); + if (marker != null) { + Lambda lambda = (Lambda) marker.child(0); + addMapInput(MapLambdaValidator.extractMapExpression("map lambda", lambda), mapInputs); + return; + } + + if (expression instanceof Lambda) { + for (ArrayItemReference argument : ((Lambda) expression).getLambdaArguments()) { + collectMapInputs(argument.getArrayExpression(), mapInputs); + } + return; + } + for (Expression child : expression.children()) { + collectMapInputs(child, mapInputs); + } + } + + /** Add only Maps whose key/value expansion would otherwise repeat computation. */ + private void addMapInput(Expression mapInput, Set mapInputs) { + if (MapLambdaValidator.requiresSingleEvaluation(mapInput)) { + mapInputs.add(mapInput); + } + } + + /** + * Count each complete {@link MapEntryArrayMap} expression for + * {@link #tryGenSharedEntryArrayAliases(Collection)}. As in {@code collectMapInputs}, only + * lambda argument arrays are traversed across a Lambda boundary. + */ + private void collectEntryArrayCounts(Expression expression, Map counts) { + if (unwrapMarker(expression) != null) { + counts.merge(expression, 1, Integer::sum); + return; + } + if (expression instanceof Lambda) { + for (ArrayItemReference argument : ((Lambda) expression).getLambdaArguments()) { + collectEntryArrayCounts(argument.getArrayExpression(), counts); + } + return; + } + for (Expression child : expression.children()) { + collectEntryArrayCounts(child, counts); + } + } + + /** Find the Map entry marker through analyzer-inserted Cast wrappers. */ + private MapEntryArrayMap unwrapMarker(Expression expression) { + while (expression instanceof Cast) { + expression = expression.child(0); + } + return expression instanceof MapEntryArrayMap ? (MapEntryArrayMap) expression : null; + } + + private static class JoinRewriteResult { + private final List newConjuncts; + private final Plan left; + private final Plan right; + + private JoinRewriteResult(List newConjuncts, Plan left, Plan right) { + this.newConjuncts = newConjuncts; + this.left = left; + this.right = right; + } + } +} 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..0b33402db718dd 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 @@ -24,12 +24,16 @@ import org.apache.doris.nereids.types.ArrayType; import org.apache.doris.nereids.types.DataType; import org.apache.doris.nereids.types.LambdaType; +import org.apache.doris.nereids.types.MapType; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableList.Builder; +import com.google.common.collect.ImmutableSet; import java.util.List; +import java.util.Locale; import java.util.Objects; +import java.util.Set; import java.util.stream.Collectors; /** @@ -39,6 +43,14 @@ */ public class Lambda extends Expression { + private static final Set MAP_ENTRY_LAMBDA_FUNCTIONS = ImmutableSet.of( + "map_all", + "map_apply", + "map_exists", + "map_filter", + "transform_keys", + "transform_values"); + private final List argumentNames; /** @@ -61,15 +73,37 @@ 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()) { + String normalizedFunctionName = functionName.toLowerCase(Locale.ROOT); + if (MAP_ENTRY_LAMBDA_FUNCTIONS.contains(normalizedFunctionName)) { + if (lambdaArgs.size() != 1) { + throw new AnalysisException(String.format( + "%s requires exactly one map argument but has %d", + functionName, lambdaArgs.size())); + } + if (argumentNames.size() != 2) { + throw new AnalysisException(String.format( + "lambda of %s requires exactly two arguments but has %d", + functionName, argumentNames.size())); + } + Expression mapExpression = lambdaArgs.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())); + } + builder.add(new ArrayItemReference(argumentNames.get(0), new MapKeys(mapExpression))); + builder.add(new ArrayItemReference(argumentNames.get(1), new MapValues(mapExpression))); + return builder.build(); + } + 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 +114,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..a9007048b54d17 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapAll.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.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(
+ *     (mapKey, mapValue) -> predicate,
+ *     map_keys(inputMap), map_values(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(MapLambdaValidator.requireLambda("map_all", arg)); + } + + private MapAll(Lambda lambda) { + super("map_all", new MapEntryArrayMap(lambda)); + } + + 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..224ad8e6ac50dc --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapApply.java @@ -0,0 +1,154 @@ +// 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.Cast; +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.shape.UnaryExpression; +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 com.google.common.collect.ImmutableList; + +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(
+ *   (mapKey, mapValue) -> struct(newKey, newValue),
+ *   map_keys(inputMap), map_values(inputMap)))
+ *
+ */ +public class MapApply extends ScalarFunction + implements UnaryExpression, CustomSignature, PropagateNullable, PreferPushDownProject, + RewriteWhenAnalyze { + + public MapApply(Expression arg) { + this(MapLambdaValidator.requireLambda("map_apply", arg)); + } + + private MapApply(Lambda lambda) { + super("map_apply", new MapEntryArrayMap(lambda)); + validateLambdaReturn(lambda); + } + + private MapApply(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public FunctionSignature customSignature() { + DataType mappedEntriesType = getArgument(0).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 = extractInputMapType(getEntryLambda(getArgument(0))); + StructType resolvedStructType = resolveNullFieldTypes(structType, inputMapType); + List fields = resolvedStructType.getFields(); + MapType resultType = MapType.of(fields.get(0).getDataType(), fields.get(1).getDataType()); + resultType.validateDataType(); + return FunctionSignature.ret(resultType).args(ArrayType.of(resolvedStructType)); + } + + @Override + public MapApply withChildren(List children) { + Preconditions.checkArgument(children.size() == 1); + return new MapApply(getFunctionParams(children)); + } + + @Override + public Expression rewriteWhenAnalyze() { + return new MapFromEntries(getArgument(0)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitMapApply(this, context); + } + + 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(); + StructType resolvedStructType = resolveNullFieldTypes(structType, extractInputMapType(lambda)); + MapType.of(resolvedStructType.getFields().get(0).getDataType(), + resolvedStructType.getFields().get(1).getDataType()).validateDataType(); + } + + private static MapType extractInputMapType(Lambda lambda) { + return (MapType) MapLambdaValidator.extractMapExpression("map_apply", lambda).getDataType(); + } + + private static Lambda getEntryLambda(Expression mappedEntries) { + while (mappedEntries instanceof Cast) { + mappedEntries = mappedEntries.child(0); + } + if (!(mappedEntries instanceof MapEntryArrayMap) + || !(mappedEntries.child(0) instanceof Lambda)) { + throw invalidReturnType(); + } + return (Lambda) mappedEntries.child(0); + } + + // Resolve only untyped fields in the two-field struct returned by the lambda. For + // map_apply((k, v) -> struct(cast(k as bigint), []), map(1, [10])), the result is + // MAP>. Keep the explicit BIGINT type and infer only the empty array + // from the input value type. + private static StructType resolveNullFieldTypes(StructType structType, MapType inputMapType) { + List fields = structType.getFields(); + StructField keyField = fields.get(0); + StructField valueField = fields.get(1); + keyField = keyField.withDataType(MapLambdaValidator.mergeNestedNullTypes( + keyField.getDataType(), inputMapType.getKeyType())); + valueField = valueField.withDataType(MapLambdaValidator.mergeNestedNullTypes( + valueField.getDataType(), inputMapType.getValueType())); + return new StructType(ImmutableList.of(keyField, valueField)); + } + + 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/MapEntryArrayMap.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapEntryArrayMap.java new file mode 100644 index 00000000000000..87e88521f88536 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapEntryArrayMap.java @@ -0,0 +1,38 @@ +// 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; + +/** Marks an ArrayMap whose first two arguments evaluate the key and value of each Map entry. */ +public final class MapEntryArrayMap extends ArrayMap { + + MapEntryArrayMap(Lambda lambda) { + super(lambda); + } + + @Override + public MapEntryArrayMap withChildren(List children) { + Preconditions.checkArgument(children.size() == 1 && children.get(0) instanceof Lambda); + return new MapEntryArrayMap((Lambda) children.get(0)); + } +} 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..17c99c6751b7c0 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapExists.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.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(
+ *     (mapKey, mapValue) -> predicate,
+ *     map_keys(inputMap), map_values(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(MapLambdaValidator.requireLambda("map_exists", arg)); + } + + private MapExists(Lambda lambda) { + super("map_exists", new MapEntryArrayMap(lambda)); + } + + 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..23c48027f4644f --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFilter.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.functions.PropagateNullable; +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 key and value arrays: + * + * 

+ * map_filter((mapKey, mapValue) -> predicate, inputMap)
+ *   ->
+ * map_filter(
+ *   inputMap,
+ *   array_map(
+ *     (mapKey, mapValue) -> predicate,
+ *     map_keys(inputMap), map_values(inputMap)))
+ *
+ */ +public class MapFilter extends ScalarFunction + implements HighOrderFunction, PropagateNullable { + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.retArgType(0).args( + MapType.of(new AnyDataType(0), new AnyDataType(1)), + ArrayType.of(BooleanType.INSTANCE))); + + private final boolean validateMapLambdaInput; + + // The argument is a bound Lambda. + public MapFilter(Expression arg) { + this(MapLambdaValidator.requireLambda("map_filter", arg)); + } + + public MapFilter(Expression map, Expression filter) { + super("map_filter", map, filter); + validateMapLambdaInput = false; + } + + private MapFilter(Lambda lambda) { + super("map_filter", + MapLambdaValidator.extractMapExpression("map_filter", lambda), + new MapEntryArrayMap(lambda)); + validateMapLambdaInput = true; + } + + private MapFilter(ScalarFunctionParams functionParams, boolean validateMapLambdaInput) { + super(functionParams); + this.validateMapLambdaInput = validateMapLambdaInput; + } + + public boolean shouldValidateMapLambdaInput() { + return validateMapLambdaInput; + } + + @Override + public MapFilter withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new MapFilter(getFunctionParams(children), validateMapLambdaInput); + } + + @Override + public List getImplSignature() { + return SIGNATURES; + } + + @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/MapFromArraysUnique.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromArraysUnique.java new file mode 100644 index 00000000000000..bd3207a2d5634a --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromArraysUnique.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 only when the input keys are known to be unique. */ +public class MapFromArraysUnique extends MapFromArrays { + + public MapFromArraysUnique(Expression keys, Expression values) { + super("%map_from_arrays_unique%", keys, values); + } + + private MapFromArraysUnique(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public MapFromArraysUnique withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new MapFromArraysUnique(getFunctionParams(children)); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaValidator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaValidator.java new file mode 100644 index 00000000000000..32e06bf6c1f3f1 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaValidator.java @@ -0,0 +1,168 @@ +// 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.Cast; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.literal.MapLiteral; +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.collect.ImmutableList; + +import java.util.List; + +/** + * Validates the internal ArrayMap used to evaluate map entries. + */ +public final class MapLambdaValidator { + + private MapLambdaValidator() { + } + + // Require a bound lambda argument. + public 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; + } + + // A Map lambda expands one Map into the original Map, map_keys(map), and map_values(map). + // Computed Maps are candidates for materialization. Slots are already materialized, and Map + // literals do not contain repeated computation. + public static boolean requiresSingleEvaluation(Expression mapExpression) { + while (mapExpression instanceof Cast) { + mapExpression = mapExpression.child(0); + } + return !(mapExpression instanceof Slot || mapExpression instanceof MapLiteral); + } + + /** + * Validate and return the Map shared by the first key and value references. + */ + public static Expression extractMapExpression(String functionName, Lambda lambda) { + List arguments = lambda.getLambdaArguments(); + if (arguments.size() < 2) { + throw new AnalysisException(String.format( + "Internal map entry lambda of %s must have key and value inputs", functionName)); + } + Expression keyArray = arguments.get(0).getArrayExpression(); + Expression valueArray = arguments.get(1).getArrayExpression(); + if (!(keyArray instanceof MapKeys) || !(valueArray instanceof MapValues)) { + throw new AnalysisException(String.format( + "Internal map entry lambda of %s must have key and value inputs", functionName)); + } + Expression keyMap = keyArray.child(0); + Expression valueMap = valueArray.child(0); + if (!keyMap.equals(valueMap)) { + throw new AnalysisException(String.format( + "Map entry inputs of %s must come from the same map", functionName)); + } + return keyMap; + } + + // Fill only NULL_TYPE positions in a Map lambda result from the matching input key or value + // type. For transform_values((k, v) -> [], map(1, [10])), infer ARRAY for the empty + // array. Apply the same inference to nested Array, Struct, and Map types. + static DataType mergeNestedNullTypes(DataType outputType, DataType inputType) { + if (outputType.isNullType()) { + return inputType; + } else if (outputType instanceof ArrayType && inputType instanceof ArrayType) { + return ArrayType.of(mergeNestedNullTypes( + ((ArrayType) outputType).getItemType(), ((ArrayType) inputType).getItemType())); + } else if (outputType instanceof MapType && inputType instanceof MapType) { + return MapType.of( + mergeNestedNullTypes( + ((MapType) outputType).getKeyType(), ((MapType) inputType).getKeyType()), + mergeNestedNullTypes( + ((MapType) outputType).getValueType(), ((MapType) inputType).getValueType())); + } else if (outputType instanceof StructType && inputType instanceof StructType) { + List outputFields = ((StructType) outputType).getFields(); + List inputFields = ((StructType) inputType).getFields(); + if (outputFields.size() != inputFields.size()) { + return outputType; + } + ImmutableList.Builder fields + = ImmutableList.builderWithExpectedSize(outputFields.size()); + for (int i = 0; i < outputFields.size(); i++) { + fields.add(outputFields.get(i).withDataType(mergeNestedNullTypes( + outputFields.get(i).getDataType(), inputFields.get(i).getDataType()))); + } + return new StructType(fields.build()); + } + return outputType; + } + + /** + * Revalidate the hidden physical arrays after optimizer rewrites. + */ + public static void validateStablePhysicalInputs(String functionName, Lambda lambda) { + List arguments = lambda.getLambdaArguments(); + if (arguments.isEmpty()) { + throw new AnalysisException(String.format( + "Internal map entry lambda of %s must have key and value inputs", functionName)); + } + // Projection CSE can replace only one of map_keys(M) and map_values(M) with a Slot. The + // analysis-time constructor already checked their common Map lineage, so repeating that + // structural check here would reject a valid partially materialized marker. Translation + // only needs the key/value driver arrays to be stable; ArrayMap checks equal lengths. + // Additional arguments are hidden arrays used to materialize nested lambda expressions. + // They are deliberately allowed to be volatile because each hidden array is evaluated + // once and then consumed through its item Slot by the owning lambda. + int driverCount = Math.min(2, arguments.size()); + for (int i = 0; i < driverCount; i++) { + ArrayItemReference argument = arguments.get(i); + if (argument.getArrayExpression().containsVolatileExpression()) { + throw new AnalysisException(String.format( + "Internal map entry input of %s must be materialized before translation", + functionName)); + } + } + } + + /** + * Validate functions that consume both the original Map and a mapped entry array. + */ + public static void validateOuterMapConsumer(String functionName, Expression mappedArray) { + Expression marker = mappedArray; + while (marker instanceof Cast) { + marker = marker.child(0); + } + // CSE can materialize the whole MapEntryArrayMap in an earlier projection layer when the + // enclosing Map function is referenced more than once. The marker was validated while + // translating that layer, so its projected slot is a valid physical input here. + if (marker instanceof Slot) { + return; + } + if (!(marker instanceof MapEntryArrayMap)) { + throw new AnalysisException(String.format( + "Mapped entry input of %s lost its internal map entry marker", functionName)); + } + Lambda lambda = (Lambda) marker.child(0); + validateStablePhysicalInputs(functionName, 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..c4130ee30310bb --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TransformKeys.java @@ -0,0 +1,82 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.nereids.trees.expressions.functions.scalar; + +import org.apache.doris.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 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(MapLambdaValidator.requireLambda("transform_keys", arg)); + } + + private TransformKeys(Lambda lambda) { + super("transform_keys", + MapLambdaValidator.extractMapExpression("transform_keys", lambda), + new MapEntryArrayMap(lambda)); + } + + private TransformKeys(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public FunctionSignature customSignature() { + MapType inputMapType = (MapType) getArgument(0).getDataType(); + ArrayType transformedKeysType = (ArrayType) getArgument(1).getDataType(); + // transform_keys((k, v) -> null, map(1, 10)) + // res_type should be: MAP instead of MAP + DataType resultKeyType = MapLambdaValidator.mergeNestedNullTypes( + transformedKeysType.getItemType(), inputMapType.getKeyType()); + transformedKeysType = ArrayType.of(resultKeyType); + MapType resultType = MapType.of(resultKeyType, inputMapType.getValueType()); + resultType.validateDataType(); + return FunctionSignature.ret(resultType).args(inputMapType, transformedKeysType); + } + + @Override + public TransformKeys withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new TransformKeys(getFunctionParams(children)); + } + + @Override + public Expression rewriteWhenAnalyze() { + return new MapFromArrays(getArgument(1), new MapValues(getArgument(0))); + } + + @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..70871fb55ee683 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TransformValues.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.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 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_arrays_unique%(
+ *   map_keys(inputMap),
+ *   array_map(
+ *     (mapKey, mapValue) -> newValue,
+ *     map_keys(inputMap), map_values(inputMap)))
+ *
+ */ +public class TransformValues extends ScalarFunction + implements CustomSignature, PropagateNullable, PreferPushDownProject, RewriteWhenAnalyze { + + public TransformValues(Expression arg) { + this(MapLambdaValidator.requireLambda("transform_values", arg)); + } + + private TransformValues(Lambda lambda) { + super("transform_values", + MapLambdaValidator.extractMapExpression("transform_values", lambda), + new MapEntryArrayMap(lambda)); + } + + private TransformValues(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public FunctionSignature customSignature() { + MapType inputMapType = (MapType) getArgument(0).getDataType(); + ArrayType transformedValuesType = (ArrayType) getArgument(1).getDataType(); + DataType resultValueType = MapLambdaValidator.mergeNestedNullTypes( + transformedValuesType.getItemType(), inputMapType.getValueType()); + transformedValuesType = ArrayType.of(resultValueType); + MapType resultType = MapType.of(inputMapType.getKeyType(), resultValueType); + resultType.validateDataType(); + return FunctionSignature.ret(resultType).args(inputMapType, transformedValuesType); + } + + @Override + public TransformValues withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new TransformValues(getFunctionParams(children)); + } + + @Override + public Expression rewriteWhenAnalyze() { + return new MapFromArraysUnique(new MapKeys(getArgument(0)), 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/main/java/org/apache/doris/nereids/util/PlanUtils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/PlanUtils.java index 11773fd4c0385d..9dd1c8d5e9ff42 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/PlanUtils.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/PlanUtils.java @@ -44,6 +44,8 @@ import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.expressions.WindowExpression; import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapEntryArrayMap; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapLambdaValidator; import org.apache.doris.nereids.trees.plans.Plan; import org.apache.doris.nereids.trees.plans.algebra.Filter; import org.apache.doris.nereids.trees.plans.algebra.Join; @@ -58,6 +60,7 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalProject; import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanVisitor; import org.apache.doris.nereids.types.DataType; +import org.apache.doris.nereids.types.MapType; import org.apache.doris.nereids.types.VarcharType; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.OriginStatement; @@ -194,19 +197,25 @@ public static List replaceExpressionByProjections(List childProjects, List targetExpressions) { - Set uniqueFunctionSlots = Sets.newHashSet(); + boolean containsMapLambda = targetExpressions.stream() + .anyMatch(target -> target.containsType(MapEntryArrayMap.class)); + Set nonRepeatableSlots = Sets.newHashSet(); for (Entry kv : ExpressionUtils.generateReplaceMap(childProjects).entrySet()) { - if (kv.getValue().containsVolatileExpression()) { - uniqueFunctionSlots.add(kv.getKey()); + Expression value = kv.getValue(); + if (value.containsVolatileExpression() + || value.containsType(MapEntryArrayMap.class) + || (containsMapLambda && value.getDataType() instanceof MapType + && MapLambdaValidator.requiresSingleEvaluation(value))) { + nonRepeatableSlots.add(kv.getKey()); } } - if (uniqueFunctionSlots.isEmpty()) { + if (nonRepeatableSlots.isEmpty()) { return true; } Set counterSet = Sets.newHashSet(); return targetExpressions.stream().noneMatch(target -> target.anyMatch( - e -> (e instanceof Slot) && uniqueFunctionSlots.contains(e) && !counterSet.add((Slot) e))); + e -> (e instanceof Slot) && nonRepeatableSlots.contains(e) && !counterSet.add((Slot) e))); } public static Plan skipProjectFilterLimit(Plan plan) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInputTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInputTest.java new file mode 100644 index 00000000000000..3d4ebe241f3a88 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInputTest.java @@ -0,0 +1,427 @@ +// 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.rules.rewrite; + +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.hint.DistributeHint; +import org.apache.doris.nereids.trees.expressions.Add; +import org.apache.doris.nereids.trees.expressions.Alias; +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.SlotReference; +import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Array; +import org.apache.doris.nereids.trees.expressions.functions.scalar.ArrayMap; +import org.apache.doris.nereids.trees.expressions.functions.scalar.CreateMap; +import org.apache.doris.nereids.trees.expressions.functions.scalar.CreateStruct; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Lambda; +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.MapContainsKey; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapEntryArrayMap; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapFromEntries; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapKeys; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapValues; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Random; +import org.apache.doris.nereids.trees.expressions.functions.scalar.StrToMap; +import org.apache.doris.nereids.trees.expressions.functions.scalar.TransformValues; +import org.apache.doris.nereids.trees.expressions.literal.BooleanLiteral; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; +import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; +import org.apache.doris.nereids.trees.plans.DistributeType; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.util.MemoPatternMatchSupported; +import org.apache.doris.nereids.util.MemoTestUtils; +import org.apache.doris.nereids.util.PlanChecker; +import org.apache.doris.nereids.util.PlanConstructor; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Optional; + +public class AddProjectForMapLambdaInputTest implements MemoPatternMatchSupported { + private final LogicalOlapScan studentScan + = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(), PlanConstructor.student); + + @Test + void testMaterializeNondeterministicMapAsOneExpression() { + Random random = new Random(); + CreateMap map = new CreateMap(random, new IntegerLiteral(10)); + LogicalProject input = project(transformValues(map, new IntegerLiteral(0)), studentScan); + + LogicalProject rewritten = (LogicalProject) PlanChecker.from( + MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .applyTopDown(new AddProjectForVolatileExpression()) + .applyTopDown(new MergeProjectable()) + .getPlan(); + + LogicalProject mapProject = (LogicalProject) rewritten.child(); + Alias mapAlias = lastAlias(mapProject); + Assertions.assertEquals(map, mapAlias.child()); + Assertions.assertEquals(studentScan, mapProject.child()); + assertTransformValuesUsesMapSlot(rewritten.getProjects().get(0).child(0), mapAlias.toSlot()); + } + + @Test + void testMaterializeMapApplyInput() { + SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); + CreateMap map = new CreateMap(studentId, new Random()); + MapFromEntries loweredMapApply = lowerMapApply(map); + LogicalProject input = project(loweredMapApply, studentScan); + + LogicalProject rewritten = (LogicalProject) PlanChecker.from( + MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .applyTopDown(new MergeProjectable()) + .getPlan(); + + LogicalProject mapProject = (LogicalProject) rewritten.child(); + Alias mapAlias = lastAlias(mapProject); + Assertions.assertEquals(map, mapAlias.child()); + + MapFromEntries result = (MapFromEntries) rewritten.getProjects().get(0).child(0); + Lambda entryLambda = (Lambda) ((MapEntryArrayMap) result.child(0)).child(0); + assertLambdaUsesMapSlot(entryLambda, mapAlias.toSlot()); + } + + @Test + void testMaterializeDeterministicMapInFilter() { + SlotReference studentName = (SlotReference) studentScan.getOutput().get(2); + StrToMap map = new StrToMap(studentName); + LogicalFilter input = new LogicalFilter<>(ImmutableSet.of( + new MapContainsKey(transformValues(map, new StringLiteral("value")), + new StringLiteral("key"))), studentScan); + + LogicalFilter rewritten = (LogicalFilter) PlanChecker.from( + MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .getPlan(); + + LogicalProject mapProject = (LogicalProject) rewritten.child(); + Alias mapAlias = lastAlias(mapProject); + Assertions.assertEquals(map, mapAlias.child()); + MapContainsKey predicate = (MapContainsKey) rewritten.getConjuncts().iterator().next(); + assertTransformValuesUsesMapSlot(predicate.child(0), mapAlias.toSlot()); + } + + @Test + void testMaterializeDirectMapEntryLambdaInput() { + SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); + CreateMap map = new CreateMap(studentId, new IntegerLiteral(1)); + Lambda lambda = bindLambda("map_all", map, BooleanLiteral.TRUE); + LogicalProject input = project(new MapAll(lambda), studentScan); + + LogicalProject rewritten = (LogicalProject) PlanChecker.from( + MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .getPlan(); + + LogicalProject mapProject = (LogicalProject) rewritten.child(); + Alias mapAlias = lastAlias(mapProject); + MapAll mapAll = (MapAll) rewritten.getProjects().get(0).child(0); + Lambda rewrittenLambda = (Lambda) ((MapEntryArrayMap) mapAll.child(0)).child(0); + assertLambdaUsesMapSlot(rewrittenLambda, mapAlias.toSlot()); + } + + @Test + void testMapRuleDoesNotMaterializeUnrelatedLambdaBody() { + SlotReference studentName = (SlotReference) studentScan.getOutput().get(2); + StrToMap map = new StrToMap(studentName); + Random random = new Random(); + Add lambdaBody = new Add(random, random); + LogicalProject input = project(transformValues(map, lambdaBody), studentScan); + + LogicalProject rewritten = (LogicalProject) PlanChecker.from( + MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .getPlan(); + + LogicalProject mapProject = (LogicalProject) rewritten.child(); + Assertions.assertEquals(studentScan.getOutput().size() + 1, mapProject.getProjects().size()); + Assertions.assertEquals(map, lastAlias(mapProject).child()); + TransformValues transformValues = (TransformValues) rewritten.getProjects().get(0).child(0); + Lambda rewrittenLambda = (Lambda) ((MapEntryArrayMap) transformValues.child(1)).child(0); + Assertions.assertEquals(lambdaBody, rewrittenLambda.getLambdaFunction()); + } + + @Test + void testMaterializeJoinMapInputOnLeft() { + LogicalOlapScan scoreScan + = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(), PlanConstructor.score); + SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); + CreateMap map = new CreateMap(studentId, new IntegerLiteral(1)); + LogicalJoin input = joinWithMap(studentScan, scoreScan, map); + + LogicalJoin rewritten = (LogicalJoin) PlanChecker.from( + MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .getPlan(); + + LogicalProject leftProject = (LogicalProject) rewritten.left(); + Alias mapAlias = lastAlias(leftProject); + Assertions.assertEquals(map, mapAlias.child()); + Assertions.assertEquals(scoreScan, rewritten.right()); + assertJoinTransformValuesUsesMapSlot(rewritten, mapAlias.toSlot()); + } + + @Test + void testMaterializeJoinMapInputOnRight() { + LogicalOlapScan scoreScan + = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(), PlanConstructor.score); + SlotReference scoreId = (SlotReference) scoreScan.getOutput().get(0); + Random random = new Random(); + CreateMap map = new CreateMap(new Add(scoreId, random), new IntegerLiteral(1)); + LogicalJoin input = joinWithMap(studentScan, scoreScan, map); + + LogicalJoin rewritten = (LogicalJoin) PlanChecker.from( + MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .applyTopDown(new AddProjectForVolatileExpression()) + .getPlan(); + + Assertions.assertEquals(studentScan, rewritten.left()); + LogicalProject rightProject = (LogicalProject) rewritten.right(); + Alias mapAlias = lastAlias(rightProject); + Assertions.assertEquals(map, mapAlias.child()); + assertJoinTransformValuesUsesMapSlot(rewritten, mapAlias.toSlot()); + } + + @Test + void testSkipJoinMapInputDependingOnBothSides() { + LogicalOlapScan scoreScan + = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(), PlanConstructor.score); + SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); + SlotReference scoreId = (SlotReference) scoreScan.getOutput().get(0); + CreateMap map = new CreateMap(studentId, scoreId); + LogicalJoin input = joinWithMap(studentScan, scoreScan, map); + + Plan rewritten = PlanChecker.from(MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .getPlan(); + + Assertions.assertEquals(input, rewritten); + } + + @Test + void testRejectVolatileJoinMapInputDependingOnBothSides() { + LogicalOlapScan scoreScan + = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(), PlanConstructor.score); + SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); + SlotReference scoreId = (SlotReference) scoreScan.getOutput().get(0); + CreateMap map = new CreateMap(new Add(studentId, new Random()), scoreId); + LogicalJoin input = joinWithMap(studentScan, scoreScan, map); + + Assertions.assertThrows(AnalysisException.class, () -> PlanChecker.from( + MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .getPlan()); + } + + @Test + void testMaterializeMapInputInNestedLambda() { + SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); + CreateMap outerMap = new CreateMap(studentId, new IntegerLiteral(10)); + Lambda outerTemplate = new Lambda(ImmutableList.of("ok", "ov"), new IntegerLiteral(0)); + List outerArguments = outerTemplate.makeArguments( + "transform_values", ImmutableList.of(outerMap)); + Slot outerKey = outerArguments.get(0).toSlot(); + Slot outerValue = outerArguments.get(1).toSlot(); + + CreateMap innerMap = new CreateMap(new Add(outerKey, new Random()), outerValue); + Lambda innerTemplate = new Lambda(ImmutableList.of("ik", "iv"), new IntegerLiteral(0)); + List innerArguments = innerTemplate.makeArguments( + "transform_values", ImmutableList.of(innerMap)); + TransformValues innerTransform = new TransformValues( + innerTemplate.withLambdaFunctionArguments(innerArguments.get(0).toSlot(), innerArguments)); + Lambda outerLambda = outerTemplate.withLambdaFunctionArguments(innerTransform, outerArguments); + LogicalProject input = project(new TransformValues(outerLambda), studentScan); + + LogicalProject rewritten = (LogicalProject) PlanChecker.from( + MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .getPlan(); + + TransformValues outerTransform = (TransformValues) rewritten.getProjects().get(0).child(0); + Lambda rewrittenOuterLambda = (Lambda) ((MapEntryArrayMap) outerTransform.child(1)).child(0); + Assertions.assertEquals(3, rewrittenOuterLambda.getLambdaArguments().size()); + ArrayItemReference hiddenArgument = rewrittenOuterLambda.getLambdaArgument(2); + Assertions.assertInstanceOf(ArrayMap.class, hiddenArgument.getArrayExpression()); + ArrayMap materializer = (ArrayMap) hiddenArgument.getArrayExpression(); + Lambda materializerLambda = (Lambda) materializer.child(0); + Assertions.assertTrue(materializerLambda.getLambdaFunction().containsType(Random.class)); + + TransformValues rewrittenInnerTransform + = (TransformValues) rewrittenOuterLambda.getLambdaFunction(); + Assertions.assertEquals(hiddenArgument.toSlot(), rewrittenInnerTransform.child(0)); + Lambda rewrittenInnerLambda + = (Lambda) ((MapEntryArrayMap) rewrittenInnerTransform.child(1)).child(0); + assertLambdaUsesMapSlot(rewrittenInnerLambda, hiddenArgument.toSlot()); + + LogicalProject rewrittenAgain = (LogicalProject) PlanChecker.from( + MemoTestUtils.createConnectContext(), rewritten) + .applyTopDown(new AddProjectForMapLambdaInput()) + .getPlan(); + TransformValues outerTransformAgain + = (TransformValues) rewrittenAgain.getProjects().get(0).child(0); + Lambda outerLambdaAgain + = (Lambda) ((MapEntryArrayMap) outerTransformAgain.child(1)).child(0); + Assertions.assertEquals(3, outerLambdaAgain.getLambdaArguments().size()); + } + + @Test + void testMaterializeMapInputInRegularArrayMapLambda() { + Array inputArray = new Array(new IntegerLiteral(1), new IntegerLiteral(2)); + Lambda outerTemplate = new Lambda(ImmutableList.of("x"), new IntegerLiteral(0)); + List outerArguments = outerTemplate.makeArguments( + "array_map", ImmutableList.of(inputArray)); + Slot outerItem = outerArguments.get(0).toSlot(); + + CreateMap innerMap = new CreateMap(new Random(), outerItem); + Lambda innerTemplate = new Lambda(ImmutableList.of("k", "v"), new IntegerLiteral(0)); + List innerArguments = innerTemplate.makeArguments( + "transform_values", ImmutableList.of(innerMap)); + TransformValues innerTransform = new TransformValues( + innerTemplate.withLambdaFunctionArguments(innerArguments.get(0).toSlot(), innerArguments)); + Lambda outerLambda = outerTemplate.withLambdaFunctionArguments(innerTransform, outerArguments); + LogicalProject input = project(new ArrayMap(outerLambda), studentScan); + + LogicalProject rewritten = (LogicalProject) PlanChecker.from( + MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .applyTopDown(new AddProjectForVolatileExpression()) + .getPlan(); + + Assertions.assertEquals(studentScan, rewritten.child()); + ArrayMap rewrittenArrayMap = (ArrayMap) rewritten.getProjects().get(0).child(0); + Lambda rewrittenOuterLambda = (Lambda) rewrittenArrayMap.child(0); + Assertions.assertEquals(2, rewrittenOuterLambda.getLambdaArguments().size()); + ArrayItemReference hiddenArgument = rewrittenOuterLambda.getLambdaArgument(1); + Assertions.assertInstanceOf(ArrayMap.class, hiddenArgument.getArrayExpression()); + ArrayMap materializer = (ArrayMap) hiddenArgument.getArrayExpression(); + Lambda materializerLambda = (Lambda) materializer.child(0); + Assertions.assertTrue(materializerLambda.getLambdaFunction().containsType(Random.class)); + + TransformValues rewrittenInnerTransform + = (TransformValues) rewrittenOuterLambda.getLambdaFunction(); + Assertions.assertEquals(hiddenArgument.toSlot(), rewrittenInnerTransform.child(0)); + Lambda rewrittenInnerLambda + = (Lambda) ((MapEntryArrayMap) rewrittenInnerTransform.child(1)).child(0); + assertLambdaUsesMapSlot(rewrittenInnerLambda, hiddenArgument.toSlot()); + } + + @Test + void testMaterializeNestedMapApplyInput() { + Array inputArray = new Array(new IntegerLiteral(1), new IntegerLiteral(2)); + Lambda outerTemplate = new Lambda(ImmutableList.of("x"), new IntegerLiteral(0)); + List outerArguments = outerTemplate.makeArguments( + "array_map", ImmutableList.of(inputArray)); + CreateMap innerMap = new CreateMap(new Random(), outerArguments.get(0).toSlot()); + MapFromEntries loweredMapApply = lowerMapApply(innerMap); + Lambda outerLambda = outerTemplate.withLambdaFunctionArguments( + loweredMapApply, outerArguments); + LogicalProject input = project(new ArrayMap(outerLambda), studentScan); + + LogicalProject rewritten = (LogicalProject) PlanChecker.from( + MemoTestUtils.createConnectContext(), input) + .applyTopDown(new AddProjectForMapLambdaInput()) + .getPlan(); + + ArrayMap rewrittenArrayMap = (ArrayMap) rewritten.getProjects().get(0).child(0); + Lambda rewrittenOuterLambda = (Lambda) rewrittenArrayMap.child(0); + // Keep x plus the materialized Map input used by the nested Map Lambda. + Assertions.assertEquals(2, rewrittenOuterLambda.getLambdaArguments().size()); + ArrayItemReference mapArgument = rewrittenOuterLambda.getLambdaArgument(1); + Assertions.assertInstanceOf(ArrayMap.class, mapArgument.getArrayExpression()); + Assertions.assertTrue(mapArgument.getArrayExpression().containsType(Random.class)); + + MapFromEntries result = (MapFromEntries) rewrittenOuterLambda.getLambdaFunction(); + Lambda entryLambda = (Lambda) ((MapEntryArrayMap) result.child(0)).child(0); + assertLambdaUsesMapSlot(entryLambda, mapArgument.toSlot()); + } + + private LogicalProject project(Expression expression, Plan child) { + return new LogicalProject(ImmutableList.of(new Alias(expression)), child); + } + + private TransformValues transformValues(Expression map, Expression body) { + return new TransformValues(bindLambda("transform_values", map, body)); + } + + private MapFromEntries lowerMapApply(Expression map) { + Lambda lambda = new Lambda(ImmutableList.of("k", "v"), new IntegerLiteral(0)); + List arguments = lambda.makeArguments("map_apply", ImmutableList.of(map)); + CreateStruct body = new CreateStruct(arguments.get(0).toSlot(), arguments.get(1).toSlot()); + MapApply mapApply = new MapApply(lambda.withLambdaFunctionArguments(body, arguments)); + return (MapFromEntries) mapApply.rewriteWhenAnalyze(); + } + + private Lambda bindLambda(String functionName, Expression map, Expression body) { + Lambda lambda = new Lambda(ImmutableList.of("k", "v"), body); + List arguments = lambda.makeArguments(functionName, ImmutableList.of(map)); + return lambda.withLambdaFunctionArguments(body, arguments); + } + + private LogicalJoin joinWithMap(Plan left, Plan right, Expression map) { + MapContainsKey predicate = new MapContainsKey( + transformValues(map, new IntegerLiteral(0)), new IntegerLiteral(1)); + return new LogicalJoin( + JoinType.INNER_JOIN, + ImmutableList.of(), + ImmutableList.of(predicate), + new DistributeHint(DistributeType.NONE), + Optional.empty(), + left, + right, + null); + } + + private Alias lastAlias(LogicalProject project) { + return (Alias) project.getProjects().get(project.getProjects().size() - 1); + } + + private void assertJoinTransformValuesUsesMapSlot(LogicalJoin join, Slot mapSlot) { + MapContainsKey predicate = (MapContainsKey) join.getOtherJoinConjuncts().get(0); + assertTransformValuesUsesMapSlot(predicate.child(0), mapSlot); + } + + private void assertTransformValuesUsesMapSlot(Expression expression, Slot mapSlot) { + TransformValues transformValues = (TransformValues) expression; + Assertions.assertEquals(mapSlot, transformValues.child(0)); + Lambda lambda = (Lambda) ((MapEntryArrayMap) transformValues.child(1)).child(0); + assertLambdaUsesMapSlot(lambda, mapSlot); + } + + private void assertLambdaUsesMapSlot(Lambda lambda, Slot mapSlot) { + Assertions.assertEquals(mapSlot, + ((MapKeys) lambda.getLambdaArgument(0).getArrayExpression()).child(0)); + Assertions.assertEquals(mapSlot, + ((MapValues) lambda.getLambdaArgument(1).getArrayExpression()).child(0)); + } + +} 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..69c7d1ddc0b6a7 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaFunctionsTest.java @@ -0,0 +1,344 @@ +// 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.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.StructType; +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 MapFilter); + assertMapEntryArray(((MapFilter) mapFilter).child(1)); + + 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 MapFromArrays); + assertMapEntryArray(transformKeys.child(0)); + Assertions.assertTrue(transformKeys.child(1) instanceof MapValues); + 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 MapFromArraysUnique); + Assertions.assertTrue(transformValues.child(0) instanceof MapKeys); + assertMapEntryArray(transformValues.child(1)); + 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 testPureNullLambdaReturnUsesInputMapType() { + MapType inputMapType = MapType.of(TinyIntType.INSTANCE, TinyIntType.INSTANCE); + + Expression transformValues = analyze( + "transform_values((k, v) -> null, map(1, 10))"); + Assertions.assertEquals(inputMapType, transformValues.getDataType()); + Assertions.assertEquals(ArrayType.of(TinyIntType.INSTANCE), + transformValues.child(1).getDataType()); + + Expression transformKeys = analyze( + "transform_keys((k, v) -> null, map(1, 10))"); + Assertions.assertEquals(inputMapType, transformKeys.getDataType()); + Assertions.assertEquals(ArrayType.of(TinyIntType.INSTANCE), + transformKeys.child(1).getDataType()); + + Expression mapApply = analyze( + "map_apply((k, v) -> struct(k, null), map(1, 10))"); + Assertions.assertEquals(inputMapType, mapApply.getDataType()); + StructType mappedEntryType = + (StructType) ((ArrayType) mapApply.child(0).getDataType()).getItemType(); + Assertions.assertEquals(TinyIntType.INSTANCE, + mappedEntryType.getFields().get(0).getDataType()); + Assertions.assertEquals(TinyIntType.INSTANCE, + mappedEntryType.getFields().get(1).getDataType()); + + Expression mapApplyNullKey = analyze( + "map_apply((k, v) -> struct(null, v), map(1, 10))"); + Assertions.assertEquals(inputMapType, mapApplyNullKey.getDataType()); + } + + @Test + public void testNestedNullLambdaReturnUsesInputMapType() { + ArrayType tinyIntArrayType = ArrayType.of(TinyIntType.INSTANCE); + MapType arrayValueMapType = MapType.of(TinyIntType.INSTANCE, tinyIntArrayType); + + Expression transformArrayValues = analyze( + "transform_values((k, v) -> [], map(1, [10]))"); + Assertions.assertEquals(arrayValueMapType, transformArrayValues.getDataType()); + Assertions.assertEquals(ArrayType.of(tinyIntArrayType), + transformArrayValues.child(1).getDataType()); + + Expression mapApplyArrayValue = analyze( + "map_apply((k, v) -> struct(k, []), map(1, [10]))"); + Assertions.assertEquals(arrayValueMapType, mapApplyArrayValue.getDataType()); + + MapType nestedMapType = MapType.of( + TinyIntType.INSTANCE, MapType.of(TinyIntType.INSTANCE, TinyIntType.INSTANCE)); + Expression transformMapValues = analyze( + "transform_values((k, v) -> map(), map(1, map(2, 20)))"); + Assertions.assertEquals(nestedMapType, transformMapValues.getDataType()); + + Expression transformStructValues = analyze( + "transform_values((k, v) -> struct(null, []), " + + "map(1, struct(10, [20])))"); + StructType transformStructType = (StructType) + ((MapType) transformStructValues.getDataType()).getValueType(); + Assertions.assertEquals(TinyIntType.INSTANCE, + transformStructType.getFields().get(0).getDataType()); + Assertions.assertEquals(tinyIntArrayType, + transformStructType.getFields().get(1).getDataType()); + + Expression mapApplyStructValue = analyze( + "map_apply((k, v) -> struct(k, struct(null, [])), " + + "map(1, struct(10, [20])))"); + StructType mapApplyStructType = (StructType) + ((MapType) mapApplyStructValue.getDataType()).getValueType(); + Assertions.assertEquals(TinyIntType.INSTANCE, + mapApplyStructType.getFields().get(0).getDataType()); + Assertions.assertEquals(tinyIntArrayType, + mapApplyStructType.getFields().get(1).getDataType()); + } + + @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 testComputedMapIsAccepted() { + SlotReference value = new SlotReference("value", IntegerType.INSTANCE); + CreateMap computedMap = new CreateMap(Literal.of(1), value); + Lambda lambda = new Lambda(ImmutableList.of("k", "v"), value); + List arguments = + lambda.makeArguments("transform_values", ImmutableList.of(computedMap)); + Lambda boundLambda = lambda.withLambdaFunctionArguments(arguments.get(1).toSlot(), arguments); + 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 MapFromArraysUnique); + } + + @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); + + AnalysisException invalidNestedValueException = Assertions.assertThrows(AnalysisException.class, + () -> analyze("map_from_arrays([1], [[]])")); + Assertions.assertTrue(invalidNestedValueException.getMessage().contains( + "Unsupported data type: map>"), + invalidNestedValueException::getMessage); + } + + @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 MapEntryArrayMap); + MapEntryArrayMap marker = (MapEntryArrayMap) expression; + Assertions.assertTrue(marker.child(0) instanceof Lambda); + + Lambda lambda = (Lambda) marker.child(0); + List arguments = lambda.getLambdaArguments(); + Assertions.assertEquals(2, arguments.size()); + Assertions.assertTrue(arguments.get(0).getArrayExpression() instanceof MapKeys); + Assertions.assertTrue(arguments.get(1).getArrayExpression() instanceof MapValues); + Assertions.assertEquals( + arguments.get(0).getArrayExpression().child(0), + arguments.get(1).getArrayExpression().child(0)); + Assertions.assertNotEquals(arguments.get(0).getExprId(), arguments.get(1).getExprId()); + + Assertions.assertTrue(marker.withChildren(marker.children()) instanceof MapEntryArrayMap); + } + +} 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..91e86065a1e004 --- /dev/null +++ b/regression-test/data/query_p0/sql_functions/map_functions/test_map_lambda.out @@ -0,0 +1,141 @@ +-- 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_null_type -- +2 \N \N + +-- !transform_keys_null_type -- +1 20 + +-- !map_apply_null_value_type -- +2 \N \N + +-- !transform_values_nested_array_null_type -- +[] + +-- !map_apply_nested_array_null_type -- +[] + +-- !transform_values_nested_map_null_type -- +0 + +-- !transform_values_nested_struct_null_type -- +{"col1":null, "col2":[]} + +-- !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 + 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..18e41f4c754cc7 --- /dev/null +++ b/regression-test/suites/query_p0/sql_functions/map_functions/test_map_lambda.groovy @@ -0,0 +1,449 @@ +// 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_null_type """ + select map_size(r), r[1], r[2] + from ( + select transform_values((k, v) -> null, mii) r + from test_map_lambda where id = 1 + ) t + """ + qt_transform_keys_null_type """ + select map_size(r), map_values(r)[1] + from ( + select transform_keys((k, v) -> null, mii) r + from test_map_lambda where id = 1 + ) t + """ + qt_map_apply_null_value_type """ + select map_size(r), r[1], r[2] + from ( + select map_apply((k, v) -> struct(k, null), mii) r + from test_map_lambda where id = 1 + ) t + """ + qt_transform_values_nested_array_null_type """ + select r[1] + from ( + select transform_values((k, v) -> [], map(1, [10])) r + ) t + """ + qt_map_apply_nested_array_null_type """ + select r[1] + from ( + select map_apply((k, v) -> struct(k, []), map(1, [10])) r + ) t + """ + qt_transform_values_nested_map_null_type """ + select map_size(r[1]) + from ( + select transform_values((k, v) -> map(), map(1, map(2, 20))) r + ) t + """ + qt_transform_values_nested_struct_null_type """ + select r[1] + from ( + select transform_values( + (k, v) -> struct(null, []), map(1, struct(10, [20]))) r + ) 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" + } + test { + sql "select map_from_arrays([1], [[]])" + exception "Unsupported data type: map>" + } + 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" + } +} From a7af50338e2b0443478fa3f8a8531044f7f1a0c7 Mon Sep 17 00:00:00 2001 From: linzhenqi Date: Mon, 24 Aug 2026 01:47:28 +0800 Subject: [PATCH 2/4] refactor --- be/src/exprs/function/function_map.cpp | 223 +++-- be/test/exprs/function/function_map_test.cpp | 223 +++++ .../glue/translator/ExpressionTranslator.java | 10 - .../doris/nereids/jobs/executor/Rewriter.java | 7 +- .../apache/doris/nereids/rules/RuleType.java | 1 - .../rules/analysis/ExpressionAnalyzer.java | 105 ++- .../AccessPathExpressionCollector.java | 21 + .../rewrite/AddProjectForMapLambdaInput.java | 786 ------------------ .../expressions/functions/scalar/Lambda.java | 35 - .../expressions/functions/scalar/MapAll.java | 12 +- .../functions/scalar/MapApply.java | 57 +- .../functions/scalar/MapExists.java | 12 +- .../functions/scalar/MapFilter.java | 41 +- .../functions/scalar/MapFromEntries.java | 6 +- ...sUnique.java => MapFromEntriesUnique.java} | 16 +- ...java => MapFromFilteredEntriesUnique.java} | 18 +- .../scalar/MapLambdaFunctionUtils.java | 129 +++ .../functions/scalar/MapLambdaValidator.java | 168 ---- .../functions/scalar/TransformKeys.java | 28 +- .../functions/scalar/TransformValues.java | 35 +- .../apache/doris/nereids/util/PlanUtils.java | 19 +- .../AddProjectForMapLambdaInputTest.java | 427 ---------- .../rules/rewrite/PruneNestedColumnTest.java | 23 + .../scalar/MapLambdaFunctionsTest.java | 88 +- .../map_functions/test_map_lambda.out | 3 + .../map_functions/test_map_lambda.groovy | 5 +- 26 files changed, 853 insertions(+), 1645 deletions(-) delete mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInput.java rename fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/{MapFromArraysUnique.java => MapFromEntriesUnique.java} (66%) rename fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/{MapEntryArrayMap.java => MapFromFilteredEntriesUnique.java} (62%) create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaFunctionUtils.java delete mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaValidator.java delete mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInputTest.java diff --git a/be/src/exprs/function/function_map.cpp b/be/src/exprs/function/function_map.cpp index 9615a5523f8214..db4abbb274f3e3 100644 --- a/be/src/exprs/function/function_map.cpp +++ b/be/src/exprs/function/function_map.cpp @@ -202,57 +202,70 @@ class FunctionMapFilter : public IFunction { 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); - ColumnPtr map_column = unpacked_map_column; - ColumnPtr predicate_column = unpacked_predicate_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 = [&](ColumnPtr& column, bool is_const) { + 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); - column = nullable->get_nested_column_ptr(); + return nullable->get_nested_column(); } + return *column; }; - merge_null_map(map_column, map_is_const); - merge_null_map(predicate_column, predicate_is_const); - const auto& map = assert_cast(*map_column); - const auto& predicate = assert_cast(*predicate_column); - if ((!map_is_const && map.size() != input_rows_count) || - (!predicate_is_const && predicate.size() != input_rows_count)) { - return Status::InvalidArgument( - "The map and lambda result offsets of function {} must be identical", name); + 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(); + 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()); } - for (size_t row = 0; row < input_rows_count; ++row) { - if (!result_null_map_data[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); - } + 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(); + } - const IColumn* predicate_data = &predicate.get_data(); - const UInt8* predicate_null_map = nullptr; - if (const auto* nullable_predicate = check_and_get_column(predicate_data)) { - predicate_null_map = nullable_predicate->get_null_map_data().data(); - predicate_data = &nullable_predicate->get_nested_column(); - } - const auto& predicate_values = assert_cast(*predicate_data).get_data(); +private: + static void 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) { + 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(); - IColumn::Selector selector; if (!map_is_const) { selector.reserve(map.get_keys().size()); } - auto result_keys = map.get_keys().clone_empty(); - auto result_values = map.get_values().clone_empty(); - auto result_offsets = ColumnArray::ColumnOffsets::create(); - result_offsets->reserve(input_rows_count); + result_offsets.reserve(result_null_map.size()); - size_t output_offset = 0; - for (size_t row = 0; row < input_rows_count; ++row) { - if (!result_null_map_data[row]) { + 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 = @@ -260,34 +273,27 @@ class FunctionMapFilter : public IFunction { 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 == nullptr || - predicate_null_map[predicate_entry] == 0) && + const bool selected = predicate_null_map[predicate_entry] == 0 && predicate_values[predicate_entry] != 0; if (selected) { - const size_t map_entry = map_begin + entry; - selector.push_back(map_entry); - ++output_offset; + selector.push_back(map_begin + entry); } } } - result_offsets->insert_value(output_offset); - } - if (map_is_const && !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()); - } else if (!map_is_const) { - map.get_keys().append_data_by_selector(result_keys, selector); - map.get_values().append_data_by_selector(result_values, selector); + result_offsets.insert_value(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)); + } + + 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(); } @@ -569,9 +575,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; } @@ -608,7 +615,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), @@ -643,6 +652,100 @@ 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(); + 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 void build_selector_and_offsets(const ColumnArray& entries, + const ColumnNullable& nullable_entries, + const ColumnNullable* nullable_array, + IColumn::Selector& selector, + ColumnArray::ColumnOffsets& filtered_offsets) { + 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(static_cast(entry)); + } + } + } + filtered_offsets.insert_value(selector.size()); + } + } +}; + class FunctionStrToMap : public IFunction { public: static constexpr auto name = "str_to_map"; @@ -1104,7 +1207,9 @@ 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(); diff --git a/be/test/exprs/function/function_map_test.cpp b/be/test/exprs/function/function_map_test.cpp index 22cdd72256fc1e..f51496c1277b00 100644 --- a/be/test/exprs/function/function_map_test.cpp +++ b/be/test/exprs/function/function_map_test.cpp @@ -51,6 +51,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 +69,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 +414,205 @@ 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_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/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java index 5dd5fc7d4b5fa5..84aea35f2c5001 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java @@ -100,9 +100,6 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.DictGetMany; 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.MapEntryArrayMap; -import org.apache.doris.nereids.trees.expressions.functions.scalar.MapFilter; -import org.apache.doris.nereids.trees.expressions.functions.scalar.MapLambdaValidator; import org.apache.doris.nereids.trees.expressions.functions.scalar.ScalarFunction; import org.apache.doris.nereids.trees.expressions.functions.udf.JavaUdaf; import org.apache.doris.nereids.trees.expressions.functions.udf.JavaUdf; @@ -535,9 +532,6 @@ public Expr visitLambda(Lambda lambda, PlanTranslatorContext context) { @Override public Expr visitArrayMap(ArrayMap arrayMap, PlanTranslatorContext context) { Lambda lambda = (Lambda) arrayMap.child(0); - if (arrayMap instanceof MapEntryArrayMap) { - MapLambdaValidator.validateStablePhysicalInputs(arrayMap.getName(), lambda); - } List arguments = new ArrayList<>(arrayMap.children().size()); arguments.add(null); int columnId = 0; @@ -727,10 +721,6 @@ public Expr visitSearchExpression(SearchExpression searchExpression, @Override public Expr visitScalarFunction(ScalarFunction function, PlanTranslatorContext context) { - if (function instanceof MapFilter && ((MapFilter) function).shouldValidateMapLambdaInput()) { - MapLambdaValidator.validateOuterMapConsumer( - function.getName(), function.getArgument(1)); - } List arguments = function.getArguments().stream() .map(arg -> arg.accept(this, context)) .collect(Collectors.toList()); 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 0dc1ca7755e2a9..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 @@ -34,7 +34,6 @@ import org.apache.doris.nereids.rules.expression.QueryColumnCollector; import org.apache.doris.nereids.rules.rewrite.AddDefaultLimit; import org.apache.doris.nereids.rules.rewrite.AddProjectForJoin; -import org.apache.doris.nereids.rules.rewrite.AddProjectForMapLambdaInput; import org.apache.doris.nereids.rules.rewrite.AddProjectForVolatileExpression; import org.apache.doris.nereids.rules.rewrite.AdjustConjunctsReturnType; import org.apache.doris.nereids.rules.rewrite.AdjustNullable; @@ -763,11 +762,7 @@ public class Rewriter extends AbstractBatchJobExecutor { topDown(new SumLiteralRewrite(), new MergePercentileToArray()) ), - topic("add projection for volatile and lambda expression", - // Materialize Map Lambda inputs and volatile expressions before merging the - // generated Projects. Separate passes avoid repeatedly adding and merging the - // same Project in one top-down rewrite. - topDown(new AddProjectForMapLambdaInput()), + topic("add projection for volatile expression", topDown(new AddProjectForVolatileExpression()), topDown(new MergeProjectable()) ), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java index 13da88ef23a299..a981dc496813b6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java @@ -228,7 +228,6 @@ public enum RuleType { PUSH_DOWN_MAX_MIN_FILTER(RuleTypeClass.REWRITE), ADD_PROJECT_FOR_JOIN(RuleTypeClass.REWRITE), - ADD_PROJECT_FOR_MAP_LAMBDA_INPUT(RuleTypeClass.REWRITE), ADD_PROJECT_FOR_VOLATILE_EXPRESSION(RuleTypeClass.REWRITE), VARIANT_SUB_PATH_PRUNING(RuleTypeClass.REWRITE), NESTED_COLUMN_PRUNING(RuleTypeClass.REWRITE), 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/rules/rewrite/AddProjectForMapLambdaInput.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInput.java deleted file mode 100644 index 6bc611215aab97..00000000000000 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInput.java +++ /dev/null @@ -1,786 +0,0 @@ -// 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.rules.rewrite; - -import org.apache.doris.common.Pair; -import org.apache.doris.nereids.exceptions.AnalysisException; -import org.apache.doris.nereids.rules.Rule; -import org.apache.doris.nereids.rules.RuleType; -import org.apache.doris.nereids.trees.expressions.Alias; -import org.apache.doris.nereids.trees.expressions.ArrayItemReference; -import org.apache.doris.nereids.trees.expressions.ArrayItemReference.ArrayItemSlot; -import org.apache.doris.nereids.trees.expressions.Cast; -import org.apache.doris.nereids.trees.expressions.ExprId; -import org.apache.doris.nereids.trees.expressions.Expression; -import org.apache.doris.nereids.trees.expressions.NamedExpression; -import org.apache.doris.nereids.trees.expressions.Slot; -import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; -import org.apache.doris.nereids.trees.expressions.functions.Function; -import org.apache.doris.nereids.trees.expressions.functions.scalar.ArrayMap; -import org.apache.doris.nereids.trees.expressions.functions.scalar.Lambda; -import org.apache.doris.nereids.trees.expressions.functions.scalar.MapEntryArrayMap; -import org.apache.doris.nereids.trees.expressions.functions.scalar.MapLambdaValidator; -import org.apache.doris.nereids.trees.plans.JoinType; -import org.apache.doris.nereids.trees.plans.Plan; -import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; -import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; -import org.apache.doris.nereids.trees.plans.logical.LogicalGenerate; -import org.apache.doris.nereids.trees.plans.logical.LogicalHaving; -import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; -import org.apache.doris.nereids.trees.plans.logical.LogicalOneRowRelation; -import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; -import org.apache.doris.nereids.trees.plans.logical.LogicalProject; -import org.apache.doris.nereids.util.ExpressionUtils; -import org.apache.doris.nereids.util.JoinUtils; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableSet; -import com.google.common.collect.Lists; -import com.google.common.collect.Maps; -import com.google.common.collect.Sets; - -import java.util.Collection; -import java.util.List; -import java.util.Map; -import java.util.Map.Entry; -import java.util.Optional; -import java.util.Set; - -/** - * Materialize computed Map inputs used by {@link MapEntryArrayMap}. - * - *

A Map entry lambda takes {@code map_keys(computedMap)} and - * {@code map_values(computedMap)} as its two input arrays. rule evaThisluates - * {@code computedMap} in a child Project and replaces all its occurrences with the same Slot: - * - * 

- * before:
- *   Project[map_from_arrays(
- *     map_keys(computedMap),
- *     MapEntryArrayMap(
- *       (mapKey, mapValue) -> valueExpression,
- *       map_keys(computedMap), map_values(computedMap)))]
- *     child
- *
- * after:
- *   Project[map_from_arrays(
- *     map_keys(materializedMapSlot),
- *     MapEntryArrayMap(
- *       (mapKey, mapValue) -> valueExpression,
- *       map_keys(materializedMapSlot), map_values(materializedMapSlot)))]
- *     Project[child.*, computedMap AS materializedMapSlot]
- *       child
- *
- * - *

Besides the basic rewrite above, this rule handles - * repeated entry arrays, nested lambdas, and Join children through dedicated helper methods below. - */ -public class AddProjectForMapLambdaInput implements RewriteRuleFactory { - - @Override - public List buildRules() { - return ImmutableList.of( - new GenerateRewrite().build(), - new OneRowRelationRewrite().build(), - new ProjectRewrite().build(), - new FilterRewrite().build(), - new HavingRewrite().build(), - new AggregateRewrite().build(), - new JoinRewrite().build() - ); - } - - private class GenerateRewrite extends OneRewriteRuleFactory { - @Override - public Rule build() { - return logicalGenerate().thenApply(ctx -> { - LogicalGenerate generate = ctx.root; - List generators = materializeNestedMapInputs(generate.getGenerators()); - Optional, LogicalProject>> - rewrittenOpt = rewriteExpressions(generate, generators); - if (rewrittenOpt.isPresent()) { - return generate.withGenerators(rewrittenOpt.get().first) - .withChildren(rewrittenOpt.get().second); - } else if (!generators.equals(generate.getGenerators())) { - return generate.withGenerators(generators); - } else { - return generate; - } - }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); - } - } - - private class OneRowRelationRewrite extends OneRewriteRuleFactory { - @Override - public Rule build() { - return logicalOneRowRelation().thenApply(ctx -> { - LogicalOneRowRelation oneRowRelation = ctx.root; - List projects = materializeNestedMapInputs(oneRowRelation.getProjects()); - List mapInputAliases = tryGenMapInputAliases(projects); - List rewrittenProjects = replaceExpressions(projects, mapInputAliases); - List entryArrayAliases = tryGenSharedEntryArrayAliases(rewrittenProjects); - if (mapInputAliases.isEmpty() && entryArrayAliases.isEmpty()) { - return projects.equals(oneRowRelation.getProjects()) - ? oneRowRelation : oneRowRelation.withProjects(projects); - } - - // A OneRowRelation has no child on which to install the usual materialization - // Project. Use the relation itself as the lowest projection, then stack the shared - // entry-array Project and the original output Project above it. - Plan child; - if (mapInputAliases.isEmpty()) { - child = oneRowRelation.withProjects(entryArrayAliases); - } else { - child = oneRowRelation.withProjects(mapInputAliases); - if (!entryArrayAliases.isEmpty()) { - child = appendProject(child, entryArrayAliases); - } - } - rewrittenProjects = replaceExpressions(rewrittenProjects, entryArrayAliases); - return new LogicalProject<>(rewrittenProjects, child); - }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); - } - } - - private class ProjectRewrite extends OneRewriteRuleFactory { - @Override - public Rule build() { - return logicalProject().thenApply(ctx -> { - LogicalProject project = ctx.root; - List projects = materializeNestedMapInputs(project.getProjects()); - Optional, LogicalProject>> - rewrittenOpt = rewriteExpressions(project, projects); - if (rewrittenOpt.isPresent()) { - return project.withProjectsAndChild(rewrittenOpt.get().first, rewrittenOpt.get().second); - } else if (!projects.equals(project.getProjects())) { - return project.withProjects(projects); - } else { - return project; - } - }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); - } - } - - private class FilterRewrite extends OneRewriteRuleFactory { - @Override - public Rule build() { - return logicalFilter().thenApply(ctx -> { - LogicalFilter filter = ctx.root; - List conjuncts = materializeNestedMapInputs(filter.getConjuncts()); - Optional, LogicalProject>> - rewrittenOpt = rewriteExpressions(filter, conjuncts); - if (rewrittenOpt.isPresent()) { - return filter.withConjunctsAndChild( - ImmutableSet.copyOf(rewrittenOpt.get().first), - rewrittenOpt.get().second); - } else if (!ImmutableSet.copyOf(conjuncts).equals(filter.getConjuncts())) { - return filter.withConjuncts(ImmutableSet.copyOf(conjuncts)); - } else { - return filter; - } - }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); - } - } - - private class HavingRewrite extends OneRewriteRuleFactory { - @Override - public Rule build() { - return logicalHaving().thenApply(ctx -> { - LogicalHaving having = ctx.root; - List conjuncts = materializeNestedMapInputs(having.getConjuncts()); - Optional, LogicalProject>> - rewrittenOpt = rewriteExpressions(having, conjuncts); - if (rewrittenOpt.isPresent()) { - return having.withConjuncts(ImmutableSet.copyOf(rewrittenOpt.get().first)) - .withChildren(rewrittenOpt.get().second); - } else if (!ImmutableSet.copyOf(conjuncts).equals(having.getConjuncts())) { - return having.withConjuncts(ImmutableSet.copyOf(conjuncts)); - } else { - return having; - } - }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); - } - } - - private class AggregateRewrite extends OneRewriteRuleFactory { - @Override - public Rule build() { - return logicalAggregate().thenApply(ctx -> { - LogicalAggregate aggregate = ctx.root; - List originalTargets = Lists.newArrayList(); - originalTargets.addAll(aggregate.getGroupByExpressions()); - originalTargets.addAll(aggregate.getOutputExpressions()); - List targets = materializeNestedMapInputs(originalTargets); - Optional, LogicalProject>> rewrittenOpt - = rewriteExpressions(aggregate, targets); - Plan newChild = rewrittenOpt.isPresent() - ? rewrittenOpt.get().second : aggregate.child(); - List newTargets = rewrittenOpt.isPresent() - ? rewrittenOpt.get().first : targets; - if (!rewrittenOpt.isPresent() && newTargets.equals(originalTargets)) { - return aggregate; - } - // rewriteExpressions treats group-by expressions and outputs as one ordered list - // so a common Map input is materialized only once. Restore the two original lists - // after replacement. - int groupBySize = aggregate.getGroupByExpressions().size(); - ImmutableList newGroupBy = ImmutableList.copyOf( - newTargets.subList(0, groupBySize)); - ImmutableList.Builder newOutputBuilder - = ImmutableList.builderWithExpectedSize(aggregate.getOutputExpressions().size()); - for (int i = groupBySize; i < newTargets.size(); i++) { - newOutputBuilder.add((NamedExpression) newTargets.get(i)); - } - return aggregate.withChildGroupByAndOutput(newGroupBy, newOutputBuilder.build(), newChild); - }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); - } - } - - private class JoinRewrite extends OneRewriteRuleFactory { - @Override - public Rule build() { - return logicalJoin().thenApply(ctx -> { - LogicalJoin join = ctx.root; - int hashOtherConjunctsSize = join.getHashJoinConjuncts().size() - + join.getOtherJoinConjuncts().size(); - int totalConjunctsSize = hashOtherConjunctsSize + join.getMarkJoinConjuncts().size(); - List allConjuncts = Lists.newArrayListWithExpectedSize(totalConjunctsSize); - allConjuncts.addAll(join.getHashJoinConjuncts()); - allConjuncts.addAll(join.getOtherJoinConjuncts()); - allConjuncts.addAll(join.getMarkJoinConjuncts()); - List originalAllConjuncts = ImmutableList.copyOf(allConjuncts); - allConjuncts = materializeNestedMapInputs(allConjuncts); - Optional rewrittenOpt = rewriteJoinExpressions(join, allConjuncts); - if (!rewrittenOpt.isPresent() && allConjuncts.equals(originalAllConjuncts)) { - return join; - } - - Plan newLeftChild = rewrittenOpt.map(result -> result.left).orElse(join.left()); - Plan newRightChild = rewrittenOpt.map(result -> result.right).orElse(join.right()); - List newAllConjuncts = rewrittenOpt - .map(result -> result.newConjuncts).orElse(allConjuncts); - List newHashOtherConjuncts = newAllConjuncts.subList(0, hashOtherConjunctsSize); - List newMarkJoinConjuncts = ImmutableList.copyOf( - newAllConjuncts.subList(hashOtherConjunctsSize, totalConjunctsSize)); - - Pair, List> pair = JoinUtils.extractExpressionForHashTable( - newLeftChild.getOutput(), newRightChild.getOutput(), newHashOtherConjuncts); - List newHashJoinConjuncts = pair.first; - List newOtherJoinConjuncts = pair.second; - JoinType joinType = join.getJoinType(); - if (joinType == JoinType.CROSS_JOIN && !newHashJoinConjuncts.isEmpty()) { - joinType = JoinType.INNER_JOIN; - } - return new LogicalJoin<>(joinType, - newHashJoinConjuncts, - newOtherJoinConjuncts, - newMarkJoinConjuncts, - join.getDistributeHint(), - join.getMarkJoinSlotReference(), - ImmutableList.of(newLeftChild, newRightChild), - join.getJoinReorderContext()); - }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); - } - } - - /** - * Rewrite expressions owned by a single-child plan and install their materialization Projects. - * - *

It first materializes computed Map inputs and replaces them in {@code targets}. It then - * materializes any {@link MapEntryArrayMap} still used more than once. These are separate - * Project layers because the second expression can depend on a Map Slot created by the first. - * The returned pair contains the rewritten targets and the top materialization Project. - */ - private Optional, LogicalProject>> rewriteExpressions( - LogicalPlan plan, Collection targets) { - // computed map materialized - List mapInputAliases = tryGenMapInputAliases(targets); - List rewrittenTargets = replaceExpressions(targets, mapInputAliases); - // MapEntryArrayMap merteialized - List entryArrayAliases = tryGenSharedEntryArrayAliases(rewrittenTargets); - if (mapInputAliases.isEmpty() && entryArrayAliases.isEmpty()) { - return Optional.empty(); - } - - Plan child = plan.child(0); - if (!mapInputAliases.isEmpty()) { - child = appendProject(child, mapInputAliases); - } - if (!entryArrayAliases.isEmpty()) { - child = appendProject(child, entryArrayAliases); - rewrittenTargets = replaceExpressions(rewrittenTargets, entryArrayAliases); - } - - return Optional.of(Pair.of(rewrittenTargets, (LogicalProject) child)); - } - - /** Add aliases without hiding any output already produced by {@code child}. */ - private LogicalProject appendProject(Plan child, List aliases) { - List projects = ImmutableList.builder() - .addAll(child.getOutput()) - .addAll(aliases) - .build(); - return new LogicalProject<>(projects, child); - } - - /** Replace each aliased expression by its Slot in all target expression trees. */ - private List replaceExpressions( - Collection expressions, List aliases) { - if (aliases.isEmpty()) { - return ImmutableList.copyOf(expressions); - } - Map replaceMap = Maps.newHashMap(); - for (NamedExpression alias : aliases) { - replaceMap.put(alias.child(0), alias.toSlot()); - } - ImmutableList.Builder builder = ImmutableList.builderWithExpectedSize(expressions.size()); - for (T expression : expressions) { - builder.add((T) ExpressionUtils.replace(expression, replaceMap)); - } - return builder.build(); - } - - /** - * Rewrite Join conjuncts using the same two materialization stages as - * {@link #rewriteExpressions(LogicalPlan, Collection)}. - * - *

Unlike a single-child plan, each generated alias must be attached to the Join child that - * contains all its input Slots. An expression referencing both children cannot be evaluated in - * either child Project, so a deterministic expression is left unchanged and a volatile one is - * rejected. Entry-array aliases are assigned after Map aliases because they may use new Slots. - */ - private Optional rewriteJoinExpressions(LogicalJoin join, - Collection targets) { - List rewrittenTargets = ImmutableList.copyOf(targets); - Plan left = join.left(); - Plan right = join.right(); - - Map> mapInputSlots = Maps.newLinkedHashMap(); - for (Expression target : rewrittenTargets) { - Set mapInputs = Sets.newLinkedHashSet(); - collectMapInputs(target, mapInputs); - for (Expression mapInput : mapInputs) { - Set inputSlots = mapInput.getInputSlots(); - mapInputSlots.computeIfAbsent(mapInput, ignored -> Sets.newLinkedHashSet()) - .addAll(inputSlots.isEmpty() ? target.getInputSlots() : inputSlots); - } - } - - ImmutableList.Builder leftAliases = ImmutableList.builder(); - ImmutableList.Builder rightAliases = ImmutableList.builder(); - Map replaceMap = Maps.newHashMap(); - Set leftOutputSet = left.getOutputSet(); - Set rightOutputSet = right.getOutputSet(); - for (Entry> entry : mapInputSlots.entrySet()) { - Set inputSlots = entry.getValue(); - Set mapInputExpressionSlots = entry.getKey().getInputSlots(); - if (!mapInputExpressionSlots.isEmpty() - && !leftOutputSet.containsAll(inputSlots) - && !rightOutputSet.containsAll(inputSlots)) { - // No child Project can reference Slots from both sides. Recalculation is safe for - // a deterministic expression, but a volatile Map would no longer have one stable - // value shared by map_keys and map_values. - if (entry.getKey().containsVolatileExpression()) { - throw new AnalysisException( - "A computed Map input containing a volatile expression cannot " - + "reference both sides of a join"); - } - continue; - } - ExprId exprId = StatementScopeIdGenerator.newExprId(); - Alias alias = new Alias( - exprId, entry.getKey(), "$_map_input_" + exprId.asInt() + "_$"); - replaceMap.put(alias.child(0), alias.toSlot()); - if (!inputSlots.isEmpty() && rightOutputSet.containsAll(inputSlots)) { - rightAliases.add(alias); - } else { - leftAliases.add(alias); - } - } - if (!replaceMap.isEmpty()) { - List leftAliasList = leftAliases.build(); - List rightAliasList = rightAliases.build(); - left = appendProjectIfNeeded(left, leftAliasList); - right = appendProjectIfNeeded(right, rightAliasList); - rewrittenTargets = replaceExpressions(rewrittenTargets, - ImmutableList.builder() - .addAll(leftAliasList) - .addAll(rightAliasList) - .build()); - } - - List entryArrayAliases = tryGenSharedEntryArrayAliases(rewrittenTargets); - ImmutableList.Builder leftEntryAliases = ImmutableList.builder(); - ImmutableList.Builder rightEntryAliases = ImmutableList.builder(); - leftOutputSet = left.getOutputSet(); - rightOutputSet = right.getOutputSet(); - for (NamedExpression alias : entryArrayAliases) { - Expression entryArray = alias.child(0); - Set inputSlots = Sets.newLinkedHashSet(entryArray.getInputSlots()); - if (inputSlots.isEmpty()) { - // As with a slot-free Map, inherit the containing conjunct's scope only to choose - // a child. The expression itself remains valid on either side. - for (Expression target : rewrittenTargets) { - if (target.anyMatch(entryArray::equals)) { - inputSlots.addAll(target.getInputSlots()); - } - } - } - Set expressionSlots = entryArray.getInputSlots(); - if (!expressionSlots.isEmpty() - && !leftOutputSet.containsAll(inputSlots) - && !rightOutputSet.containsAll(inputSlots)) { - if (entryArray.containsVolatileExpression()) { - throw new AnalysisException( - "A shared Map entry array containing a volatile expression cannot " - + "reference both sides of a join"); - } - continue; - } - if (!inputSlots.isEmpty() && rightOutputSet.containsAll(inputSlots)) { - rightEntryAliases.add(alias); - } else { - leftEntryAliases.add(alias); - } - } - List leftEntryAliasList = leftEntryAliases.build(); - List rightEntryAliasList = rightEntryAliases.build(); - if (!leftEntryAliasList.isEmpty() || !rightEntryAliasList.isEmpty()) { - left = appendProjectIfNeeded(left, leftEntryAliasList); - right = appendProjectIfNeeded(right, rightEntryAliasList); - rewrittenTargets = replaceExpressions(rewrittenTargets, - ImmutableList.builder() - .addAll(leftEntryAliasList) - .addAll(rightEntryAliasList) - .build()); - } - - if (replaceMap.isEmpty() && leftEntryAliasList.isEmpty() && rightEntryAliasList.isEmpty()) { - return Optional.empty(); - } - return Optional.of(new JoinRewriteResult(rewrittenTargets, left, right)); - } - - /** Avoid creating an identity Project when one side of a Join has no aliases. */ - private Plan appendProjectIfNeeded(Plan child, List aliases) { - if (aliases.isEmpty()) { - return child; - } - List projects = ImmutableList.builder() - .addAll(child.getOutput()) - .addAll(aliases) - .build(); - return new LogicalProject<>(projects, child); - } - - /** - * Find and alias each distinct computed Map consumed by a {@link MapEntryArrayMap}. - * - *

This method turns the expressions found by {@link #collectMapInputs(Expression, Set)} into - * aliases. Slots and Map literals are excluded because they need no materialization. - */ - private List tryGenMapInputAliases( - Collection targets) { - Set mapInputs = Sets.newLinkedHashSet(); - for (Expression target : targets) { - collectMapInputs(target, mapInputs); - } - - ImmutableList.Builder aliases - = ImmutableList.builderWithExpectedSize(mapInputs.size()); - for (Expression mapInput : mapInputs) { - ExprId exprId = StatementScopeIdGenerator.newExprId(); - aliases.add(new Alias(exprId, mapInput, "$_map_input_" + exprId.asInt() + "_$")); - } - return aliases.build(); - } - - /** - * Find repeated {@link MapEntryArrayMap} expressions and create one shared alias for each. - * - *

This is used by the current safe lowering of {@code map_apply}; the implementation does - * not use the optional fast lowering into two independent two-parameter ArrayMaps. The original - * two-parameter lambda is evaluated first and produces {@code ARRAY<STRUCT>}: - * - * 

-     * mappedEntries = MapEntryArrayMap(
-     *   (mapKey, mapValue) -> struct(newKey, newValue),
-     *   map_keys(inputMap), map_values(inputMap))
-     * map_from_arrays(
-     *   array_map(mappedEntry -> mappedEntry[1], mappedEntries),
-     *   array_map(mappedEntry -> mappedEntry[2], mappedEntries))
-     *
- * - *

The two extraction ArrayMaps have one parameter because they iterate the resulting Struct - * array, not the original Map. They do not copy or reevaluate the original lambda body. - */ - private List tryGenSharedEntryArrayAliases( - Collection targets) { - Map entryArrayCounts = Maps.newLinkedHashMap(); - for (Expression target : targets) { - collectEntryArrayCounts(target, entryArrayCounts); - } - - ImmutableList.Builder aliases = ImmutableList.builder(); - for (Entry entry : entryArrayCounts.entrySet()) { - if (entry.getValue() > 1) { - ExprId exprId = StatementScopeIdGenerator.newExprId(); - aliases.add(new Alias( - exprId, entry.getKey(), "$_map_entries_" + exprId.asInt() + "_$")); - } - } - return aliases.build(); - } - - /** Apply nested-lambda materialization independently to every target expression. */ - private List materializeNestedMapInputs(Collection expressions) { - ImmutableList.Builder builder = ImmutableList.builderWithExpectedSize(expressions.size()); - for (T expression : expressions) { - builder.add((T) materializeNestedMapInputs(expression)); - } - return builder.build(); - } - - /** - * Materialize computed Maps that depend on lambda item Slots inside the owning ArrayMap. - * - *

Consider: - * - * 

-     * select transform_values(
-     *   (outer_k, outer_v) -> transform_values((inner_k, inner_v) -> inner_k, map(outer_k + random(), outer_v)),
-     *   map(1, 10));
-     *
- * - * A relation Project cannot evaluate {@code map(outer_k + random(), outer_v)} because {@code outer_k} and - * {@code outer_v} exist only while the outer lambda is running. The outer ArrayMap is rewritten to - * carry a hidden array whose item is that Map: - * - * 
-     * outer inputs before:
-     *   outer_k <- map_keys(outerMap)
-     *   outer_v <- map_values(outerMap)
-     *
-     * outer inputs after:
-     *   outer_k <- map_keys(outerMap)
-     *   outer_v <- map_values(outerMap)
-     *   materializedInnerMap
-     *      - array_map((outerKey, outerValue) -> map(outerKey + random(), outerValue),
-     *                   map_keys(outerMap), map_values(outerMap))
-     *
-     * outer body after:
-     *   transform_values((innerKey, innerValue) -> innerKey, materializedInnerMap)
-     *
- * - *

Traversal is bottom-up. For each ArrayMap, computed Maps in its body become hidden input - * arrays; repeated entry arrays are handled afterward because they may use those hidden inputs. - */ - private Expression materializeNestedMapInputs(Expression expression) { - ImmutableList.Builder children - = ImmutableList.builderWithExpectedSize(expression.arity()); - boolean changed = false; - for (Expression child : expression.children()) { - Expression rewrittenChild = materializeNestedMapInputs(child); - children.add(rewrittenChild); - changed |= rewrittenChild != child; - } - Expression rewritten = changed ? expression.withChildren(children.build()) : expression; - if (!(rewritten instanceof ArrayMap)) { - return rewritten; - } - - Lambda lambda = (Lambda) rewritten.child(0); - Set mapInputs = Sets.newLinkedHashSet(); - collectMapInputs(lambda.getLambdaFunction(), mapInputs); - - List sourceArguments = lambda.getLambdaArguments(); - List argumentNames = Lists.newArrayList(lambda.getLambdaArgumentNames()); - List arguments = Lists.newArrayList(sourceArguments); - Expression lambdaBody = lambda.getLambdaFunction(); - for (Expression mapInput : mapInputs) { - Pair materialized = buildLambdaMaterializer(mapInput, sourceArguments); - ArrayItemReference hiddenArgument = new ArrayItemReference(materialized.second, materialized.first); - argumentNames.add(materialized.second); - arguments.add(hiddenArgument); - Map replaceMap = Maps.newHashMap(); - replaceMap.put(mapInput, hiddenArgument.toSlot()); - lambdaBody = ExpressionUtils.replace(lambdaBody, replaceMap); - } - - List entryArrayAliases = tryGenSharedEntryArrayAliases( - ImmutableList.of(lambdaBody)); - for (NamedExpression entryArrayAlias : entryArrayAliases) { - Expression entryArray = entryArrayAlias.child(0); - Pair materialized = buildLambdaMaterializer(entryArray, arguments); - ArrayItemReference hiddenArgument = new ArrayItemReference(materialized.second, materialized.first); - argumentNames.add(materialized.second); - arguments.add(hiddenArgument); - Map replaceMap = Maps.newHashMap(); - replaceMap.put(entryArray, hiddenArgument.toSlot()); - lambdaBody = ExpressionUtils.replace(lambdaBody, replaceMap); - } - if (arguments.size() == sourceArguments.size()) { - return rewritten; - } - - // A shared entry-array materializer can embed an earlier Map materializer. In that case the - // final body references only the entry-array argument. Keep all user arguments, but remove - // optimizer-added arguments no longer referenced by the final body to avoid evaluating the - // embedded Map expression a second time. - Set referencedArgumentIds = collectReferencedArgumentIds(lambdaBody); - ImmutableList.Builder retainedNames = ImmutableList.builder(); - ImmutableList.Builder retainedArguments = ImmutableList.builder(); - for (int i = 0; i < arguments.size(); i++) { - ArrayItemReference argument = arguments.get(i); - if (i < sourceArguments.size() - || referencedArgumentIds.contains(argument.getExprId())) { - retainedNames.add(argumentNames.get(i)); - retainedArguments.add(argument); - } - } - return rewritten.withChildren(ImmutableList.of( - new Lambda(retainedNames.build(), lambdaBody, retainedArguments.build()))); - } - - /** - * Build an ArrayMap that evaluates {@code expression} once per entry of the enclosing lambda. - * - *

Only enclosing arguments referenced by the expression are forwarded. For - * {@code map(ok + random(), ov)}, the generated lambda receives copies of {@code ok} and - * {@code ov}, with fresh ExprIds, and its body is rebound to those copies. If the expression - * only captures relation Slots, one enclosing array is still forwarded as a row-count and - * offset driver; all arrays of one ArrayMap have identical entry offsets. - * - * @return the materializing ArrayMap and the name of the hidden item argument that will expose - * each materialized result to the original lambda body - */ - private Pair buildLambdaMaterializer( - Expression expression, List sourceArguments) { - Set referencedArgumentIds = collectReferencedArgumentIds(expression); - - List selectedArguments = sourceArguments.stream() - .filter(argument -> referencedArgumentIds.contains(argument.getExprId())) - .collect(ImmutableList.toImmutableList()); - if (selectedArguments.isEmpty()) { - // ArrayMap needs an array to define the entry count even when the expression only - // captures relation slots. Any current lambda input has the same entry offsets. - selectedArguments = ImmutableList.of(sourceArguments.get(0)); - } - - Map replaceMap = Maps.newHashMap(); - ImmutableList.Builder materializerNames - = ImmutableList.builderWithExpectedSize(selectedArguments.size()); - ImmutableList.Builder materializerArguments - = ImmutableList.builderWithExpectedSize(selectedArguments.size()); - for (ArrayItemReference sourceArgument : selectedArguments) { - ExprId exprId = StatementScopeIdGenerator.newExprId(); - String name = "$_map_materialize_arg_" + exprId.asInt() + "_$"; - ArrayItemReference materializerArgument = new ArrayItemReference( - exprId, name, sourceArgument.getArrayExpression()); - materializerNames.add(name); - materializerArguments.add(materializerArgument); - replaceMap.put(sourceArgument.toSlot(), materializerArgument.toSlot()); - } - - Expression materializerBody = ExpressionUtils.replace(expression, replaceMap); - Lambda materializerLambda = new Lambda( - materializerNames.build(), materializerBody, materializerArguments.build()); - ExprId hiddenExprId = StatementScopeIdGenerator.newExprId(); - String hiddenName = "$_map_input_" + hiddenExprId.asInt() + "_$"; - return Pair.of(new ArrayMap(materializerLambda), hiddenName); - } - - /** Return ExprIds of lambda item Slots referenced by an expression. */ - private Set collectReferencedArgumentIds(Expression expression) { - Set referencedArgumentIds = Sets.newHashSet(); - expression.foreach(node -> { - if (node instanceof ArrayItemSlot) { - referencedArgumentIds.add(((ArrayItemSlot) node).getExprId()); - } - }); - return referencedArgumentIds; - } - - /** Traverse an expression and collect the Map input of every {@link MapEntryArrayMap} marker. */ - private void collectMapInputs(Expression expression, Set mapInputs) { - MapEntryArrayMap marker = unwrapMarker(expression); - if (marker != null) { - Lambda lambda = (Lambda) marker.child(0); - addMapInput(MapLambdaValidator.extractMapExpression("map lambda", lambda), mapInputs); - return; - } - - if (expression instanceof Lambda) { - for (ArrayItemReference argument : ((Lambda) expression).getLambdaArguments()) { - collectMapInputs(argument.getArrayExpression(), mapInputs); - } - return; - } - for (Expression child : expression.children()) { - collectMapInputs(child, mapInputs); - } - } - - /** Add only Maps whose key/value expansion would otherwise repeat computation. */ - private void addMapInput(Expression mapInput, Set mapInputs) { - if (MapLambdaValidator.requiresSingleEvaluation(mapInput)) { - mapInputs.add(mapInput); - } - } - - /** - * Count each complete {@link MapEntryArrayMap} expression for - * {@link #tryGenSharedEntryArrayAliases(Collection)}. As in {@code collectMapInputs}, only - * lambda argument arrays are traversed across a Lambda boundary. - */ - private void collectEntryArrayCounts(Expression expression, Map counts) { - if (unwrapMarker(expression) != null) { - counts.merge(expression, 1, Integer::sum); - return; - } - if (expression instanceof Lambda) { - for (ArrayItemReference argument : ((Lambda) expression).getLambdaArguments()) { - collectEntryArrayCounts(argument.getArrayExpression(), counts); - } - return; - } - for (Expression child : expression.children()) { - collectEntryArrayCounts(child, counts); - } - } - - /** Find the Map entry marker through analyzer-inserted Cast wrappers. */ - private MapEntryArrayMap unwrapMarker(Expression expression) { - while (expression instanceof Cast) { - expression = expression.child(0); - } - return expression instanceof MapEntryArrayMap ? (MapEntryArrayMap) expression : null; - } - - private static class JoinRewriteResult { - private final List newConjuncts; - private final Plan left; - private final Plan right; - - private JoinRewriteResult(List newConjuncts, Plan left, Plan right) { - this.newConjuncts = newConjuncts; - this.left = left; - this.right = right; - } - } -} 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 0b33402db718dd..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 @@ -24,16 +24,12 @@ import org.apache.doris.nereids.types.ArrayType; import org.apache.doris.nereids.types.DataType; import org.apache.doris.nereids.types.LambdaType; -import org.apache.doris.nereids.types.MapType; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableList.Builder; -import com.google.common.collect.ImmutableSet; import java.util.List; -import java.util.Locale; import java.util.Objects; -import java.util.Set; import java.util.stream.Collectors; /** @@ -42,15 +38,6 @@ * After bind, x -> x : arguments("x") -> children: Expression(x) ArrayItemReference(x) */ public class Lambda extends Expression { - - private static final Set MAP_ENTRY_LAMBDA_FUNCTIONS = ImmutableSet.of( - "map_all", - "map_apply", - "map_exists", - "map_filter", - "transform_keys", - "transform_values"); - private final List argumentNames; /** @@ -78,28 +65,6 @@ public Lambda(List argumentNames, List children) { */ public ImmutableList makeArguments(String functionName, List lambdaArgs) { Builder builder = new ImmutableList.Builder<>(); - String normalizedFunctionName = functionName.toLowerCase(Locale.ROOT); - if (MAP_ENTRY_LAMBDA_FUNCTIONS.contains(normalizedFunctionName)) { - if (lambdaArgs.size() != 1) { - throw new AnalysisException(String.format( - "%s requires exactly one map argument but has %d", - functionName, lambdaArgs.size())); - } - if (argumentNames.size() != 2) { - throw new AnalysisException(String.format( - "lambda of %s requires exactly two arguments but has %d", - functionName, argumentNames.size())); - } - Expression mapExpression = lambdaArgs.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())); - } - builder.add(new ArrayItemReference(argumentNames.get(0), new MapKeys(mapExpression))); - builder.add(new ArrayItemReference(argumentNames.get(1), new MapValues(mapExpression))); - return builder.build(); - } 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") && lambdaArgs.size() == 1 && argumentNames.size() == 2) { 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 index a9007048b54d17..c703fad6e7dddc 100644 --- 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 @@ -40,8 +40,8 @@ * -> * array_match_all( * array_map( - * (mapKey, mapValue) -> predicate, - * map_keys(inputMap), map_values(inputMap))) + * entry -> predicate(entry[1], entry[2]), + * map_entries(inputMap))) * */ public class MapAll extends ScalarFunction @@ -52,11 +52,15 @@ public class MapAll extends ScalarFunction // The argument is a bound Lambda. public MapAll(Expression arg) { - this(MapLambdaValidator.requireLambda("map_all", arg)); + this(MapLambdaFunctionUtils.requireLambda("map_all", arg)); } private MapAll(Lambda lambda) { - super("map_all", new MapEntryArrayMap(lambda)); + this(MapLambdaFunctionUtils.rewrite(lambda, (body, key, value, entry) -> body)); + } + + private MapAll(MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda) { + super("map_all", rewrittenLambda.toArrayMap()); } private MapAll(ScalarFunctionParams functionParams) { 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 index 224ad8e6ac50dc..46cf7f7b775f20 100644 --- 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 @@ -19,13 +19,11 @@ import org.apache.doris.catalog.FunctionSignature; import org.apache.doris.nereids.exceptions.AnalysisException; -import org.apache.doris.nereids.trees.expressions.Cast; 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.shape.UnaryExpression; import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; import org.apache.doris.nereids.types.ArrayType; import org.apache.doris.nereids.types.DataType; @@ -48,21 +46,24 @@ * map_apply((mapKey, mapValue) -> struct(newKey, newValue), inputMap) * -> * map_from_entries(array_map( - * (mapKey, mapValue) -> struct(newKey, newValue), - * map_keys(inputMap), map_values(inputMap))) + * entry -> struct(newKey(entry[1], entry[2]), newValue(entry[1], entry[2])), + * map_entries(inputMap))) * */ public class MapApply extends ScalarFunction - implements UnaryExpression, CustomSignature, PropagateNullable, PreferPushDownProject, + implements CustomSignature, PropagateNullable, PreferPushDownProject, RewriteWhenAnalyze { public MapApply(Expression arg) { - this(MapLambdaValidator.requireLambda("map_apply", arg)); + this(MapLambdaFunctionUtils.requireLambda("map_apply", arg)); } private MapApply(Lambda lambda) { - super("map_apply", new MapEntryArrayMap(lambda)); - validateLambdaReturn(lambda); + this(validateAndRewrite(lambda)); + } + + private MapApply(MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda) { + super("map_apply", rewrittenLambda.getMapExpression(), rewrittenLambda.toArrayMap()); } private MapApply(ScalarFunctionParams functionParams) { @@ -71,7 +72,7 @@ private MapApply(ScalarFunctionParams functionParams) { @Override public FunctionSignature customSignature() { - DataType mappedEntriesType = getArgument(0).getDataType(); + DataType mappedEntriesType = getArgument(1).getDataType(); if (!(mappedEntriesType instanceof ArrayType) || !(((ArrayType) mappedEntriesType).getItemType() instanceof StructType)) { throw invalidReturnType(); @@ -80,23 +81,23 @@ public FunctionSignature customSignature() { if (structType.getFields().size() != 2) { throw invalidReturnType(); } - MapType inputMapType = extractInputMapType(getEntryLambda(getArgument(0))); + MapType inputMapType = (MapType) getArgument(0).getDataType(); StructType resolvedStructType = resolveNullFieldTypes(structType, inputMapType); List fields = resolvedStructType.getFields(); MapType resultType = MapType.of(fields.get(0).getDataType(), fields.get(1).getDataType()); resultType.validateDataType(); - return FunctionSignature.ret(resultType).args(ArrayType.of(resolvedStructType)); + return FunctionSignature.ret(resultType).args(inputMapType, ArrayType.of(resolvedStructType)); } @Override public MapApply withChildren(List children) { - Preconditions.checkArgument(children.size() == 1); + Preconditions.checkArgument(children.size() == 2); return new MapApply(getFunctionParams(children)); } @Override public Expression rewriteWhenAnalyze() { - return new MapFromEntries(getArgument(0)); + return new MapFromEntries(getArgument(1)); } @Override @@ -104,7 +105,14 @@ public R accept(ExpressionVisitor visitor, C context) { return visitor.visitMapApply(this, context); } - private static void validateLambdaReturn(Lambda lambda) { + private static MapLambdaFunctionUtils.RewrittenMapLambda validateAndRewrite(Lambda lambda) { + MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda = MapLambdaFunctionUtils.rewrite( + lambda, (body, key, value, entry) -> body); + validateLambdaReturn(lambda, (MapType) rewrittenLambda.getMapExpression().getDataType()); + return rewrittenLambda; + } + + private static void validateLambdaReturn(Lambda lambda, MapType inputMapType) { Expression lambdaBody = lambda.getLambdaFunction(); if (!(lambdaBody.getDataType() instanceof StructType) || ((StructType) lambdaBody.getDataType()).getFields().size() != 2 @@ -112,26 +120,11 @@ private static void validateLambdaReturn(Lambda lambda) { throw invalidReturnType(); } StructType structType = (StructType) lambdaBody.getDataType(); - StructType resolvedStructType = resolveNullFieldTypes(structType, extractInputMapType(lambda)); + StructType resolvedStructType = resolveNullFieldTypes(structType, inputMapType); MapType.of(resolvedStructType.getFields().get(0).getDataType(), resolvedStructType.getFields().get(1).getDataType()).validateDataType(); } - private static MapType extractInputMapType(Lambda lambda) { - return (MapType) MapLambdaValidator.extractMapExpression("map_apply", lambda).getDataType(); - } - - private static Lambda getEntryLambda(Expression mappedEntries) { - while (mappedEntries instanceof Cast) { - mappedEntries = mappedEntries.child(0); - } - if (!(mappedEntries instanceof MapEntryArrayMap) - || !(mappedEntries.child(0) instanceof Lambda)) { - throw invalidReturnType(); - } - return (Lambda) mappedEntries.child(0); - } - // Resolve only untyped fields in the two-field struct returned by the lambda. For // map_apply((k, v) -> struct(cast(k as bigint), []), map(1, [10])), the result is // MAP>. Keep the explicit BIGINT type and infer only the empty array @@ -140,9 +133,9 @@ private static StructType resolveNullFieldTypes(StructType structType, MapType i List fields = structType.getFields(); StructField keyField = fields.get(0); StructField valueField = fields.get(1); - keyField = keyField.withDataType(MapLambdaValidator.mergeNestedNullTypes( + keyField = keyField.withDataType(MapLambdaFunctionUtils.mergeNestedNullTypes( keyField.getDataType(), inputMapType.getKeyType())); - valueField = valueField.withDataType(MapLambdaValidator.mergeNestedNullTypes( + valueField = valueField.withDataType(MapLambdaFunctionUtils.mergeNestedNullTypes( valueField.getDataType(), inputMapType.getValueType())); return new StructType(ImmutableList.of(keyField, valueField)); } 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 index 17c99c6751b7c0..36687cf81be88d 100644 --- 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 @@ -40,8 +40,8 @@ * -> * array_match_any( * array_map( - * (mapKey, mapValue) -> predicate, - * map_keys(inputMap), map_values(inputMap))) + * entry -> predicate(entry[1], entry[2]), + * map_entries(inputMap))) * */ public class MapExists extends ScalarFunction @@ -52,11 +52,15 @@ public class MapExists extends ScalarFunction /** Constructor with a bound Lambda argument. */ public MapExists(Expression arg) { - this(MapLambdaValidator.requireLambda("map_exists", arg)); + this(MapLambdaFunctionUtils.requireLambda("map_exists", arg)); } private MapExists(Lambda lambda) { - super("map_exists", new MapEntryArrayMap(lambda)); + this(MapLambdaFunctionUtils.rewrite(lambda, (body, key, value, entry) -> body)); + } + + private MapExists(MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda) { + super("map_exists", rewrittenLambda.toArrayMap()); } private MapExists(ScalarFunctionParams functionParams) { 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 index 23c48027f4644f..b5613b0ebe7a24 100644 --- 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 @@ -20,6 +20,8 @@ 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; @@ -34,30 +36,33 @@ /** * Scalar function map_filter. * - *

The Map lambda is evaluated by an ArrayMap over the Map's key and value arrays: + *

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

  * map_filter((mapKey, mapValue) -> predicate, inputMap)
  *   ->
- * map_filter(
- *   inputMap,
+ * %map_from_filtered_entries_unique%(
  *   array_map(
- *     (mapKey, mapValue) -> predicate,
- *     map_keys(inputMap), map_values(inputMap)))
+ *     entry -> if(predicate(entry[1], entry[2]), entry, null),
+ *     map_entries(inputMap)))
  *
*/ public class MapFilter extends ScalarFunction - implements HighOrderFunction, PropagateNullable { + 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(MapLambdaValidator.requireLambda("map_filter", arg)); + this(MapLambdaFunctionUtils.requireLambda("map_filter", arg)); } public MapFilter(Expression map, Expression filter) { @@ -66,9 +71,14 @@ public MapFilter(Expression map, Expression filter) { } 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", - MapLambdaValidator.extractMapExpression("map_filter", lambda), - new MapEntryArrayMap(lambda)); + rewrittenLambda.getMapExpression(), rewrittenLambda.toArrayMap()); validateMapLambdaInput = true; } @@ -77,10 +87,6 @@ private MapFilter(ScalarFunctionParams functionParams, boolean validateMapLambda this.validateMapLambdaInput = validateMapLambdaInput; } - public boolean shouldValidateMapLambdaInput() { - return validateMapLambdaInput; - } - @Override public MapFilter withChildren(List children) { Preconditions.checkArgument(children.size() == 2); @@ -89,7 +95,14 @@ public MapFilter withChildren(List children) { @Override public List getImplSignature() { - return SIGNATURES; + return validateMapLambdaInput ? MAP_LAMBDA_SIGNATURES : SIGNATURES; + } + + @Override + public Expression rewriteWhenAnalyze() { + return validateMapLambdaInput + ? new MapFromFilteredEntriesUnique(getArgument(1)) + : this; } @Override 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/MapFromArraysUnique.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromEntriesUnique.java similarity index 66% rename from fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromArraysUnique.java rename to fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromEntriesUnique.java index bd3207a2d5634a..24176e0b6a28e8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromArraysUnique.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromEntriesUnique.java @@ -23,20 +23,20 @@ import java.util.List; -/** Internal Map constructor used only when the input keys are known to be unique. */ -public class MapFromArraysUnique extends MapFromArrays { +/** Internal Map constructor used when entry keys are known to be unique. */ +public class MapFromEntriesUnique extends MapFromEntries { - public MapFromArraysUnique(Expression keys, Expression values) { - super("%map_from_arrays_unique%", keys, values); + public MapFromEntriesUnique(Expression entries) { + super("%map_from_entries_unique%", entries); } - private MapFromArraysUnique(ScalarFunctionParams functionParams) { + private MapFromEntriesUnique(ScalarFunctionParams functionParams) { super(functionParams); } @Override - public MapFromArraysUnique withChildren(List children) { - Preconditions.checkArgument(children.size() == 2); - return new MapFromArraysUnique(getFunctionParams(children)); + 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/MapEntryArrayMap.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromFilteredEntriesUnique.java similarity index 62% rename from fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapEntryArrayMap.java rename to fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromFilteredEntriesUnique.java index 87e88521f88536..1f7cd306299f78 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapEntryArrayMap.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapFromFilteredEntriesUnique.java @@ -23,16 +23,20 @@ import java.util.List; -/** Marks an ArrayMap whose first two arguments evaluate the key and value of each Map entry. */ -public final class MapEntryArrayMap extends ArrayMap { +/** Internal Map constructor that drops null entries produced by a Map-filter Lambda. */ +public class MapFromFilteredEntriesUnique extends MapFromEntries { - MapEntryArrayMap(Lambda lambda) { - super(lambda); + public MapFromFilteredEntriesUnique(Expression entries) { + super("%map_from_filtered_entries_unique%", entries); + } + + private MapFromFilteredEntriesUnique(ScalarFunctionParams functionParams) { + super(functionParams); } @Override - public MapEntryArrayMap withChildren(List children) { - Preconditions.checkArgument(children.size() == 1 && children.get(0) instanceof Lambda); - return new MapEntryArrayMap((Lambda) children.get(0)); + 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..55da184752161c --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaFunctionUtils.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.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 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 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; + } + + /** Fill only NULL_TYPE positions from the corresponding input Map field type. */ + static DataType mergeNestedNullTypes(DataType outputType, DataType inputType) { + if (outputType.isNullType()) { + return inputType; + } else if (outputType instanceof ArrayType && inputType instanceof ArrayType) { + return ArrayType.of(mergeNestedNullTypes( + ((ArrayType) outputType).getItemType(), ((ArrayType) inputType).getItemType())); + } else if (outputType instanceof MapType && inputType instanceof MapType) { + return MapType.of( + mergeNestedNullTypes( + ((MapType) outputType).getKeyType(), ((MapType) inputType).getKeyType()), + mergeNestedNullTypes( + ((MapType) outputType).getValueType(), ((MapType) inputType).getValueType())); + } else if (outputType instanceof StructType && inputType instanceof StructType) { + List outputFields = ((StructType) outputType).getFields(); + List inputFields = ((StructType) inputType).getFields(); + if (outputFields.size() != inputFields.size()) { + return outputType; + } + ImmutableList.Builder fields + = ImmutableList.builderWithExpectedSize(outputFields.size()); + for (int i = 0; i < outputFields.size(); i++) { + fields.add(outputFields.get(i).withDataType(mergeNestedNullTypes( + outputFields.get(i).getDataType(), inputFields.get(i).getDataType()))); + } + return new StructType(fields.build()); + } + return outputType; + } + + 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/MapLambdaValidator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaValidator.java deleted file mode 100644 index 32e06bf6c1f3f1..00000000000000 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/MapLambdaValidator.java +++ /dev/null @@ -1,168 +0,0 @@ -// 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.Cast; -import org.apache.doris.nereids.trees.expressions.Expression; -import org.apache.doris.nereids.trees.expressions.Slot; -import org.apache.doris.nereids.trees.expressions.literal.MapLiteral; -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.collect.ImmutableList; - -import java.util.List; - -/** - * Validates the internal ArrayMap used to evaluate map entries. - */ -public final class MapLambdaValidator { - - private MapLambdaValidator() { - } - - // Require a bound lambda argument. - public 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; - } - - // A Map lambda expands one Map into the original Map, map_keys(map), and map_values(map). - // Computed Maps are candidates for materialization. Slots are already materialized, and Map - // literals do not contain repeated computation. - public static boolean requiresSingleEvaluation(Expression mapExpression) { - while (mapExpression instanceof Cast) { - mapExpression = mapExpression.child(0); - } - return !(mapExpression instanceof Slot || mapExpression instanceof MapLiteral); - } - - /** - * Validate and return the Map shared by the first key and value references. - */ - public static Expression extractMapExpression(String functionName, Lambda lambda) { - List arguments = lambda.getLambdaArguments(); - if (arguments.size() < 2) { - throw new AnalysisException(String.format( - "Internal map entry lambda of %s must have key and value inputs", functionName)); - } - Expression keyArray = arguments.get(0).getArrayExpression(); - Expression valueArray = arguments.get(1).getArrayExpression(); - if (!(keyArray instanceof MapKeys) || !(valueArray instanceof MapValues)) { - throw new AnalysisException(String.format( - "Internal map entry lambda of %s must have key and value inputs", functionName)); - } - Expression keyMap = keyArray.child(0); - Expression valueMap = valueArray.child(0); - if (!keyMap.equals(valueMap)) { - throw new AnalysisException(String.format( - "Map entry inputs of %s must come from the same map", functionName)); - } - return keyMap; - } - - // Fill only NULL_TYPE positions in a Map lambda result from the matching input key or value - // type. For transform_values((k, v) -> [], map(1, [10])), infer ARRAY for the empty - // array. Apply the same inference to nested Array, Struct, and Map types. - static DataType mergeNestedNullTypes(DataType outputType, DataType inputType) { - if (outputType.isNullType()) { - return inputType; - } else if (outputType instanceof ArrayType && inputType instanceof ArrayType) { - return ArrayType.of(mergeNestedNullTypes( - ((ArrayType) outputType).getItemType(), ((ArrayType) inputType).getItemType())); - } else if (outputType instanceof MapType && inputType instanceof MapType) { - return MapType.of( - mergeNestedNullTypes( - ((MapType) outputType).getKeyType(), ((MapType) inputType).getKeyType()), - mergeNestedNullTypes( - ((MapType) outputType).getValueType(), ((MapType) inputType).getValueType())); - } else if (outputType instanceof StructType && inputType instanceof StructType) { - List outputFields = ((StructType) outputType).getFields(); - List inputFields = ((StructType) inputType).getFields(); - if (outputFields.size() != inputFields.size()) { - return outputType; - } - ImmutableList.Builder fields - = ImmutableList.builderWithExpectedSize(outputFields.size()); - for (int i = 0; i < outputFields.size(); i++) { - fields.add(outputFields.get(i).withDataType(mergeNestedNullTypes( - outputFields.get(i).getDataType(), inputFields.get(i).getDataType()))); - } - return new StructType(fields.build()); - } - return outputType; - } - - /** - * Revalidate the hidden physical arrays after optimizer rewrites. - */ - public static void validateStablePhysicalInputs(String functionName, Lambda lambda) { - List arguments = lambda.getLambdaArguments(); - if (arguments.isEmpty()) { - throw new AnalysisException(String.format( - "Internal map entry lambda of %s must have key and value inputs", functionName)); - } - // Projection CSE can replace only one of map_keys(M) and map_values(M) with a Slot. The - // analysis-time constructor already checked their common Map lineage, so repeating that - // structural check here would reject a valid partially materialized marker. Translation - // only needs the key/value driver arrays to be stable; ArrayMap checks equal lengths. - // Additional arguments are hidden arrays used to materialize nested lambda expressions. - // They are deliberately allowed to be volatile because each hidden array is evaluated - // once and then consumed through its item Slot by the owning lambda. - int driverCount = Math.min(2, arguments.size()); - for (int i = 0; i < driverCount; i++) { - ArrayItemReference argument = arguments.get(i); - if (argument.getArrayExpression().containsVolatileExpression()) { - throw new AnalysisException(String.format( - "Internal map entry input of %s must be materialized before translation", - functionName)); - } - } - } - - /** - * Validate functions that consume both the original Map and a mapped entry array. - */ - public static void validateOuterMapConsumer(String functionName, Expression mappedArray) { - Expression marker = mappedArray; - while (marker instanceof Cast) { - marker = marker.child(0); - } - // CSE can materialize the whole MapEntryArrayMap in an earlier projection layer when the - // enclosing Map function is referenced more than once. The marker was validated while - // translating that layer, so its projected slot is a valid physical input here. - if (marker instanceof Slot) { - return; - } - if (!(marker instanceof MapEntryArrayMap)) { - throw new AnalysisException(String.format( - "Mapped entry input of %s lost its internal map entry marker", functionName)); - } - Lambda lambda = (Lambda) marker.child(0); - validateStablePhysicalInputs(functionName, 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 index c4130ee30310bb..9637b27b78d50e 100644 --- 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 @@ -27,8 +27,11 @@ 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 com.google.common.collect.ImmutableList; import java.util.List; @@ -37,13 +40,17 @@ public class TransformKeys extends ScalarFunction implements CustomSignature, PropagateNullable, PreferPushDownProject, RewriteWhenAnalyze { public TransformKeys(Expression arg) { - this(MapLambdaValidator.requireLambda("transform_keys", 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", - MapLambdaValidator.extractMapExpression("transform_keys", lambda), - new MapEntryArrayMap(lambda)); + rewrittenLambda.getMapExpression(), rewrittenLambda.toArrayMap()); } private TransformKeys(ScalarFunctionParams functionParams) { @@ -53,15 +60,18 @@ private TransformKeys(ScalarFunctionParams functionParams) { @Override public FunctionSignature customSignature() { MapType inputMapType = (MapType) getArgument(0).getDataType(); - ArrayType transformedKeysType = (ArrayType) getArgument(1).getDataType(); + ArrayType mappedEntriesType = (ArrayType) getArgument(1).getDataType(); + StructType entryType = (StructType) mappedEntriesType.getItemType(); + List fields = entryType.getFields(); // transform_keys((k, v) -> null, map(1, 10)) // res_type should be: MAP instead of MAP - DataType resultKeyType = MapLambdaValidator.mergeNestedNullTypes( - transformedKeysType.getItemType(), inputMapType.getKeyType()); - transformedKeysType = ArrayType.of(resultKeyType); + DataType resultKeyType = MapLambdaFunctionUtils.mergeNestedNullTypes( + fields.get(0).getDataType(), inputMapType.getKeyType()); + StructType resolvedEntryType = new StructType(ImmutableList.of( + fields.get(0).withDataType(resultKeyType), fields.get(1))); MapType resultType = MapType.of(resultKeyType, inputMapType.getValueType()); resultType.validateDataType(); - return FunctionSignature.ret(resultType).args(inputMapType, transformedKeysType); + return FunctionSignature.ret(resultType).args(inputMapType, ArrayType.of(resolvedEntryType)); } @Override @@ -72,7 +82,7 @@ public TransformKeys withChildren(List children) { @Override public Expression rewriteWhenAnalyze() { - return new MapFromArrays(getArgument(1), new MapValues(getArgument(0))); + return new MapFromEntries(getArgument(1)); } @Override 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 index 70871fb55ee683..e27c6e4650379a 100644 --- 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 @@ -27,8 +27,11 @@ 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 com.google.common.collect.ImmutableList; import java.util.List; @@ -40,24 +43,27 @@ * 
  * transform_values((mapKey, mapValue) -> newValue, inputMap)
  *   ->
- * %map_from_arrays_unique%(
- *   map_keys(inputMap),
+ * %map_from_entries_unique%(
  *   array_map(
- *     (mapKey, mapValue) -> newValue,
- *     map_keys(inputMap), map_values(inputMap)))
+ *     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(MapLambdaValidator.requireLambda("transform_values", 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", - MapLambdaValidator.extractMapExpression("transform_values", lambda), - new MapEntryArrayMap(lambda)); + rewrittenLambda.getMapExpression(), rewrittenLambda.toArrayMap()); } private TransformValues(ScalarFunctionParams functionParams) { @@ -67,13 +73,16 @@ private TransformValues(ScalarFunctionParams functionParams) { @Override public FunctionSignature customSignature() { MapType inputMapType = (MapType) getArgument(0).getDataType(); - ArrayType transformedValuesType = (ArrayType) getArgument(1).getDataType(); - DataType resultValueType = MapLambdaValidator.mergeNestedNullTypes( - transformedValuesType.getItemType(), inputMapType.getValueType()); - transformedValuesType = ArrayType.of(resultValueType); + ArrayType mappedEntriesType = (ArrayType) getArgument(1).getDataType(); + StructType entryType = (StructType) mappedEntriesType.getItemType(); + List fields = entryType.getFields(); + DataType resultValueType = MapLambdaFunctionUtils.mergeNestedNullTypes( + fields.get(1).getDataType(), inputMapType.getValueType()); + StructType resolvedEntryType = new StructType(ImmutableList.of( + fields.get(0), fields.get(1).withDataType(resultValueType))); MapType resultType = MapType.of(inputMapType.getKeyType(), resultValueType); resultType.validateDataType(); - return FunctionSignature.ret(resultType).args(inputMapType, transformedValuesType); + return FunctionSignature.ret(resultType).args(inputMapType, ArrayType.of(resolvedEntryType)); } @Override @@ -84,7 +93,7 @@ public TransformValues withChildren(List children) { @Override public Expression rewriteWhenAnalyze() { - return new MapFromArraysUnique(new MapKeys(getArgument(0)), getArgument(1)); + return new MapFromEntriesUnique(getArgument(1)); } @Override diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/PlanUtils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/PlanUtils.java index 9dd1c8d5e9ff42..11773fd4c0385d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/PlanUtils.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/PlanUtils.java @@ -44,8 +44,6 @@ import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.expressions.WindowExpression; import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; -import org.apache.doris.nereids.trees.expressions.functions.scalar.MapEntryArrayMap; -import org.apache.doris.nereids.trees.expressions.functions.scalar.MapLambdaValidator; import org.apache.doris.nereids.trees.plans.Plan; import org.apache.doris.nereids.trees.plans.algebra.Filter; import org.apache.doris.nereids.trees.plans.algebra.Join; @@ -60,7 +58,6 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalProject; import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanVisitor; import org.apache.doris.nereids.types.DataType; -import org.apache.doris.nereids.types.MapType; import org.apache.doris.nereids.types.VarcharType; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.OriginStatement; @@ -197,25 +194,19 @@ public static List replaceExpressionByProjections(List childProjects, List targetExpressions) { - boolean containsMapLambda = targetExpressions.stream() - .anyMatch(target -> target.containsType(MapEntryArrayMap.class)); - Set nonRepeatableSlots = Sets.newHashSet(); + Set uniqueFunctionSlots = Sets.newHashSet(); for (Entry kv : ExpressionUtils.generateReplaceMap(childProjects).entrySet()) { - Expression value = kv.getValue(); - if (value.containsVolatileExpression() - || value.containsType(MapEntryArrayMap.class) - || (containsMapLambda && value.getDataType() instanceof MapType - && MapLambdaValidator.requiresSingleEvaluation(value))) { - nonRepeatableSlots.add(kv.getKey()); + if (kv.getValue().containsVolatileExpression()) { + uniqueFunctionSlots.add(kv.getKey()); } } - if (nonRepeatableSlots.isEmpty()) { + if (uniqueFunctionSlots.isEmpty()) { return true; } Set counterSet = Sets.newHashSet(); return targetExpressions.stream().noneMatch(target -> target.anyMatch( - e -> (e instanceof Slot) && nonRepeatableSlots.contains(e) && !counterSet.add((Slot) e))); + e -> (e instanceof Slot) && uniqueFunctionSlots.contains(e) && !counterSet.add((Slot) e))); } public static Plan skipProjectFilterLimit(Plan plan) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInputTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInputTest.java deleted file mode 100644 index 3d4ebe241f3a88..00000000000000 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInputTest.java +++ /dev/null @@ -1,427 +0,0 @@ -// 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.rules.rewrite; - -import org.apache.doris.nereids.exceptions.AnalysisException; -import org.apache.doris.nereids.hint.DistributeHint; -import org.apache.doris.nereids.trees.expressions.Add; -import org.apache.doris.nereids.trees.expressions.Alias; -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.SlotReference; -import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; -import org.apache.doris.nereids.trees.expressions.functions.scalar.Array; -import org.apache.doris.nereids.trees.expressions.functions.scalar.ArrayMap; -import org.apache.doris.nereids.trees.expressions.functions.scalar.CreateMap; -import org.apache.doris.nereids.trees.expressions.functions.scalar.CreateStruct; -import org.apache.doris.nereids.trees.expressions.functions.scalar.Lambda; -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.MapContainsKey; -import org.apache.doris.nereids.trees.expressions.functions.scalar.MapEntryArrayMap; -import org.apache.doris.nereids.trees.expressions.functions.scalar.MapFromEntries; -import org.apache.doris.nereids.trees.expressions.functions.scalar.MapKeys; -import org.apache.doris.nereids.trees.expressions.functions.scalar.MapValues; -import org.apache.doris.nereids.trees.expressions.functions.scalar.Random; -import org.apache.doris.nereids.trees.expressions.functions.scalar.StrToMap; -import org.apache.doris.nereids.trees.expressions.functions.scalar.TransformValues; -import org.apache.doris.nereids.trees.expressions.literal.BooleanLiteral; -import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; -import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; -import org.apache.doris.nereids.trees.plans.DistributeType; -import org.apache.doris.nereids.trees.plans.JoinType; -import org.apache.doris.nereids.trees.plans.Plan; -import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; -import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; -import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; -import org.apache.doris.nereids.trees.plans.logical.LogicalProject; -import org.apache.doris.nereids.util.MemoPatternMatchSupported; -import org.apache.doris.nereids.util.MemoTestUtils; -import org.apache.doris.nereids.util.PlanChecker; -import org.apache.doris.nereids.util.PlanConstructor; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableSet; -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Test; - -import java.util.List; -import java.util.Optional; - -public class AddProjectForMapLambdaInputTest implements MemoPatternMatchSupported { - private final LogicalOlapScan studentScan - = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(), PlanConstructor.student); - - @Test - void testMaterializeNondeterministicMapAsOneExpression() { - Random random = new Random(); - CreateMap map = new CreateMap(random, new IntegerLiteral(10)); - LogicalProject input = project(transformValues(map, new IntegerLiteral(0)), studentScan); - - LogicalProject rewritten = (LogicalProject) PlanChecker.from( - MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .applyTopDown(new AddProjectForVolatileExpression()) - .applyTopDown(new MergeProjectable()) - .getPlan(); - - LogicalProject mapProject = (LogicalProject) rewritten.child(); - Alias mapAlias = lastAlias(mapProject); - Assertions.assertEquals(map, mapAlias.child()); - Assertions.assertEquals(studentScan, mapProject.child()); - assertTransformValuesUsesMapSlot(rewritten.getProjects().get(0).child(0), mapAlias.toSlot()); - } - - @Test - void testMaterializeMapApplyInput() { - SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); - CreateMap map = new CreateMap(studentId, new Random()); - MapFromEntries loweredMapApply = lowerMapApply(map); - LogicalProject input = project(loweredMapApply, studentScan); - - LogicalProject rewritten = (LogicalProject) PlanChecker.from( - MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .applyTopDown(new MergeProjectable()) - .getPlan(); - - LogicalProject mapProject = (LogicalProject) rewritten.child(); - Alias mapAlias = lastAlias(mapProject); - Assertions.assertEquals(map, mapAlias.child()); - - MapFromEntries result = (MapFromEntries) rewritten.getProjects().get(0).child(0); - Lambda entryLambda = (Lambda) ((MapEntryArrayMap) result.child(0)).child(0); - assertLambdaUsesMapSlot(entryLambda, mapAlias.toSlot()); - } - - @Test - void testMaterializeDeterministicMapInFilter() { - SlotReference studentName = (SlotReference) studentScan.getOutput().get(2); - StrToMap map = new StrToMap(studentName); - LogicalFilter input = new LogicalFilter<>(ImmutableSet.of( - new MapContainsKey(transformValues(map, new StringLiteral("value")), - new StringLiteral("key"))), studentScan); - - LogicalFilter rewritten = (LogicalFilter) PlanChecker.from( - MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .getPlan(); - - LogicalProject mapProject = (LogicalProject) rewritten.child(); - Alias mapAlias = lastAlias(mapProject); - Assertions.assertEquals(map, mapAlias.child()); - MapContainsKey predicate = (MapContainsKey) rewritten.getConjuncts().iterator().next(); - assertTransformValuesUsesMapSlot(predicate.child(0), mapAlias.toSlot()); - } - - @Test - void testMaterializeDirectMapEntryLambdaInput() { - SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); - CreateMap map = new CreateMap(studentId, new IntegerLiteral(1)); - Lambda lambda = bindLambda("map_all", map, BooleanLiteral.TRUE); - LogicalProject input = project(new MapAll(lambda), studentScan); - - LogicalProject rewritten = (LogicalProject) PlanChecker.from( - MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .getPlan(); - - LogicalProject mapProject = (LogicalProject) rewritten.child(); - Alias mapAlias = lastAlias(mapProject); - MapAll mapAll = (MapAll) rewritten.getProjects().get(0).child(0); - Lambda rewrittenLambda = (Lambda) ((MapEntryArrayMap) mapAll.child(0)).child(0); - assertLambdaUsesMapSlot(rewrittenLambda, mapAlias.toSlot()); - } - - @Test - void testMapRuleDoesNotMaterializeUnrelatedLambdaBody() { - SlotReference studentName = (SlotReference) studentScan.getOutput().get(2); - StrToMap map = new StrToMap(studentName); - Random random = new Random(); - Add lambdaBody = new Add(random, random); - LogicalProject input = project(transformValues(map, lambdaBody), studentScan); - - LogicalProject rewritten = (LogicalProject) PlanChecker.from( - MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .getPlan(); - - LogicalProject mapProject = (LogicalProject) rewritten.child(); - Assertions.assertEquals(studentScan.getOutput().size() + 1, mapProject.getProjects().size()); - Assertions.assertEquals(map, lastAlias(mapProject).child()); - TransformValues transformValues = (TransformValues) rewritten.getProjects().get(0).child(0); - Lambda rewrittenLambda = (Lambda) ((MapEntryArrayMap) transformValues.child(1)).child(0); - Assertions.assertEquals(lambdaBody, rewrittenLambda.getLambdaFunction()); - } - - @Test - void testMaterializeJoinMapInputOnLeft() { - LogicalOlapScan scoreScan - = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(), PlanConstructor.score); - SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); - CreateMap map = new CreateMap(studentId, new IntegerLiteral(1)); - LogicalJoin input = joinWithMap(studentScan, scoreScan, map); - - LogicalJoin rewritten = (LogicalJoin) PlanChecker.from( - MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .getPlan(); - - LogicalProject leftProject = (LogicalProject) rewritten.left(); - Alias mapAlias = lastAlias(leftProject); - Assertions.assertEquals(map, mapAlias.child()); - Assertions.assertEquals(scoreScan, rewritten.right()); - assertJoinTransformValuesUsesMapSlot(rewritten, mapAlias.toSlot()); - } - - @Test - void testMaterializeJoinMapInputOnRight() { - LogicalOlapScan scoreScan - = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(), PlanConstructor.score); - SlotReference scoreId = (SlotReference) scoreScan.getOutput().get(0); - Random random = new Random(); - CreateMap map = new CreateMap(new Add(scoreId, random), new IntegerLiteral(1)); - LogicalJoin input = joinWithMap(studentScan, scoreScan, map); - - LogicalJoin rewritten = (LogicalJoin) PlanChecker.from( - MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .applyTopDown(new AddProjectForVolatileExpression()) - .getPlan(); - - Assertions.assertEquals(studentScan, rewritten.left()); - LogicalProject rightProject = (LogicalProject) rewritten.right(); - Alias mapAlias = lastAlias(rightProject); - Assertions.assertEquals(map, mapAlias.child()); - assertJoinTransformValuesUsesMapSlot(rewritten, mapAlias.toSlot()); - } - - @Test - void testSkipJoinMapInputDependingOnBothSides() { - LogicalOlapScan scoreScan - = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(), PlanConstructor.score); - SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); - SlotReference scoreId = (SlotReference) scoreScan.getOutput().get(0); - CreateMap map = new CreateMap(studentId, scoreId); - LogicalJoin input = joinWithMap(studentScan, scoreScan, map); - - Plan rewritten = PlanChecker.from(MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .getPlan(); - - Assertions.assertEquals(input, rewritten); - } - - @Test - void testRejectVolatileJoinMapInputDependingOnBothSides() { - LogicalOlapScan scoreScan - = new LogicalOlapScan(StatementScopeIdGenerator.newRelationId(), PlanConstructor.score); - SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); - SlotReference scoreId = (SlotReference) scoreScan.getOutput().get(0); - CreateMap map = new CreateMap(new Add(studentId, new Random()), scoreId); - LogicalJoin input = joinWithMap(studentScan, scoreScan, map); - - Assertions.assertThrows(AnalysisException.class, () -> PlanChecker.from( - MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .getPlan()); - } - - @Test - void testMaterializeMapInputInNestedLambda() { - SlotReference studentId = (SlotReference) studentScan.getOutput().get(0); - CreateMap outerMap = new CreateMap(studentId, new IntegerLiteral(10)); - Lambda outerTemplate = new Lambda(ImmutableList.of("ok", "ov"), new IntegerLiteral(0)); - List outerArguments = outerTemplate.makeArguments( - "transform_values", ImmutableList.of(outerMap)); - Slot outerKey = outerArguments.get(0).toSlot(); - Slot outerValue = outerArguments.get(1).toSlot(); - - CreateMap innerMap = new CreateMap(new Add(outerKey, new Random()), outerValue); - Lambda innerTemplate = new Lambda(ImmutableList.of("ik", "iv"), new IntegerLiteral(0)); - List innerArguments = innerTemplate.makeArguments( - "transform_values", ImmutableList.of(innerMap)); - TransformValues innerTransform = new TransformValues( - innerTemplate.withLambdaFunctionArguments(innerArguments.get(0).toSlot(), innerArguments)); - Lambda outerLambda = outerTemplate.withLambdaFunctionArguments(innerTransform, outerArguments); - LogicalProject input = project(new TransformValues(outerLambda), studentScan); - - LogicalProject rewritten = (LogicalProject) PlanChecker.from( - MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .getPlan(); - - TransformValues outerTransform = (TransformValues) rewritten.getProjects().get(0).child(0); - Lambda rewrittenOuterLambda = (Lambda) ((MapEntryArrayMap) outerTransform.child(1)).child(0); - Assertions.assertEquals(3, rewrittenOuterLambda.getLambdaArguments().size()); - ArrayItemReference hiddenArgument = rewrittenOuterLambda.getLambdaArgument(2); - Assertions.assertInstanceOf(ArrayMap.class, hiddenArgument.getArrayExpression()); - ArrayMap materializer = (ArrayMap) hiddenArgument.getArrayExpression(); - Lambda materializerLambda = (Lambda) materializer.child(0); - Assertions.assertTrue(materializerLambda.getLambdaFunction().containsType(Random.class)); - - TransformValues rewrittenInnerTransform - = (TransformValues) rewrittenOuterLambda.getLambdaFunction(); - Assertions.assertEquals(hiddenArgument.toSlot(), rewrittenInnerTransform.child(0)); - Lambda rewrittenInnerLambda - = (Lambda) ((MapEntryArrayMap) rewrittenInnerTransform.child(1)).child(0); - assertLambdaUsesMapSlot(rewrittenInnerLambda, hiddenArgument.toSlot()); - - LogicalProject rewrittenAgain = (LogicalProject) PlanChecker.from( - MemoTestUtils.createConnectContext(), rewritten) - .applyTopDown(new AddProjectForMapLambdaInput()) - .getPlan(); - TransformValues outerTransformAgain - = (TransformValues) rewrittenAgain.getProjects().get(0).child(0); - Lambda outerLambdaAgain - = (Lambda) ((MapEntryArrayMap) outerTransformAgain.child(1)).child(0); - Assertions.assertEquals(3, outerLambdaAgain.getLambdaArguments().size()); - } - - @Test - void testMaterializeMapInputInRegularArrayMapLambda() { - Array inputArray = new Array(new IntegerLiteral(1), new IntegerLiteral(2)); - Lambda outerTemplate = new Lambda(ImmutableList.of("x"), new IntegerLiteral(0)); - List outerArguments = outerTemplate.makeArguments( - "array_map", ImmutableList.of(inputArray)); - Slot outerItem = outerArguments.get(0).toSlot(); - - CreateMap innerMap = new CreateMap(new Random(), outerItem); - Lambda innerTemplate = new Lambda(ImmutableList.of("k", "v"), new IntegerLiteral(0)); - List innerArguments = innerTemplate.makeArguments( - "transform_values", ImmutableList.of(innerMap)); - TransformValues innerTransform = new TransformValues( - innerTemplate.withLambdaFunctionArguments(innerArguments.get(0).toSlot(), innerArguments)); - Lambda outerLambda = outerTemplate.withLambdaFunctionArguments(innerTransform, outerArguments); - LogicalProject input = project(new ArrayMap(outerLambda), studentScan); - - LogicalProject rewritten = (LogicalProject) PlanChecker.from( - MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .applyTopDown(new AddProjectForVolatileExpression()) - .getPlan(); - - Assertions.assertEquals(studentScan, rewritten.child()); - ArrayMap rewrittenArrayMap = (ArrayMap) rewritten.getProjects().get(0).child(0); - Lambda rewrittenOuterLambda = (Lambda) rewrittenArrayMap.child(0); - Assertions.assertEquals(2, rewrittenOuterLambda.getLambdaArguments().size()); - ArrayItemReference hiddenArgument = rewrittenOuterLambda.getLambdaArgument(1); - Assertions.assertInstanceOf(ArrayMap.class, hiddenArgument.getArrayExpression()); - ArrayMap materializer = (ArrayMap) hiddenArgument.getArrayExpression(); - Lambda materializerLambda = (Lambda) materializer.child(0); - Assertions.assertTrue(materializerLambda.getLambdaFunction().containsType(Random.class)); - - TransformValues rewrittenInnerTransform - = (TransformValues) rewrittenOuterLambda.getLambdaFunction(); - Assertions.assertEquals(hiddenArgument.toSlot(), rewrittenInnerTransform.child(0)); - Lambda rewrittenInnerLambda - = (Lambda) ((MapEntryArrayMap) rewrittenInnerTransform.child(1)).child(0); - assertLambdaUsesMapSlot(rewrittenInnerLambda, hiddenArgument.toSlot()); - } - - @Test - void testMaterializeNestedMapApplyInput() { - Array inputArray = new Array(new IntegerLiteral(1), new IntegerLiteral(2)); - Lambda outerTemplate = new Lambda(ImmutableList.of("x"), new IntegerLiteral(0)); - List outerArguments = outerTemplate.makeArguments( - "array_map", ImmutableList.of(inputArray)); - CreateMap innerMap = new CreateMap(new Random(), outerArguments.get(0).toSlot()); - MapFromEntries loweredMapApply = lowerMapApply(innerMap); - Lambda outerLambda = outerTemplate.withLambdaFunctionArguments( - loweredMapApply, outerArguments); - LogicalProject input = project(new ArrayMap(outerLambda), studentScan); - - LogicalProject rewritten = (LogicalProject) PlanChecker.from( - MemoTestUtils.createConnectContext(), input) - .applyTopDown(new AddProjectForMapLambdaInput()) - .getPlan(); - - ArrayMap rewrittenArrayMap = (ArrayMap) rewritten.getProjects().get(0).child(0); - Lambda rewrittenOuterLambda = (Lambda) rewrittenArrayMap.child(0); - // Keep x plus the materialized Map input used by the nested Map Lambda. - Assertions.assertEquals(2, rewrittenOuterLambda.getLambdaArguments().size()); - ArrayItemReference mapArgument = rewrittenOuterLambda.getLambdaArgument(1); - Assertions.assertInstanceOf(ArrayMap.class, mapArgument.getArrayExpression()); - Assertions.assertTrue(mapArgument.getArrayExpression().containsType(Random.class)); - - MapFromEntries result = (MapFromEntries) rewrittenOuterLambda.getLambdaFunction(); - Lambda entryLambda = (Lambda) ((MapEntryArrayMap) result.child(0)).child(0); - assertLambdaUsesMapSlot(entryLambda, mapArgument.toSlot()); - } - - private LogicalProject project(Expression expression, Plan child) { - return new LogicalProject(ImmutableList.of(new Alias(expression)), child); - } - - private TransformValues transformValues(Expression map, Expression body) { - return new TransformValues(bindLambda("transform_values", map, body)); - } - - private MapFromEntries lowerMapApply(Expression map) { - Lambda lambda = new Lambda(ImmutableList.of("k", "v"), new IntegerLiteral(0)); - List arguments = lambda.makeArguments("map_apply", ImmutableList.of(map)); - CreateStruct body = new CreateStruct(arguments.get(0).toSlot(), arguments.get(1).toSlot()); - MapApply mapApply = new MapApply(lambda.withLambdaFunctionArguments(body, arguments)); - return (MapFromEntries) mapApply.rewriteWhenAnalyze(); - } - - private Lambda bindLambda(String functionName, Expression map, Expression body) { - Lambda lambda = new Lambda(ImmutableList.of("k", "v"), body); - List arguments = lambda.makeArguments(functionName, ImmutableList.of(map)); - return lambda.withLambdaFunctionArguments(body, arguments); - } - - private LogicalJoin joinWithMap(Plan left, Plan right, Expression map) { - MapContainsKey predicate = new MapContainsKey( - transformValues(map, new IntegerLiteral(0)), new IntegerLiteral(1)); - return new LogicalJoin( - JoinType.INNER_JOIN, - ImmutableList.of(), - ImmutableList.of(predicate), - new DistributeHint(DistributeType.NONE), - Optional.empty(), - left, - right, - null); - } - - private Alias lastAlias(LogicalProject project) { - return (Alias) project.getProjects().get(project.getProjects().size() - 1); - } - - private void assertJoinTransformValuesUsesMapSlot(LogicalJoin join, Slot mapSlot) { - MapContainsKey predicate = (MapContainsKey) join.getOtherJoinConjuncts().get(0); - assertTransformValuesUsesMapSlot(predicate.child(0), mapSlot); - } - - private void assertTransformValuesUsesMapSlot(Expression expression, Slot mapSlot) { - TransformValues transformValues = (TransformValues) expression; - Assertions.assertEquals(mapSlot, transformValues.child(0)); - Lambda lambda = (Lambda) ((MapEntryArrayMap) transformValues.child(1)).child(0); - assertLambdaUsesMapSlot(lambda, mapSlot); - } - - private void assertLambdaUsesMapSlot(Lambda lambda, Slot mapSlot) { - Assertions.assertEquals(mapSlot, - ((MapKeys) lambda.getLambdaArgument(0).getArrayExpression()).child(0)); - Assertions.assertEquals(mapSlot, - ((MapValues) lambda.getLambdaArgument(1).getArrayExpression()).child(0)); - } - -} 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 index 69c7d1ddc0b6a7..08a56a4c01fc41 100644 --- 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 @@ -30,6 +30,7 @@ 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; @@ -88,8 +89,8 @@ public void testMapLambdaWrappersAndTypes() { tupleMapApply.getDataType()); Expression mapFilter = analyze("map_filter((k, v) -> v > 10, map(1, 10, 2, 20))"); - Assertions.assertTrue(mapFilter instanceof MapFilter); - assertMapEntryArray(((MapFilter) mapFilter).child(1)); + 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]))"); @@ -97,17 +98,15 @@ public void testMapLambdaWrappersAndTypes() { 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 MapFromArrays); + Assertions.assertTrue(transformKeys instanceof MapFromEntries); assertMapEntryArray(transformKeys.child(0)); - Assertions.assertTrue(transformKeys.child(1) instanceof MapValues); 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 MapFromArraysUnique); - Assertions.assertTrue(transformValues.child(0) instanceof MapKeys); - assertMapEntryArray(transformValues.child(1)); + Assertions.assertTrue(transformValues instanceof MapFromEntriesUnique); + assertMapEntryArray(transformValues.child(0)); Assertions.assertEquals(MapType.of(TinyIntType.INSTANCE, SmallIntType.INSTANCE), transformValues.getDataType()); } @@ -130,20 +129,19 @@ public void testPureNullLambdaReturnUsesInputMapType() { Expression transformValues = analyze( "transform_values((k, v) -> null, map(1, 10))"); Assertions.assertEquals(inputMapType, transformValues.getDataType()); - Assertions.assertEquals(ArrayType.of(TinyIntType.INSTANCE), - transformValues.child(1).getDataType()); + Assertions.assertEquals(TinyIntType.INSTANCE, + mappedEntryType(transformValues).getFields().get(1).getDataType()); Expression transformKeys = analyze( "transform_keys((k, v) -> null, map(1, 10))"); Assertions.assertEquals(inputMapType, transformKeys.getDataType()); - Assertions.assertEquals(ArrayType.of(TinyIntType.INSTANCE), - transformKeys.child(1).getDataType()); + Assertions.assertEquals(TinyIntType.INSTANCE, + mappedEntryType(transformKeys).getFields().get(0).getDataType()); Expression mapApply = analyze( "map_apply((k, v) -> struct(k, null), map(1, 10))"); Assertions.assertEquals(inputMapType, mapApply.getDataType()); - StructType mappedEntryType = - (StructType) ((ArrayType) mapApply.child(0).getDataType()).getItemType(); + StructType mappedEntryType = mappedEntryType(mapApply); Assertions.assertEquals(TinyIntType.INSTANCE, mappedEntryType.getFields().get(0).getDataType()); Assertions.assertEquals(TinyIntType.INSTANCE, @@ -162,8 +160,8 @@ public void testNestedNullLambdaReturnUsesInputMapType() { Expression transformArrayValues = analyze( "transform_values((k, v) -> [], map(1, [10]))"); Assertions.assertEquals(arrayValueMapType, transformArrayValues.getDataType()); - Assertions.assertEquals(ArrayType.of(tinyIntArrayType), - transformArrayValues.child(1).getDataType()); + Assertions.assertEquals(tinyIntArrayType, + mappedEntryType(transformArrayValues).getFields().get(1).getDataType()); Expression mapApplyArrayValue = analyze( "map_apply((k, v) -> struct(k, []), map(1, [10]))"); @@ -203,21 +201,34 @@ public void testNestedLambdaCanCaptureImmediateOuterScope() { 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); - Lambda lambda = new Lambda(ImmutableList.of("k", "v"), value); - List arguments = - lambda.makeArguments("transform_values", ImmutableList.of(computedMap)); - Lambda boundLambda = lambda.withLambdaFunctionArguments(arguments.get(1).toSlot(), arguments); + 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 MapFromArraysUnique); + Assertions.assertTrue(nondeterministicMap instanceof MapFromEntriesUnique); + assertMapEntryArray(nondeterministicMap.child(0)); } @Test @@ -269,11 +280,10 @@ public void testMapFromArraysCanBeAnalyzed() { Assertions.assertTrue(complexKeyException.getMessage().contains( "MAP key type must be a primitive type"), complexKeyException::getMessage); - AnalysisException invalidNestedValueException = Assertions.assertThrows(AnalysisException.class, - () -> analyze("map_from_arrays([1], [[]])")); - Assertions.assertTrue(invalidNestedValueException.getMessage().contains( - "Unsupported data type: map>"), - invalidNestedValueException::getMessage); + Expression nestedNullValueMap = analyze("map_from_arrays([1], [[]])"); + Assertions.assertEquals( + MapType.of(TinyIntType.INSTANCE, ArrayType.of(TinyIntType.INSTANCE)), + nestedNullValueMap.getDataType()); } @Test @@ -324,21 +334,25 @@ private void assertMapEntryArray(Expression expression) { while (expression instanceof Cast) { expression = expression.child(0); } - Assertions.assertTrue(expression instanceof MapEntryArrayMap); - MapEntryArrayMap marker = (MapEntryArrayMap) expression; - Assertions.assertTrue(marker.child(0) instanceof Lambda); + Assertions.assertTrue(expression instanceof ArrayMap); + ArrayMap arrayMap = (ArrayMap) expression; + Assertions.assertTrue(arrayMap.child(0) instanceof Lambda); - Lambda lambda = (Lambda) marker.child(0); + Lambda lambda = (Lambda) arrayMap.child(0); List arguments = lambda.getLambdaArguments(); - Assertions.assertEquals(2, arguments.size()); - Assertions.assertTrue(arguments.get(0).getArrayExpression() instanceof MapKeys); - Assertions.assertTrue(arguments.get(1).getArrayExpression() instanceof MapValues); - Assertions.assertEquals( - arguments.get(0).getArrayExpression().child(0), - arguments.get(1).getArrayExpression().child(0)); - Assertions.assertNotEquals(arguments.get(0).getExprId(), arguments.get(1).getExprId()); + 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); + } - Assertions.assertTrue(marker.withChildren(marker.children()) instanceof MapEntryArrayMap); + private StructType mappedEntryType(Expression mapExpression) { + return (StructType) ((ArrayType) mapExpression.child(0).getDataType()).getItemType(); } } 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 index 91e86065a1e004..99699c9d4a6f85 100644 --- 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 @@ -139,3 +139,6 @@ x:a y:b -- !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 index 18e41f4c754cc7..69b5ab494ff5e0 100644 --- 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 @@ -422,10 +422,7 @@ suite("test_map_lambda", "p0") { sql "select map_from_arrays([[1]], [10])" exception "MAP key type must be a primitive type" } - test { - sql "select map_from_arrays([1], [[]])" - exception "Unsupported data type: map>" - } + 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" From 9ee7c92e7348857927c5a9c92c789f0c33e9c1ec Mon Sep 17 00:00:00 2001 From: linzhenqi Date: Thu, 27 Aug 2026 02:40:57 +0800 Subject: [PATCH 3/4] fix: remove function-specific NullType handling from map lambdas --- .../functions/scalar/MapApply.java | 28 ++----- .../scalar/MapLambdaFunctionUtils.java | 35 --------- .../functions/scalar/TransformKeys.java | 10 +-- .../functions/scalar/TransformValues.java | 8 +- .../scalar/MapLambdaFunctionsTest.java | 77 ------------------- .../map_functions/test_map_lambda.out | 22 ------ .../map_functions/test_map_lambda.groovy | 46 ----------- 7 files changed, 9 insertions(+), 217 deletions(-) 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 index 46cf7f7b775f20..f0afeef273b209 100644 --- 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 @@ -32,7 +32,6 @@ import org.apache.doris.nereids.types.StructType; import com.google.common.base.Preconditions; -import com.google.common.collect.ImmutableList; import java.util.List; @@ -82,11 +81,10 @@ public FunctionSignature customSignature() { throw invalidReturnType(); } MapType inputMapType = (MapType) getArgument(0).getDataType(); - StructType resolvedStructType = resolveNullFieldTypes(structType, inputMapType); - List fields = resolvedStructType.getFields(); + List fields = structType.getFields(); MapType resultType = MapType.of(fields.get(0).getDataType(), fields.get(1).getDataType()); resultType.validateDataType(); - return FunctionSignature.ret(resultType).args(inputMapType, ArrayType.of(resolvedStructType)); + return FunctionSignature.ret(resultType).args(inputMapType, mappedEntriesType); } @Override @@ -108,11 +106,11 @@ public R accept(ExpressionVisitor visitor, C context) { private static MapLambdaFunctionUtils.RewrittenMapLambda validateAndRewrite(Lambda lambda) { MapLambdaFunctionUtils.RewrittenMapLambda rewrittenLambda = MapLambdaFunctionUtils.rewrite( lambda, (body, key, value, entry) -> body); - validateLambdaReturn(lambda, (MapType) rewrittenLambda.getMapExpression().getDataType()); + validateLambdaReturn(lambda); return rewrittenLambda; } - private static void validateLambdaReturn(Lambda lambda, MapType inputMapType) { + private static void validateLambdaReturn(Lambda lambda) { Expression lambdaBody = lambda.getLambdaFunction(); if (!(lambdaBody.getDataType() instanceof StructType) || ((StructType) lambdaBody.getDataType()).getFields().size() != 2 @@ -120,24 +118,8 @@ private static void validateLambdaReturn(Lambda lambda, MapType inputMapType) { throw invalidReturnType(); } StructType structType = (StructType) lambdaBody.getDataType(); - StructType resolvedStructType = resolveNullFieldTypes(structType, inputMapType); - MapType.of(resolvedStructType.getFields().get(0).getDataType(), - resolvedStructType.getFields().get(1).getDataType()).validateDataType(); - } - - // Resolve only untyped fields in the two-field struct returned by the lambda. For - // map_apply((k, v) -> struct(cast(k as bigint), []), map(1, [10])), the result is - // MAP>. Keep the explicit BIGINT type and infer only the empty array - // from the input value type. - private static StructType resolveNullFieldTypes(StructType structType, MapType inputMapType) { List fields = structType.getFields(); - StructField keyField = fields.get(0); - StructField valueField = fields.get(1); - keyField = keyField.withDataType(MapLambdaFunctionUtils.mergeNestedNullTypes( - keyField.getDataType(), inputMapType.getKeyType())); - valueField = valueField.withDataType(MapLambdaFunctionUtils.mergeNestedNullTypes( - valueField.getDataType(), inputMapType.getValueType())); - return new StructType(ImmutableList.of(keyField, valueField)); + MapType.of(fields.get(0).getDataType(), fields.get(1).getDataType()).validateDataType(); } private static AnalysisException invalidReturnType() { 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 index 55da184752161c..9bd8224c903f0e 100644 --- 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 @@ -22,11 +22,6 @@ 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 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 com.google.common.collect.ImmutableList; @@ -63,36 +58,6 @@ static Lambda requireLambda(String functionName, Expression expression) { return (Lambda) expression; } - /** Fill only NULL_TYPE positions from the corresponding input Map field type. */ - static DataType mergeNestedNullTypes(DataType outputType, DataType inputType) { - if (outputType.isNullType()) { - return inputType; - } else if (outputType instanceof ArrayType && inputType instanceof ArrayType) { - return ArrayType.of(mergeNestedNullTypes( - ((ArrayType) outputType).getItemType(), ((ArrayType) inputType).getItemType())); - } else if (outputType instanceof MapType && inputType instanceof MapType) { - return MapType.of( - mergeNestedNullTypes( - ((MapType) outputType).getKeyType(), ((MapType) inputType).getKeyType()), - mergeNestedNullTypes( - ((MapType) outputType).getValueType(), ((MapType) inputType).getValueType())); - } else if (outputType instanceof StructType && inputType instanceof StructType) { - List outputFields = ((StructType) outputType).getFields(); - List inputFields = ((StructType) inputType).getFields(); - if (outputFields.size() != inputFields.size()) { - return outputType; - } - ImmutableList.Builder fields - = ImmutableList.builderWithExpectedSize(outputFields.size()); - for (int i = 0; i < outputFields.size(); i++) { - fields.add(outputFields.get(i).withDataType(mergeNestedNullTypes( - outputFields.get(i).getDataType(), inputFields.get(i).getDataType()))); - } - return new StructType(fields.build()); - } - return outputType; - } - private static Expression extractMapExpression(Lambda lambda) { List arguments = lambda.getLambdaArguments(); Preconditions.checkArgument(arguments.size() == 1, 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 index 9637b27b78d50e..d5b832075bc743 100644 --- 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 @@ -31,7 +31,6 @@ import org.apache.doris.nereids.types.StructType; import com.google.common.base.Preconditions; -import com.google.common.collect.ImmutableList; import java.util.List; @@ -63,15 +62,10 @@ public FunctionSignature customSignature() { ArrayType mappedEntriesType = (ArrayType) getArgument(1).getDataType(); StructType entryType = (StructType) mappedEntriesType.getItemType(); List fields = entryType.getFields(); - // transform_keys((k, v) -> null, map(1, 10)) - // res_type should be: MAP instead of MAP - DataType resultKeyType = MapLambdaFunctionUtils.mergeNestedNullTypes( - fields.get(0).getDataType(), inputMapType.getKeyType()); - StructType resolvedEntryType = new StructType(ImmutableList.of( - fields.get(0).withDataType(resultKeyType), fields.get(1))); + DataType resultKeyType = fields.get(0).getDataType(); MapType resultType = MapType.of(resultKeyType, inputMapType.getValueType()); resultType.validateDataType(); - return FunctionSignature.ret(resultType).args(inputMapType, ArrayType.of(resolvedEntryType)); + return FunctionSignature.ret(resultType).args(inputMapType, mappedEntriesType); } @Override 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 index e27c6e4650379a..47f549d09177b5 100644 --- 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 @@ -31,7 +31,6 @@ import org.apache.doris.nereids.types.StructType; import com.google.common.base.Preconditions; -import com.google.common.collect.ImmutableList; import java.util.List; @@ -76,13 +75,10 @@ public FunctionSignature customSignature() { ArrayType mappedEntriesType = (ArrayType) getArgument(1).getDataType(); StructType entryType = (StructType) mappedEntriesType.getItemType(); List fields = entryType.getFields(); - DataType resultValueType = MapLambdaFunctionUtils.mergeNestedNullTypes( - fields.get(1).getDataType(), inputMapType.getValueType()); - StructType resolvedEntryType = new StructType(ImmutableList.of( - fields.get(0), fields.get(1).withDataType(resultValueType))); + DataType resultValueType = fields.get(1).getDataType(); MapType resultType = MapType.of(inputMapType.getKeyType(), resultValueType); resultType.validateDataType(); - return FunctionSignature.ret(resultType).args(inputMapType, ArrayType.of(resolvedEntryType)); + return FunctionSignature.ret(resultType).args(inputMapType, mappedEntriesType); } @Override 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 index 08a56a4c01fc41..b86920ebe1dfe2 100644 --- 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 @@ -37,7 +37,6 @@ 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.StructType; import org.apache.doris.nereids.types.TinyIntType; import org.apache.doris.nereids.util.MemoTestUtils; import org.apache.doris.utframe.TestWithFeService; @@ -122,78 +121,6 @@ public void testMapExistsAndAllRewriteToArrayMatch() { assertMapEntryArray(mapAll.child(0)); } - @Test - public void testPureNullLambdaReturnUsesInputMapType() { - MapType inputMapType = MapType.of(TinyIntType.INSTANCE, TinyIntType.INSTANCE); - - Expression transformValues = analyze( - "transform_values((k, v) -> null, map(1, 10))"); - Assertions.assertEquals(inputMapType, transformValues.getDataType()); - Assertions.assertEquals(TinyIntType.INSTANCE, - mappedEntryType(transformValues).getFields().get(1).getDataType()); - - Expression transformKeys = analyze( - "transform_keys((k, v) -> null, map(1, 10))"); - Assertions.assertEquals(inputMapType, transformKeys.getDataType()); - Assertions.assertEquals(TinyIntType.INSTANCE, - mappedEntryType(transformKeys).getFields().get(0).getDataType()); - - Expression mapApply = analyze( - "map_apply((k, v) -> struct(k, null), map(1, 10))"); - Assertions.assertEquals(inputMapType, mapApply.getDataType()); - StructType mappedEntryType = mappedEntryType(mapApply); - Assertions.assertEquals(TinyIntType.INSTANCE, - mappedEntryType.getFields().get(0).getDataType()); - Assertions.assertEquals(TinyIntType.INSTANCE, - mappedEntryType.getFields().get(1).getDataType()); - - Expression mapApplyNullKey = analyze( - "map_apply((k, v) -> struct(null, v), map(1, 10))"); - Assertions.assertEquals(inputMapType, mapApplyNullKey.getDataType()); - } - - @Test - public void testNestedNullLambdaReturnUsesInputMapType() { - ArrayType tinyIntArrayType = ArrayType.of(TinyIntType.INSTANCE); - MapType arrayValueMapType = MapType.of(TinyIntType.INSTANCE, tinyIntArrayType); - - Expression transformArrayValues = analyze( - "transform_values((k, v) -> [], map(1, [10]))"); - Assertions.assertEquals(arrayValueMapType, transformArrayValues.getDataType()); - Assertions.assertEquals(tinyIntArrayType, - mappedEntryType(transformArrayValues).getFields().get(1).getDataType()); - - Expression mapApplyArrayValue = analyze( - "map_apply((k, v) -> struct(k, []), map(1, [10]))"); - Assertions.assertEquals(arrayValueMapType, mapApplyArrayValue.getDataType()); - - MapType nestedMapType = MapType.of( - TinyIntType.INSTANCE, MapType.of(TinyIntType.INSTANCE, TinyIntType.INSTANCE)); - Expression transformMapValues = analyze( - "transform_values((k, v) -> map(), map(1, map(2, 20)))"); - Assertions.assertEquals(nestedMapType, transformMapValues.getDataType()); - - Expression transformStructValues = analyze( - "transform_values((k, v) -> struct(null, []), " - + "map(1, struct(10, [20])))"); - StructType transformStructType = (StructType) - ((MapType) transformStructValues.getDataType()).getValueType(); - Assertions.assertEquals(TinyIntType.INSTANCE, - transformStructType.getFields().get(0).getDataType()); - Assertions.assertEquals(tinyIntArrayType, - transformStructType.getFields().get(1).getDataType()); - - Expression mapApplyStructValue = analyze( - "map_apply((k, v) -> struct(k, struct(null, [])), " - + "map(1, struct(10, [20])))"); - StructType mapApplyStructType = (StructType) - ((MapType) mapApplyStructValue.getDataType()).getValueType(); - Assertions.assertEquals(TinyIntType.INSTANCE, - mapApplyStructType.getFields().get(0).getDataType()); - Assertions.assertEquals(tinyIntArrayType, - mapApplyStructType.getFields().get(1).getDataType()); - } - @Test public void testNestedLambdaCanCaptureImmediateOuterScope() { Expression nested = analyze("map_exists((x, v) -> " @@ -351,8 +278,4 @@ private void assertMapEntryArray(Expression expression) { Assertions.assertTrue(arrayMap.withChildren(arrayMap.children()) instanceof ArrayMap); } - private StructType mappedEntryType(Expression mapExpression) { - return (StructType) ((ArrayType) mapExpression.child(0).getDataType()).getItemType(); - } - } 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 index 99699c9d4a6f85..8bdd1713035e4d 100644 --- 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 @@ -47,27 +47,6 @@ -- !transform_values -- 11 21 --- !transform_values_null_type -- -2 \N \N - --- !transform_keys_null_type -- -1 20 - --- !map_apply_null_value_type -- -2 \N \N - --- !transform_values_nested_array_null_type -- -[] - --- !map_apply_nested_array_null_type -- -[] - --- !transform_values_nested_map_null_type -- -0 - --- !transform_values_nested_struct_null_type -- -{"col1":null, "col2":[]} - -- !transform_values_string -- x:a y:b @@ -141,4 +120,3 @@ x:a y:b -- !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 index 69b5ab494ff5e0..9d0266da3dfb91 100644 --- 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 @@ -156,52 +156,6 @@ suite("test_map_lambda", "p0") { from test_map_lambda where id = 1 ) t """ - qt_transform_values_null_type """ - select map_size(r), r[1], r[2] - from ( - select transform_values((k, v) -> null, mii) r - from test_map_lambda where id = 1 - ) t - """ - qt_transform_keys_null_type """ - select map_size(r), map_values(r)[1] - from ( - select transform_keys((k, v) -> null, mii) r - from test_map_lambda where id = 1 - ) t - """ - qt_map_apply_null_value_type """ - select map_size(r), r[1], r[2] - from ( - select map_apply((k, v) -> struct(k, null), mii) r - from test_map_lambda where id = 1 - ) t - """ - qt_transform_values_nested_array_null_type """ - select r[1] - from ( - select transform_values((k, v) -> [], map(1, [10])) r - ) t - """ - qt_map_apply_nested_array_null_type """ - select r[1] - from ( - select map_apply((k, v) -> struct(k, []), map(1, [10])) r - ) t - """ - qt_transform_values_nested_map_null_type """ - select map_size(r[1]) - from ( - select transform_values((k, v) -> map(), map(1, map(2, 20))) r - ) t - """ - qt_transform_values_nested_struct_null_type """ - select r[1] - from ( - select transform_values( - (k, v) -> struct(null, []), map(1, struct(10, [20]))) r - ) t - """ qt_transform_values_string """ select r['a'], r['b'] from ( From a6e16d36007e68b5d891d8924ca90d92695f829b Mon Sep 17 00:00:00 2001 From: linzhenqi Date: Thu, 27 Aug 2026 03:05:15 +0800 Subject: [PATCH 4/4] [fix](be) Reject oversized map selector positions ### What problem does this PR solve? Issue Number: None Related PR: #66968 Problem Summary: Map offsets use 64-bit positions while IColumn::Selector stores 32-bit indexes. map_filter and the filtered-entry map constructor could therefore truncate a position at 2^32 and silently copy the wrong key and value. Validate the nested entry count before building a selector and return INVALID_ARGUMENT when an index cannot be represented. Add an O(1)-memory boundary unit test for the overflow case. ### Release note Reject map operations whose nested positions exceed the selector range instead of returning corrupted data. ### Check List (For Author) - Test: Unit Test (FunctionMapTest.* under ASAN) - Behavior changed: Yes (oversized selector inputs now return INVALID_ARGUMENT) - Does this need documentation: No --- be/src/exprs/function/function_map.cpp | 48 ++++++++++++++------ be/test/exprs/function/function_map_test.cpp | 34 ++++++++++++++ 2 files changed, 67 insertions(+), 15 deletions(-) diff --git a/be/src/exprs/function/function_map.cpp b/be/src/exprs/function/function_map.cpp index db4abbb274f3e3..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 @@ -223,8 +224,9 @@ class FunctionMapFilter : public IFunction { IColumn::Selector selector; auto result_offsets = ColumnArray::ColumnOffsets::create(); - build_selector_and_offsets(map, map_is_const, predicate, predicate_is_const, - result_null_map_data, selector, *result_offsets); + 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(); @@ -249,11 +251,18 @@ class FunctionMapFilter : public IFunction { } private: - static void 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) { + 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 = @@ -282,6 +291,7 @@ class FunctionMapFilter : public IFunction { } result_offsets.insert_value(selector.size()); } + return Status::OK(); } static Status check_arguments(const ColumnMap& map, bool map_is_const, @@ -697,8 +707,8 @@ class FunctionMapFromFilteredEntries : public IFunction { if (nullable_entries.has_null()) { IColumn::Selector selector; auto filtered_offsets = ColumnArray::ColumnOffsets::create(); - build_selector_and_offsets(entries, nullable_entries, nullable_array, selector, - *filtered_offsets); + 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(); @@ -723,11 +733,18 @@ class FunctionMapFromFilteredEntries : public IFunction { } private: - static void build_selector_and_offsets(const ColumnArray& entries, - const ColumnNullable& nullable_entries, - const ColumnNullable* nullable_array, - IColumn::Selector& selector, - ColumnArray::ColumnOffsets& filtered_offsets) { + 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()); @@ -737,12 +754,13 @@ class FunctionMapFromFilteredEntries : public IFunction { 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(static_cast(entry)); + selector.push_back(entry); } } } filtered_offsets.insert_value(selector.size()); } + return Status::OK(); } }; diff --git a/be/test/exprs/function/function_map_test.cpp b/be/test/exprs/function/function_map_test.cpp index f51496c1277b00..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" @@ -557,6 +559,38 @@ TEST(FunctionMapTest, map_from_entries_unique_skips_deduplication) { 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(