Fix lambda type resolution

src/Analyzer/LambdaNode.h

@@ -86,7 +86,12 @@ class LambdaNode final : public IQueryTreeNode
 
     DataTypePtr getResultType() const override
     {
-        return getExpression()->getResultType();
+        return result_type;
+    }
+
+    void resolve(DataTypePtr lambda_type)
+    {
+        result_type = std::move(lambda_type);
     }
 
     void dumpTreeImpl(WriteBuffer & buffer, FormatState & format_state, size_t indent) const override;
@@ -102,6 +107,7 @@ class LambdaNode final : public IQueryTreeNode
 
 private:
     Names argument_names;
+    DataTypePtr result_type;
 
     static constexpr size_t arguments_child_index = 0;
     static constexpr size_t expression_child_index = 1;

src/Analyzer/Passes/QueryAnalysisPass.cpp

@@ -5085,8 +5085,11 @@ ProjectionNames QueryAnalyzer::resolveFunction(QueryTreeNodePtr & node, Identifi
                 arguments_projection_names[function_lambda_argument_index] = lambda_argument_projection_name_buffer.str();
             }
 
-            argument_types[function_lambda_argument_index] = std::make_shared<DataTypeFunction>(function_data_type_argument_types, lambda_to_resolve->getResultType());
-            argument_columns[function_lambda_argument_index].type = argument_types[function_lambda_argument_index];
+            auto lambda_resolved_type = std::make_shared<DataTypeFunction>(function_data_type_argument_types, lambda_to_resolve_typed.getExpression()->getResultType());
+            lambda_to_resolve_typed.resolve(lambda_resolved_type);
+
+            argument_types[function_lambda_argument_index] = lambda_resolved_type;
+            argument_columns[function_lambda_argument_index].type = lambda_resolved_type;
             function_arguments[function_lambda_argument_index] = std::move(lambda_to_resolve);
         }
 

src/Analyzer/QueryTreePassManager.cpp

@@ -115,13 +115,23 @@ class ValidationChecker : public InDepthQueryTreeVisitor<ValidationChecker>
 
         for (size_t i = 0; i < expected_argument_types_size; ++i)
         {
-            // Skip lambdas
-            if (WhichDataType(expected_argument_types[i]).isFunction())
-                continue;
-
             const auto & expected_argument_type = expected_argument_types[i];
             const auto & actual_argument_type = actual_argument_columns[i].type;
 
+            if (!expected_argument_type)
+                throw Exception(ErrorCodes::LOGICAL_ERROR,
+                    "Function {} expected argument {} type is not set after running {} pass",
+                    function->toAST()->formatForErrorMessage(),
+                    i + 1,
+                    pass_name);
+
+            if (!actual_argument_type)
+                throw Exception(ErrorCodes::LOGICAL_ERROR,
+                    "Function {} actual argument {} type is not set after running {} pass",
+                    function->toAST()->formatForErrorMessage(),
+                    i + 1,
+                    pass_name);
+
             if (!expected_argument_type->equals(*actual_argument_type))
             {
                 /// Aggregate functions remove low cardinality for their argument types

src/Planner/PlannerActionsVisitor.cpp

@@ -599,11 +599,16 @@ PlannerActionsVisitorImpl::NodeNameAndNodeMinLevel PlannerActionsVisitorImpl::vi
 
     auto lambda_node_name = calculateActionNodeName(node, *planner_context);
     auto function_capture = std::make_shared<FunctionCaptureOverloadResolver>(
-        lambda_actions, captured_column_names, lambda_arguments_names_and_types, result_type, lambda_expression_node_name);
+        lambda_actions, captured_column_names, lambda_arguments_names_and_types, lambda_node.getExpression()->getResultType(), lambda_expression_node_name);
     actions_stack.pop_back();
 
     // TODO: Pass IFunctionBase here not FunctionCaptureOverloadResolver.
-    actions_stack[level].addFunctionIfNecessary(lambda_node_name, std::move(lambda_children), function_capture);
+    const auto * actions_node = actions_stack[level].addFunctionIfNecessary(lambda_node_name, std::move(lambda_children), function_capture);
+
+    if (!result_type->equals(*actions_node->result_type))
+        throw Exception(ErrorCodes::LOGICAL_ERROR,
+            "Lambda resolved type {} is not equal to type from actions DAG {}",
+            result_type, actions_node->result_type);
 
     size_t actions_stack_size = actions_stack.size();
     for (size_t i = level + 1; i < actions_stack_size; ++i)

tests/queries/0_stateless/02707_analyzer_nested_lambdas_types.reference

@@ -0,0 +1,5 @@
+[0]	[[0]]
+[0]	[[0]]
+[0]	[[1]]
+[0]	[[1]]
+[0]	[[1]]

tests/queries/0_stateless/02707_analyzer_nested_lambdas_types.sql

@@ -0,0 +1,24 @@
+SELECT
+    range(1),
+    arrayMap(x -> arrayMap(x -> x, range(x)), [1])
+SETTINGS allow_experimental_analyzer = 0;
+
+SELECT
+    range(1),
+    arrayMap(x -> arrayMap(x -> x, range(x)), [1])
+SETTINGS allow_experimental_analyzer = 1;
+
+SELECT
+    range(1),
+    arrayMap(x -> arrayMap(x -> 1, range(x)), [1])
+SETTINGS allow_experimental_analyzer = 0;
+
+SELECT
+    range(1),
+    arrayMap(x -> arrayMap(x -> 1, range(x)), [1])
+SETTINGS allow_experimental_analyzer = 1;
+
+SELECT
+    range(1),
+    arrayMap(x -> arrayMap(y -> 1, range(x)), [1])
+SETTINGS allow_experimental_analyzer = 1;