From 50c12d24c82988fea82521d42851aa3c609e32c7 Mon Sep 17 00:00:00 2001 From: JeanArhancet Date: Sun, 21 Jul 2024 16:30:40 +0200 Subject: [PATCH] feat: add length scalar function --- core/expr.rs | 174 ++++++++++------------------------ core/function.rs | 3 + core/vdbe.rs | 40 +++++++- testing/scalar-functions.test | 20 ++++ 4 files changed, 111 insertions(+), 126 deletions(-) diff --git a/core/expr.rs b/core/expr.rs index 858e059af..fbdfff86d 100644 --- a/core/expr.rs +++ b/core/expr.rs @@ -222,10 +222,8 @@ pub fn translate_expr( args, filter_over: _, } => { - let func_type: Option = match normalize_ident(name.0.as_str()).as_str().parse() { - Ok(func) => Some(func), - Err(_) => None, - }; + let func_type: Option = normalize_ident(name.0.as_str()).as_str().parse().ok(); + match func_type { Some(Func::Agg(_)) => { anyhow::bail!("Parse error: aggregation function in non-aggregation context") @@ -236,12 +234,16 @@ pub fn translate_expr( let args = if let Some(args) = args { if args.len() < 2 { anyhow::bail!( - "Parse error: coalesce function with less than 2 arguments" + "Parse error: {} function with less than 2 arguments", + srf.to_string() ); } args } else { - anyhow::bail!("Parse error: coalesce function with no arguments"); + anyhow::bail!( + "Parse error: {} function with no arguments", + srf.to_string() + ); }; // coalesce function is implemented as a series of not null checks @@ -267,16 +269,20 @@ pub fn translate_expr( let args = if let Some(args) = args { if args.len() < 2 { anyhow::bail!( - "Parse error: like function with less than 2 arguments" + "Parse error: {} function with less than 2 arguments", + srf.to_string() ); } args } else { - anyhow::bail!("Parse error: like function with no arguments"); + anyhow::bail!( + "Parse error: {} function with no arguments", + srf.to_string() + ); }; for arg in args { let reg = program.alloc_register(); - let _ = translate_expr(program, select, arg, reg)?; + let _ = translate_expr(program, select, &arg, reg)?; match arg { ast::Expr::Literal(_) => program.mark_last_insn_constant(), _ => {} @@ -285,161 +291,81 @@ pub fn translate_expr( program.emit_insn(Insn::Function { start_reg: target_register + 1, dest: target_register, - func: SingleRowFunc::Like, + func: srf, }); Ok(target_register) } - SingleRowFunc::Abs => { + SingleRowFunc::Abs + | SingleRowFunc::Lower + | SingleRowFunc::Upper + | SingleRowFunc::Length => { let args = if let Some(args) = args { if args.len() != 1 { anyhow::bail!( - "Parse error: abs function with not exactly 1 argument" + "Parse error: {} function with not exactly 1 argument", + srf.to_string() ); } args } else { - anyhow::bail!("Parse error: abs function with no arguments"); + anyhow::bail!( + "Parse error: {} function with no arguments", + srf.to_string() + ); }; let regs = program.alloc_register(); - let _ = translate_expr(program, select, &args[0], regs)?; + translate_expr(program, select, &args[0], regs)?; program.emit_insn(Insn::Function { start_reg: regs, dest: target_register, - func: SingleRowFunc::Abs, + func: srf, }); - - Ok(target_register) - } - SingleRowFunc::Upper => { - let args = if let Some(args) = args { - if args.len() != 1 { - anyhow::bail!( - "Parse error: upper function with not exactly 1 argument" - ); - } - args - } else { - anyhow::bail!("Parse error: upper function with no arguments"); - }; - - let regs = program.alloc_register(); - let _ = translate_expr(program, select, &args[0], regs)?; - program.emit_insn(Insn::Function { - start_reg: regs, - dest: target_register, - func: SingleRowFunc::Upper, - }); - - Ok(target_register) - } - SingleRowFunc::Lower => { - let args = if let Some(args) = args { - if args.len() != 1 { - anyhow::bail!( - "Parse error: lower function with not exactly 1 argument" - ); - } - args - } else { - anyhow::bail!("Parse error: lower function with no arguments"); - }; - - let regs = program.alloc_register(); - let _ = translate_expr(program, select, &args[0], regs)?; - program.emit_insn(Insn::Function { - start_reg: regs, - dest: target_register, - func: SingleRowFunc::Lower, - }); - Ok(target_register) } SingleRowFunc::Random => { if args.is_some() { - anyhow::bail!("Parse error: random function with arguments"); + anyhow::bail!( + "Parse error: {} function with arguments", + srf.to_string() + ); } let regs = program.alloc_register(); - program.emit_insn(Insn::Function { start_reg: regs, dest: target_register, - func: SingleRowFunc::Random, + func: srf, }); Ok(target_register) } - SingleRowFunc::Trim => { + SingleRowFunc::Trim | SingleRowFunc::Round => { let args = if let Some(args) = args { if args.len() > 2 { anyhow::bail!( - "Parse error: trim function with more than 2 arguments" + "Parse error: {} function with more than 2 arguments", + srf.to_string() ); } args } else { - anyhow::bail!("Parse error: trim function with no arguments"); + anyhow::bail!( + "Parse error: {} function with no arguments", + srf.to_string() + ); }; - if args.len() == 1 { - let regs = program.alloc_register(); - let _ = translate_expr(program, select, &args[0], regs)?; - program.emit_insn(Insn::Function { - start_reg: regs, - dest: target_register, - func: SingleRowFunc::Trim, - }); - } else { - for arg in args { - let reg = program.alloc_register(); - let _ = translate_expr(program, select, arg, reg)?; - match arg { - ast::Expr::Literal(_) => program.mark_last_insn_constant(), - _ => {} - } + for arg in args.iter() { + let reg = program.alloc_register(); + translate_expr(program, select, arg, reg)?; + if let ast::Expr::Literal(_) = arg { + program.mark_last_insn_constant(); } - program.emit_insn(Insn::Function { - start_reg: target_register + 1, - dest: target_register, - func: SingleRowFunc::Trim, - }); - } - Ok(target_register) - } - SingleRowFunc::Round => { - let args = if let Some(args) = args { - if args.len() > 2 { - anyhow::bail!( - "Parse error: round function with more than 2 arguments" - ); - } - args - } else { - anyhow::bail!("Parse error: round function with no arguments"); - }; - - if args.len() == 1 { - let regs = program.alloc_register(); - let _ = translate_expr(program, select, &args[0], regs)?; - program.emit_insn(Insn::Function { - start_reg: regs, - dest: target_register, - func: SingleRowFunc::Round, - }); - } else { - for arg in args { - let reg = program.alloc_register(); - let _ = translate_expr(program, select, arg, reg)?; - match arg { - ast::Expr::Literal(_) => program.mark_last_insn_constant(), - _ => {} - } - } - program.emit_insn(Insn::Function { - start_reg: target_register + 1, - dest: target_register, - func: SingleRowFunc::Round, - }); } + program.emit_insn(Insn::Function { + start_reg: target_register + 1, + dest: target_register, + func: srf, + }); Ok(target_register) } } diff --git a/core/function.rs b/core/function.rs index 0c3e4b271..224e96500 100644 --- a/core/function.rs +++ b/core/function.rs @@ -37,6 +37,7 @@ pub enum SingleRowFunc { Random, Trim, Round, + Length, } impl ToString for SingleRowFunc { @@ -50,6 +51,7 @@ impl ToString for SingleRowFunc { SingleRowFunc::Random => "random".to_string(), SingleRowFunc::Trim => "trim".to_string(), SingleRowFunc::Round => "round".to_string(), + SingleRowFunc::Length => "length".to_string(), } } } @@ -80,6 +82,7 @@ impl FromStr for Func { "random" => Ok(Func::SingleRow(SingleRowFunc::Random)), "trim" => Ok(Func::SingleRow(SingleRowFunc::Trim)), "round" => Ok(Func::SingleRow(SingleRowFunc::Round)), + "length" => Ok(Func::SingleRow(SingleRowFunc::Length)), _ => Err(()), } } diff --git a/core/vdbe.rs b/core/vdbe.rs index 383f5cdec..40d56aea5 100644 --- a/core/vdbe.rs +++ b/core/vdbe.rs @@ -1289,6 +1289,11 @@ impl Program { } state.pc += 1; } + SingleRowFunc::Length => { + let reg_value = state.registers[*start_reg].borrow_mut(); + state.registers[*dest] = exec_length(reg_value); + state.pc += 1; + } SingleRowFunc::Random => { state.registers[*dest] = exec_random(); state.pc += 1; @@ -1856,12 +1861,23 @@ fn exec_lower(reg: &OwnedValue) -> Option { } } +fn exec_length(reg: &OwnedValue) -> OwnedValue { + match reg { + OwnedValue::Text(_) | OwnedValue::Integer(_) | OwnedValue::Float(_) => { + OwnedValue::Integer(reg.to_string().len() as i64) + } + OwnedValue::Blob(blob) => OwnedValue::Integer(blob.len() as i64), + _ => reg.to_owned(), + } +} + fn exec_upper(reg: &OwnedValue) -> Option { match reg { OwnedValue::Text(t) => Some(OwnedValue::Text(Rc::new(t.to_uppercase()))), t => Some(t.to_owned()), } } + fn exec_abs(reg: &OwnedValue) -> Option { match reg { OwnedValue::Integer(x) => { @@ -1951,11 +1967,31 @@ fn exec_if(reg: &OwnedValue, null_reg: &OwnedValue, not: bool) -> bool { #[cfg(test)] mod tests { use super::{ - exec_abs, exec_if, exec_like, exec_lower, exec_random, exec_round, exec_trim, exec_upper, - OwnedValue, + exec_abs, exec_if, exec_length, exec_like, exec_lower, exec_random, exec_round, exec_trim, + exec_upper, OwnedValue, }; use std::rc::Rc; + #[test] + fn test_length() { + let input_str = OwnedValue::Text(Rc::new(String::from("bob"))); + let expected_len = OwnedValue::Integer(3); + assert_eq!(exec_length(&input_str), expected_len); + + let input_integer = OwnedValue::Integer(123); + let expected_len = OwnedValue::Integer(3); + assert_eq!(exec_length(&input_integer), expected_len); + + let input_float = OwnedValue::Float(123.456); + let expected_len = OwnedValue::Integer(7); + assert_eq!(exec_length(&input_float), expected_len); + + // Expected byte array for "example" + let expected_blob = OwnedValue::Blob(Rc::new(vec![101, 120, 97, 109, 112, 108, 101])); + let expected_len = OwnedValue::Integer(7); + assert_eq!(exec_length(&expected_blob), expected_len); + } + #[test] fn test_trim() { let input_str = OwnedValue::Text(Rc::new(String::from(" Bob and Alice "))); diff --git a/testing/scalar-functions.test b/testing/scalar-functions.test index 5cdb864b6..a3b3d2114 100644 --- a/testing/scalar-functions.test +++ b/testing/scalar-functions.test @@ -122,3 +122,23 @@ do_execsql_test round-float-zero-precision { do_execsql_test round-null-precision { SELECT round(123.456, null); } {} + +do_execsql_test length-text { + SELECT length('limbo'); +} {5} + +do_execsql_test length-integer { + SELECT length(12345); +} {5} + +do_execsql_test length-float { + SELECT length(123.456); +} {7} + +do_execsql_test length-null { + SELECT length(NULL); +} {} + +do_execsql_test length-empty-text { + SELECT length(''); +} {0} \ No newline at end of file