[BugFix] Fix array_map common-expression crash (StarRocks#52909)

stdpain · web-flow · commit 8bda0cbc9f23 · 2024-11-15T10:35:03.000+08:00
Signed-off-by: stdpain &lt;drfeng08@gmail.com&gt;
diff --git a/be/src/exprs/array_map_expr.cpp b/be/src/exprs/array_map_expr.cpp
@@ -60,6 +60,20 @@ Status ArrayMapExpr::prepare(RuntimeState* state, ExprContext* context) {
 
     return Status::OK();
 }
+Status ArrayMapExpr::open(RuntimeState* state, ExprContext* context, FunctionContext::FunctionStateScope scope) {
+    RETURN_IF_ERROR(Expr::open(state, context, scope));
+    for (auto [_, expr] : _outer_common_exprs) {
+        RETURN_IF_ERROR(expr->open(state, context, scope));
+    }
+    return Status::OK();
+}
+
+void ArrayMapExpr::close(RuntimeState* state, ExprContext* context, FunctionContext::FunctionStateScope scope) {
+    for (auto [_, expr] : _outer_common_exprs) {
+        expr->close(state, context, scope);
+    }
+    Expr::close(state, context, scope);
+}
 
 template <bool all_const_input, bool independent_lambda_expr>
 StatusOr<ColumnPtr> ArrayMapExpr::evaluate_lambda_expr(ExprContext* context, Chunk* chunk,
diff --git a/be/src/exprs/array_map_expr.h b/be/src/exprs/array_map_expr.h
@@ -37,6 +37,8 @@ class ArrayMapExpr final : public Expr {
     explicit ArrayMapExpr(TypeDescriptor type);
 
     Status prepare(RuntimeState* state, ExprContext* context) override;
+    Status open(RuntimeState* state, ExprContext* context, FunctionContext::FunctionStateScope scope) override;
+    void close(RuntimeState* state, ExprContext* context, FunctionContext::FunctionStateScope scope) override;
     Expr* clone(ObjectPool* pool) const override { return pool->add(new ArrayMapExpr(*this)); }
 
     StatusOr<ColumnPtr> evaluate_checked(ExprContext* context, Chunk* ptr) override;
diff --git a/test/sql/test_array_fn/R/test_array_map_2 b/test/sql/test_array_fn/R/test_array_map_2
@@ -50,6 +50,10 @@ select array_length(array_map((x,y)->x > y, @arr,@arr)) from table(generate_seri
 1000000
 1000000
 -- !result
+select count(*) from array_map_test where array_map((a1) -> concat(split(id, 'd'), split(id, 'd')), arr_str) is not null;
+-- result:
+1
+-- !result
 -- name: test_array_map_3
 CREATE TABLE `t` (
   `k` bigint NOT NULL COMMENT "",
diff --git a/test/sql/test_array_fn/T/test_array_map_2 b/test/sql/test_array_fn/T/test_array_map_2
@@ -22,6 +22,7 @@ select count(array_map(x->array_length(array_concat(arr_str,[])), arr_largeint))
 
 set @arr=array_repeat("12345",1000000);
 select array_length(array_map((x,y)->x > y, @arr,@arr)) from table(generate_series(1,10,1));
+select count(*) from array_map_test where array_map((a1) -> concat(split(id, 'd'), split(id, 'd')), arr_str) is not null;
 
 -- name: test_array_map_3
 CREATE TABLE `t` (
@@ -72,4 +73,4 @@ select array_map(arg0 -> array_map(arg1 -> array_map(arg2 -> array_map(arg3 -> a
 select array_map(arg0 -> array_map(arg1 -> array_map(arg2 -> array_map(arg3 -> array_map(arg4->array_length(arg2) + arg4, arg3), arg2), arg1), arg0), [[[[[1,2]]]]] );
 
 set @arr=array_generate(1,10000);
-select /*+ SET_VAR(query_mem_limit=104857600)*/array_sum(array_map(x->array_contains(@arr,x), array_generate(1,100000)));
+select /*+ SET_VAR(query_mem_limit=104857600)*/array_sum(array_map(x->array_contains(@arr,x), array_generate(1,100000)));
