ChenMiaoi commented on code in PR #50151:
URL: https://github.com/apache/doris/pull/50151#discussion_r2055733521


##########
be/src/vec/functions/function_jsonb.cpp:
##########
@@ -1758,16 +1834,31 @@ class FunctionJsonSearch : public IFunction {
             }
 
             // an error occurs if any path argument is not a valid path 
expression.
+            // root_path_str = get_start_string(i);
             std::string root_path_str = "$";
+            if (!start_null_check(i)) {
+                root_path_str = get_start_string(i);
+            }
             JsonbPath root_path;
-            root_path.seek(root_path_str.c_str(), root_path_str.size());
+            if (!root_path.seek(root_path_str.c_str(), root_path_str.size())) {
+                return Status::InvalidArgument(
+                        "the start_path argument {} is not a valid json path", 
root_path_str);
+            }
             std::vector<JsonbPath*> paths;
             paths.push_back(&root_path);
 
             if (!search_is_const) {
                 state_ptr = std::make_shared<LikeState>();
                 state_ptr->is_like_pattern = true;
                 const auto& search_str = col_search_string->get_data_at(i);
+
+                // The default is \ if the escape_char argument is missing or 
NULL.

Review Comment:
   done it



##########
be/src/vec/functions/function_jsonb.cpp:
##########
@@ -1919,38 +1983,232 @@ class FunctionJsonSearch : public IFunction {
             };
         }
 
-        // search_str
         ColumnPtr col_search;
-        bool search_is_const = false;
-        std::tie(col_search, search_is_const) =
-                unpack_if_const(block.get_by_position(arguments[2]).column);
+        *search_is_const = false;
 
-        const ColumnString* col_search_string =
-                check_and_get_column<ColumnString>(col_search.get());
-        if (auto* nullable = 
check_and_get_column<ColumnNullable>(col_search.get())) {
-            col_search_string =
-                    
check_and_get_column<ColumnString>(*nullable->get_nested_column_ptr());
-        }
-        if (!col_search_string) {
-            return Status::RuntimeError("Illegal arg pattern {} should be 
ColumnString",
-                                        col_search->get_name());
-        }
-        if (search_is_const) {
-            CheckNullFun search_null_check = always_not_null;
+        RETURN_IF_ERROR(parse_column_args(context, block, arguments, 2, 
search_is_const, col_search,
+                                          col_search_string));
+
+        if (*search_is_const) {
+            search_null_check = always_not_null;
             if (col_search->is_null_at(0)) {
                 search_null_check = always_null;
             }
-            RETURN_IF_ERROR(execute_vector<true>(
+        } else {
+            search_null_check = [col_search](size_t i) { return 
col_search->is_null_at(i); };
+        }
+
+        return Status::OK();
+    }
+};
+
+template <typename Impl>
+class FunctionJsonSearch : public IFunction {
+public:
+    static constexpr auto name = "json_search";
+    static FunctionPtr create() { return 
std::make_shared<FunctionJsonSearch<Impl>>(); }
+
+    String get_name() const override { return name; }
+    bool is_variadic() const override { return true; }
+    DataTypes get_variadic_argument_types_impl() const override {
+        return Impl::get_variadic_argument_types();
+    }
+    size_t get_number_of_arguments() const override {
+        return get_variadic_argument_types_impl().size();
+    }
+    DataTypePtr get_return_type_impl(const DataTypes& arguments) const 
override {
+        return make_nullable(std::make_shared<DataTypeJsonb>());
+    }
+    bool use_default_implementation_for_nulls() const override { return false; 
}
+
+    Status open(FunctionContext* context, FunctionContext::FunctionStateScope 
scope) override {
+        if (scope != FunctionContext::THREAD_LOCAL) {
+            return Status::OK();
+        }
+        if (context->is_col_constant(2)) {
+            std::shared_ptr<LikeState> state = std::make_shared<LikeState>();
+            state->is_like_pattern = true;
+            const auto pattern_col = context->get_constant_col(2)->column_ptr;
+            const auto& pattern = pattern_col->get_data_at(0);
+            RETURN_IF_ERROR(
+                    FunctionLike::construct_like_const_state(context, pattern, 
state, false));
+            context->set_function_state(scope, state);
+        }
+        return Status::OK();
+    }
+
+    Status execute_impl(FunctionContext* context, Block& block, const 
ColumnNumbers& arguments,
+                        uint32_t result, size_t input_rows_count) const 
override {
+        return Impl::execute_impl(context, block, arguments, result, 
input_rows_count);
+    }
+};
+
+#define JSON_SEARCH_NORMAL_PREPARE()                                           
                  \
+    JsonSearchUtil::CheckNullFun json_null_check = 
JsonSearchUtil::always_not_null;              \
+    JsonSearchUtil::GetJsonStringRefFun get_json_fun;                          
                  \
+                                                                               
                  \
+    JsonSearchUtil::CheckNullFun one_null_check = 
JsonSearchUtil::always_not_null;               \
+    JsonSearchUtil::OneFun one_check = JsonSearchUtil::always_one;             
                  \
+                                                                               
                  \
+    JsonSearchUtil::CheckNullFun search_null_check = 
JsonSearchUtil::always_not_null;            \
+    const ColumnString* col_search_string;                                     
                  \
+    bool search_is_const = false;                                              
                  \
+                                                                               
                  \
+    RETURN_IF_ERROR(JsonSearchUtil::execute_prepare(                           
                  \
+            context, block, arguments, json_null_check, get_json_fun, 
one_null_check, one_check, \
+            search_null_check, col_search_string, &search_is_const))
+
+struct JsonSearchNormal {
+    static DataTypes get_variadic_argument_types() {
+        return {
+                std::make_shared<DataTypeString>(),
+                std::make_shared<DataTypeString>(),
+                std::make_shared<DataTypeString>(),
+        };
+    }
+
+    static Status execute_impl(FunctionContext* context, Block& block,
+                               const ColumnNumbers& arguments, uint32_t result,
+                               size_t input_rows_count) {
+        if (arguments.size() != 3) {
+            return Status::InvalidArgument("wrong arguments for function 
json_search");
+        }
+
+        JSON_SEARCH_NORMAL_PREPARE();
+
+        if (search_is_const) {
+            RETURN_IF_ERROR(JsonSearchUtil::execute_vector<true>(
                     block, input_rows_count, json_null_check, get_json_fun, 
one_null_check,
                     one_check, search_null_check, col_search_string, context, 
result));
         } else {
-            CheckNullFun search_null_check = [col_search](size_t i) {
-                return col_search->is_null_at(i);
-            };
-            RETURN_IF_ERROR(execute_vector<false>(
+            RETURN_IF_ERROR(JsonSearchUtil::execute_vector<false>(
                     block, input_rows_count, json_null_check, get_json_fun, 
one_null_check,
                     one_check, search_null_check, col_search_string, context, 
result));
         }
+
+        return Status::OK();
+    }
+};
+
+#define JSON_SEARCH_EXTRA_PREPARE(arg_name, index, null_check_fn, get_type, 
get_fn)         \
+    JsonSearchUtil::CheckNullFun arg_name##_null_check = 
JsonSearchUtil::always_null;       \
+    get_type get_##arg_name##_string;                                          
             \
+                                                                               
             \
+    ColumnPtr col_##arg_name;                                                  
             \
+    bool arg_name##_is_const = false;                                          
             \
+    const ColumnString* col_##arg_name##_string;                               
             \
+                                                                               
             \
+    RETURN_IF_ERROR(JsonSearchUtil::parse_column_args(context, block, 
arguments, index,     \
+                                                      &arg_name##_is_const, 
col_##arg_name, \
+                                                      
col_##arg_name##_string));            \
+                                                                               
             \
+    if (arg_name##_is_const) {                                                 
             \
+        if (col_##arg_name->is_null_at(0)) {                                   
             \
+            arg_name##_null_check = JsonSearchUtil::always_null;               
             \
+        }                                                                      
             \
+    } else {                                                                   
             \
+        arg_name##_null_check = null_check_fn;                                 
             \
+        get_##arg_name##_string = get_fn;                                      
             \
+    }
+
+struct JsonSearchEscape {
+    static DataTypes get_variadic_argument_types() {
+        return {
+                std::make_shared<DataTypeString>(),
+                std::make_shared<DataTypeString>(),
+                std::make_shared<DataTypeString>(),
+                std::make_shared<DataTypeString>(),
+        };
+    }
+
+    static Status execute_impl(FunctionContext* context, Block& block,
+                               const ColumnNumbers& arguments, uint32_t result,
+                               size_t input_rows_count) {
+        if (arguments.size() != 4) {
+            return Status::InvalidArgument("wrong arguments for function 
json_search");
+        }
+
+        JSON_SEARCH_NORMAL_PREPARE();
+
+        JSON_SEARCH_EXTRA_PREPARE(
+                escape, 3, [col_escape](size_t i) { return 
col_escape->is_null_at(i); },
+                JsonSearchUtil::GetJsonEscapeFun,
+                [col_escape_string](size_t i, char* escape_char) {
+                    auto escape_string = 
col_escape_string->get_data_at(i).to_string();
+                    if (escape_string.length() != 1) {
+                        return Status::RuntimeError("Illegal arg pattern {} 
should be char",
+                                                    
col_escape_string->get_name());
+                    }
+                    *escape_char = escape_string.at(0);
+                    return Status::OK();
+                });
+
+        if (search_is_const) {
+            RETURN_IF_ERROR(JsonSearchUtil::execute_vector<true>(
+                    block, input_rows_count, json_null_check, get_json_fun, 
one_null_check,
+                    one_check, search_null_check, col_search_string, context, 
result,
+                    escape_null_check, get_escape_string));
+        } else {
+            RETURN_IF_ERROR(JsonSearchUtil::execute_vector<false>(
+                    block, input_rows_count, json_null_check, get_json_fun, 
one_null_check,
+                    one_check, search_null_check, col_search_string, context, 
result,
+                    escape_null_check, get_escape_string));
+        }
+
+        return Status::OK();
+    }
+};
+
+struct JsonSearchStartPath {
+    static DataTypes get_variadic_argument_types() {
+        return {
+                std::make_shared<DataTypeString>(), 
std::make_shared<DataTypeString>(),
+                std::make_shared<DataTypeString>(), 
std::make_shared<DataTypeString>(),
+                std::make_shared<DataTypeString>(),
+        };
+    }
+
+    static Status execute_impl(FunctionContext* context, Block& block,
+                               const ColumnNumbers& arguments, uint32_t result,
+                               size_t input_rows_count) {
+        if (arguments.size() != 5) {
+            return Status::InvalidArgument("wrong arguments for function 
json_search");
+        }
+
+        JSON_SEARCH_NORMAL_PREPARE();
+
+        JSON_SEARCH_EXTRA_PREPARE(
+                escape, 3, [col_escape](size_t i) { return 
col_escape->is_null_at(i); },
+                JsonSearchUtil::GetJsonEscapeFun,
+                [col_escape_string](size_t i, char* escape_char) {
+                    auto escape_string = 
col_escape_string->get_data_at(i).to_string();

Review Comment:
   done it



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org
For additional commands, e-mail: commits-h...@doris.apache.org

Reply via email to