[BugFix] jit can't repreapre/reopen some exprs,like UDF, dict

Signed-off-by: Zhuhe Fang <fzhedu@gmail.com>

be/src/exprs/dictionary_get_expr.cpp

@@ -74,6 +74,10 @@ StatusOr<ColumnPtr> DictionaryGetExpr::evaluate_checked(ExprContext* context, Ch
 
 Status DictionaryGetExpr::prepare(RuntimeState* state, ExprContext* context) {
     RETURN_IF_ERROR(Expr::prepare(state, context));
+    if (_is_prepared) {
+        return Status::OK();
+    }
+    _is_prepared = true;
     _runtime_state = state;
 
     _schema = StorageEngine::instance()->dictionary_cache_manager()->get_dictionary_schema_by_id(

be/src/exprs/dictionary_get_expr.h

@@ -44,6 +44,7 @@ class DictionaryGetExpr final : public Expr {
     ChunkPtr _key_chunk = nullptr;
     ChunkPtr _value_chunk = nullptr;
     ColumnPtr _struct_column = nullptr;
+    bool _is_prepared = false;
 };
 
 } // namespace starrocks

be/src/exprs/expr.cpp

@@ -227,6 +227,9 @@ Status Expr::create_expr_tree(ObjectPool* pool, const TExpr& texpr, ExprContext*
 }
 
 Status Expr::rewrite_jit_exprs(std::vector<ExprContext*>& expr_ctxs, ObjectPool* pool, RuntimeState* state) {
+    if (!state->is_jit_enabled()) {
+        return Status::OK();
+    }
     std::vector<ExprContext*> tmp;
     for (auto ctx : expr_ctxs) {
         auto* prev_expr = ctx->root();
@@ -242,14 +245,15 @@ Status Expr::rewrite_jit_exprs(std::vector<ExprContext*>& expr_ctxs, ObjectPool*
         if (replaced) {
             // The node was replaced, so we need to update the context.
             new_ctx = pool->add(new ExprContext(*root_expr));
+            if (state != nullptr) {
+                RETURN_IF_ERROR((*root_expr)->prepare(state, new_ctx));
+                RETURN_IF_ERROR((*root_expr)->open(state, new_ctx, FunctionContext::THREAD_LOCAL));
+            }
         }
         tmp.emplace_back(new_ctx);
     }
     expr_ctxs.swap(tmp);
-    if (state != nullptr) {
-        RETURN_IF_ERROR(Expr::prepare(expr_ctxs, state));
-        RETURN_IF_ERROR(Expr::open(expr_ctxs, state));
-    }
+
     return Status::OK();
 }
 

be/src/exprs/java_function_call_expr.cpp

@@ -125,9 +125,13 @@ JavaFunctionCallExpr::~JavaFunctionCallExpr() {
 
 // TODO support prepare UDF
 Status JavaFunctionCallExpr::prepare(RuntimeState* state, ExprContext* context) {
-    _runtime_state = state;
     // init Expr::prepare
     RETURN_IF_ERROR(Expr::prepare(state, context));
+    if (_is_prepared) {
+        return Status::OK();
+    }
+    _is_prepared = true;
+    _runtime_state = state;
 
     if (!_fn.__isset.fid) {
         return Status::InternalError("Not Found function id for " + _fn.name.function_name);

be/src/exprs/java_function_call_expr.h

@@ -45,5 +45,6 @@ class JavaFunctionCallExpr final : public Expr {
     std::shared_ptr<JavaUDFContext> _func_desc;
     std::shared_ptr<UDFFunctionCallHelper> _call_helper;
     bool _is_returning_random_value;
+    bool _is_prepared = false;
 };
 } // namespace starrocks