From dbaef4e1cd93f8b4ee7e222f2b1af3059f2cdfad Mon Sep 17 00:00:00 2001 From: Ashish Kolhe <35160958+oasisk@users.noreply.github.com> Date: Mon, 10 Apr 2023 15:55:23 +0530 Subject: [PATCH] fix: query results issue (#535) * fix query results * fixed clippy warnings * fixes --- src/service/functions.rs | 29 +++++++--------- .../search/datafusion/transform_udf.rs | 33 +++++++++---------- src/service/search/mod.rs | 7 ++-- web/src/components/functions/add.vue | 3 +- 4 files changed, 32 insertions(+), 40 deletions(-) diff --git a/src/service/functions.rs b/src/service/functions.rs index c1c926b0eb..136a77554e 100644 --- a/src/service/functions.rs +++ b/src/service/functions.rs @@ -55,8 +55,8 @@ pub async fn register_function( } trans.stream_type = s_type; - let js_func = trans.function.to_owned(); - if trans.trans_type == 0 && !js_func.contains('(') && !js_func.contains(')') { + let func = trans.function.to_owned(); + if trans.trans_type == 0 && !func.contains('(') && !func.contains(')') { msg.push_str(" not valid function"); is_err = true; } @@ -70,11 +70,8 @@ pub async fn register_function( } else { trans.stream_name = stream_name.to_string(); trans.name = name.to_string(); - if trans.trans_type == 0 { - extract_num_args(&mut trans); - } else { - trans.num_args = 1; - } + extract_num_args(&mut trans); + db::functions::set( org_id.as_str(), Some(stream_name.to_string()), @@ -93,16 +90,12 @@ pub async fn register_function( None => { let mut is_err = false; trans.name = name.to_string(); - let js_func = trans.function.to_owned(); - if trans.trans_type == 0 && !js_func.contains('(') && !js_func.contains(')') { + let func = trans.function.to_owned(); + if trans.trans_type == 0 && !func.contains('(') && !func.contains(')') { is_err = true; } if !is_err { - if trans.trans_type == 0 { - extract_num_args(&mut trans); - } else { - trans.num_args = 2; - } + extract_num_args(&mut trans); db::functions::set(org_id.as_str(), None, None, name.as_str(), trans) .await .unwrap(); @@ -158,10 +151,10 @@ pub async fn delete_function( } fn extract_num_args(trans: &mut Transform) { - let js_func = trans.function.to_owned(); - let start_stream = js_func.find('(').unwrap(); - let end_stream = js_func.find(')').unwrap(); - let args = &js_func[start_stream + 1..end_stream].trim(); + let func = trans.function.to_owned(); + let start_stream = func.find('(').unwrap(); + let end_stream = func.find(')').unwrap(); + let args = &func[start_stream + 1..end_stream].trim(); if args.is_empty() { trans.num_args = 0; } else { diff --git a/src/service/search/datafusion/transform_udf.rs b/src/service/search/datafusion/transform_udf.rs index 9720f4ef60..ff6e1cfba6 100644 --- a/src/service/search/datafusion/transform_udf.rs +++ b/src/service/search/datafusion/transform_udf.rs @@ -21,7 +21,6 @@ use datafusion::{ physical_plan::functions::make_scalar_function, prelude::create_udf, }; - use mlua::{Function, Lua, LuaSerdeExt, MultiValue}; use std::sync::Arc; use vrl::{prelude::BTreeMap, TargetValueRef, VrlRuntime}; @@ -75,13 +74,15 @@ pub async fn get_all_transform(org_id: &str) -> Vec datafusion::logical_expr::ScalarUDF { +fn get_udf_vrl(fn_name: String, func: &str, num_args: u8) -> datafusion::logical_expr::ScalarUDF { let local_fn_name = fn_name; - let local_js_func = js_func.to_owned(); + let mut local_func = func.trim().to_owned(); + + if local_func.starts_with("function(") { + if let Some(idx) = local_func.find(')') { + local_func = local_func[idx + 1..].to_owned(); + } + } let pow_calc = move |args: &[ArrayRef]| { let len = args[0].len(); @@ -95,9 +96,9 @@ fn get_udf_vrl( .as_any() .downcast_ref::() .expect("cast failed"); - obj_str.push_str(&format!("col{j} = \"{}\" \n", col.value(i))); + obj_str.push_str(&format!("col{} = \"{}\" \n", j + 1, col.value(i))); } - obj_str.push_str(&format!(" \n {}", &local_js_func)); + obj_str.push_str(&format!(" \n {}", &local_func)); let func = compile_vrl_function(&obj_str).unwrap(); data_vec.insert(i, apply_vrl_fn(&mut runtime, func)); } @@ -156,18 +157,14 @@ pub fn apply_vrl_fn(runtime: &mut Runtime, program: vrl::Program) -> String { } } -fn get_udf_lua( - fn_name: String, - js_func: &str, - num_args: u8, -) -> datafusion::logical_expr::ScalarUDF { +fn get_udf_lua(fn_name: String, func: &str, num_args: u8) -> datafusion::logical_expr::ScalarUDF { let local_fn_name = fn_name; - let local_js_func = js_func.to_owned(); + let local_func = func.to_owned(); let pow_calc = move |args: &[ArrayRef]| { //Lua let lua = Lua::new(); //Register Lua Function - let func = load_transform(&lua, local_js_func.to_string()); + let func = load_transform(&lua, local_func.to_string()); let len = args[0].len(); let mut data_vec = vec![]; @@ -193,8 +190,8 @@ fn get_udf_lua( pow_udf } -fn load_transform(lua: &Lua, js_func: String) -> Function { - lua.load(&js_func).eval().unwrap() +fn load_transform(lua: &Lua, func: String) -> Function { + lua.load(&func).eval().unwrap() } fn lua_transform(lua: &Lua, func: &Function, args: &[ArrayRef], stream: usize) -> String { diff --git a/src/service/search/mod.rs b/src/service/search/mod.rs index 6a5c2d0be7..90496f50c0 100644 --- a/src/service/search/mod.rs +++ b/src/service/search/mod.rs @@ -744,7 +744,10 @@ fn apply_query_fn(req: &cluster_rpc::SearchRequest, resp: &[json::Value]) -> Vec } } } + return ret_resp; }; - }; - ret_resp + resp.to_vec() + } else { + resp.to_vec() + } } diff --git a/web/src/components/functions/add.vue b/web/src/components/functions/add.vue index 07ea168727..7dadadd869 100644 --- a/web/src/components/functions/add.vue +++ b/web/src/components/functions/add.vue @@ -202,8 +202,7 @@ end`; suffixCode.value = ""; } else { prefixCode.value = `function()`; - suffixCode.value = ` -end`; + suffixCode.value = ""; } }