[improve](profile) add timer for record udf execute time (apache#41779)

## Proposed changes <img width="314" alt="image" src="https://github.com/user-attachments/assets/d4f4e2fc-6736-4660-a2c7-12e01eb1a4d1">
zhangstar333 · Nov 21, 2024 · fe88b07 · fe88b07
1 parent e2c4b6f
commit fe88b07
Show file tree

Hide file tree

Showing 12 changed files with 33 additions and 8 deletions.
diff --git a/be/src/pipeline/pipeline_task.h b/be/src/pipeline/pipeline_task.h
@@ -224,6 +224,8 @@ class PipelineTask {
 
     RuntimeState* runtime_state() const { return _state; }
 
+    RuntimeProfile* get_task_profile() const { return _task_profile.get(); }
+
     std::string task_name() const { return fmt::format("task{}({})", _index, _pipeline->_name); }
 
     void stop_if_finished() {

diff --git a/be/src/udf/udf.h b/be/src/udf/udf.h
@@ -26,6 +26,7 @@
 #include <vector>
 
 #include "runtime/types.h"
+#include "util/runtime_profile.h"
 #include "vec/common/arena.h"
 
 namespace doris {
@@ -88,6 +89,12 @@ class FunctionContext {
         _jsonb_string_as_string = jsonb_string_as_string;
     }
 
+    void set_udf_execute_timer(RuntimeProfile::Counter* udf_execute_timer) {
+        _udf_execute_timer = udf_execute_timer;
+    }
+
+    RuntimeProfile::Counter* get_udf_execute_timer() { return _udf_execute_timer; }
+
     // Cast flag, when enable string_as_jsonb_string, string casting to jsonb will not parse string
     // instead just insert a string literal
     bool string_as_jsonb_string() const { return _string_as_jsonb_string; }
@@ -176,6 +183,8 @@ class FunctionContext {
 
     std::vector<std::shared_ptr<doris::ColumnPtrWrapper>> _constant_cols;
 
+    //udf execute timer
+    RuntimeProfile::Counter* _udf_execute_timer = nullptr;
     bool _check_overflow_for_decimal = false;
 
     bool _string_as_jsonb_string = false;

diff --git a/be/src/vec/exprs/vcase_expr.cpp b/be/src/vec/exprs/vcase_expr.cpp
@@ -86,7 +86,7 @@ Status VCaseExpr::open(RuntimeState* state, VExprContext* context,
     for (auto& i : _children) {
         RETURN_IF_ERROR(i->open(state, context, scope));
     }
-    RETURN_IF_ERROR(VExpr::init_function_context(context, scope, _function));
+    RETURN_IF_ERROR(VExpr::init_function_context(state, context, scope, _function));
     if (scope == FunctionContext::FRAGMENT_LOCAL) {
         RETURN_IF_ERROR(VExpr::get_const_col(context, nullptr));
     }

diff --git a/be/src/vec/exprs/vcast_expr.cpp b/be/src/vec/exprs/vcast_expr.cpp
@@ -89,7 +89,7 @@ doris::Status VCastExpr::open(doris::RuntimeState* state, VExprContext* context,
     for (auto& i : _children) {
         RETURN_IF_ERROR(i->open(state, context, scope));
     }
-    RETURN_IF_ERROR(VExpr::init_function_context(context, scope, _function));
+    RETURN_IF_ERROR(VExpr::init_function_context(state, context, scope, _function));
     if (scope == FunctionContext::FRAGMENT_LOCAL) {
         RETURN_IF_ERROR(VExpr::get_const_col(context, nullptr));
     }

diff --git a/be/src/vec/exprs/vectorized_fn_call.cpp b/be/src/vec/exprs/vectorized_fn_call.cpp
@@ -27,6 +27,7 @@
 #include "common/config.h"
 #include "common/consts.h"
 #include "common/status.h"
+#include "pipeline/pipeline_task.h"
 #include "runtime/runtime_state.h"
 #include "udf/udf.h"
 #include "vec/columns/column.h"
@@ -125,7 +126,7 @@ Status VectorizedFnCall::open(RuntimeState* state, VExprContext* context,
     for (auto& i : _children) {
         RETURN_IF_ERROR(i->open(state, context, scope));
     }
-    RETURN_IF_ERROR(VExpr::init_function_context(context, scope, _function));
+    RETURN_IF_ERROR(VExpr::init_function_context(state, context, scope, _function));
     if (scope == FunctionContext::FRAGMENT_LOCAL) {
         RETURN_IF_ERROR(VExpr::get_const_col(context, nullptr));
     }

diff --git a/be/src/vec/exprs/vexpr.cpp b/be/src/vec/exprs/vexpr.cpp
@@ -31,6 +31,7 @@
 #include "common/config.h"
 #include "common/exception.h"
 #include "common/status.h"
+#include "pipeline/pipeline_task.h"
 #include "runtime/define_primitive_type.h"
 #include "vec/columns/column_vector.h"
 #include "vec/columns/columns_number.h"
@@ -562,7 +563,7 @@ void VExpr::register_function_context(RuntimeState* state, VExprContext* context
     _fn_context_index = context->register_function_context(state, _type, arg_types);
 }
 
-Status VExpr::init_function_context(VExprContext* context,
+Status VExpr::init_function_context(RuntimeState* state, VExprContext* context,
                                     FunctionContext::FunctionStateScope scope,
                                     const FunctionBasePtr& function) const {
     FunctionContext* fn_ctx = context->fn_context(_fn_context_index);
@@ -574,6 +575,12 @@ Status VExpr::init_function_context(VExprContext* context,
             constant_cols.push_back(const_col);
         }
         fn_ctx->set_constant_cols(constant_cols);
+    } else {
+        if (function->is_udf_function()) {
+            auto* timer = ADD_TIMER(state->get_task()->get_task_profile(),
+                                    "UDF[" + function->get_name() + "]");
+            fn_ctx->set_udf_execute_timer(timer);
+        }
     }
 
     if (scope == FunctionContext::FRAGMENT_LOCAL) {

diff --git a/be/src/vec/exprs/vexpr.h b/be/src/vec/exprs/vexpr.h
@@ -286,7 +286,8 @@ class VExpr {
     /// 1. Set constant columns result of function arguments.
     /// 2. Call function's prepare() to initialize function state, fragment-local or
     /// thread-local according the input `FunctionStateScope` argument.
-    Status init_function_context(VExprContext* context, FunctionContext::FunctionStateScope scope,
+    Status init_function_context(RuntimeState* state, VExprContext* context,
+                                 FunctionContext::FunctionStateScope scope,
                                  const FunctionBasePtr& function) const;
 
     /// Helper function to close function context, fragment-local or thread-local according

diff --git a/be/src/vec/exprs/vin_predicate.cpp b/be/src/vec/exprs/vin_predicate.cpp
@@ -91,7 +91,7 @@ Status VInPredicate::open(RuntimeState* state, VExprContext* context,
     for (auto& child : _children) {
         RETURN_IF_ERROR(child->open(state, context, scope));
     }
-    RETURN_IF_ERROR(VExpr::init_function_context(context, scope, _function));
+    RETURN_IF_ERROR(VExpr::init_function_context(state, context, scope, _function));
     if (scope == FunctionContext::FRAGMENT_LOCAL) {
         RETURN_IF_ERROR(VExpr::get_const_col(context, nullptr));
     }

diff --git a/be/src/vec/exprs/vmatch_predicate.cpp b/be/src/vec/exprs/vmatch_predicate.cpp
@@ -109,7 +109,7 @@ Status VMatchPredicate::open(RuntimeState* state, VExprContext* context,
     for (int i = 0; i < _children.size(); ++i) {
         RETURN_IF_ERROR(_children[i]->open(state, context, scope));
     }
-    RETURN_IF_ERROR(VExpr::init_function_context(context, scope, _function));
+    RETURN_IF_ERROR(VExpr::init_function_context(state, context, scope, _function));
     if (scope == FunctionContext::THREAD_LOCAL || scope == FunctionContext::FRAGMENT_LOCAL) {
         context->fn_context(_fn_context_index)->set_function_state(scope, _inverted_index_ctx);
     }

diff --git a/be/src/vec/functions/function.h b/be/src/vec/functions/function.h
@@ -210,6 +210,8 @@ class IFunctionBase {
 
     virtual bool is_use_default_implementation_for_constants() const = 0;
 
+    virtual bool is_udf_function() const { return false; }
+
     /// The property of monotonicity for a certain range.
     struct Monotonicity {
         bool is_monotonic = false; /// Is the function monotonous (nondecreasing or nonincreasing).

diff --git a/be/src/vec/functions/function_java_udf.cpp b/be/src/vec/functions/function_java_udf.cpp
@@ -44,6 +44,7 @@ Status JavaFunctionCall::open(FunctionContext* context, FunctionContext::Functio
     RETURN_IF_ERROR(JniUtil::GetJNIEnv(&env));
 
     if (scope == FunctionContext::FunctionStateScope::THREAD_LOCAL) {
+        SCOPED_TIMER(context->get_udf_execute_timer());
         std::shared_ptr<JniContext> jni_ctx = std::make_shared<JniContext>();
         context->set_function_state(FunctionContext::THREAD_LOCAL, jni_ctx);
 
@@ -96,7 +97,7 @@ Status JavaFunctionCall::execute_impl(FunctionContext* context, Block& block,
     RETURN_IF_ERROR(JniUtil::GetJNIEnv(&env));
     JniContext* jni_ctx = reinterpret_cast<JniContext*>(
             context->get_function_state(FunctionContext::THREAD_LOCAL));
-
+    SCOPED_TIMER(context->get_udf_execute_timer());
     std::unique_ptr<long[]> input_table;
     RETURN_IF_ERROR(JniConnector::to_java_table(&block, num_rows, arguments, input_table));
     auto input_table_schema = JniConnector::parse_table_schema(&block, arguments, true);

diff --git a/be/src/vec/functions/function_java_udf.h b/be/src/vec/functions/function_java_udf.h
@@ -107,6 +107,8 @@ class JavaFunctionCall : public IFunctionBase {
 
     bool is_use_default_implementation_for_constants() const override { return true; }
 
+    bool is_udf_function() const override { return true; }
+
 private:
     const TFunction& fn_;
     const DataTypes _argument_types;