diff --git a/core/function.rs b/core/function.rs index 7cc6d9cc3..8937726b9 100644 --- a/core/function.rs +++ b/core/function.rs @@ -159,6 +159,7 @@ pub enum VectorFunc { VectorDistanceCos, VectorDistanceEuclidean, VectorConcat, + VectorSlice, } impl VectorFunc { @@ -178,6 +179,7 @@ impl Display for VectorFunc { // We use `distance_l2` to reduce user input Self::VectorDistanceEuclidean => "vector_distance_l2".to_string(), Self::VectorConcat => "vector_concat".to_string(), + Self::VectorSlice => "subvector".to_string(), }; write!(f, "{str}") } @@ -841,6 +843,7 @@ impl Func { "vector_distance_cos" => Ok(Self::Vector(VectorFunc::VectorDistanceCos)), "vector_distance_l2" => Ok(Self::Vector(VectorFunc::VectorDistanceEuclidean)), "vector_concat" => Ok(Self::Vector(VectorFunc::VectorConcat)), + "subvector" => Ok(Self::Vector(VectorFunc::VectorSlice)), _ => crate::bail_parse_error!("no such function: {}", name), } } diff --git a/core/translate/expr.rs b/core/translate/expr.rs index 6de758324..2d7cdc26b 100644 --- a/core/translate/expr.rs +++ b/core/translate/expr.rs @@ -929,6 +929,16 @@ pub fn translate_expr( emit_function_call(program, func_ctx, &[regs, regs + 1], target_register)?; Ok(target_register) } + VectorFunc::VectorSlice => { + let args = expect_arguments_exact!(args, 3, vector_func); + let regs = program.alloc_registers(3); + translate_expr(program, referenced_tables, &args[0], regs, resolver)?; + translate_expr(program, referenced_tables, &args[1], regs + 1, resolver)?; + translate_expr(program, referenced_tables, &args[2], regs + 2, resolver)?; + + emit_function_call(program, func_ctx, &[regs, regs + 2], target_register)?; + Ok(target_register) + } }, Func::Scalar(srf) => { match srf { diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index c67320f09..515e6675e 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -16,7 +16,7 @@ use crate::types::{ use crate::util::normalize_ident; use crate::vdbe::insn::InsertFlags; use crate::vdbe::registers_to_ref_values; -use crate::vector::vector_concat; +use crate::vector::{subvector, vector_concat}; use crate::{ error::{ LimboError, SQLITE_CONSTRAINT, SQLITE_CONSTRAINT_NOTNULL, SQLITE_CONSTRAINT_PRIMARYKEY, @@ -4562,6 +4562,10 @@ pub fn op_function( let result = vector_concat(&state.registers[*start_reg..*start_reg + arg_count])?; state.registers[*dest] = Register::Value(result); } + VectorFunc::VectorSlice => { + let result = subvector(&state.registers[*start_reg..*start_reg + arg_count])?; + state.registers[*dest] = Register::Value(result) + } }, crate::function::Func::External(f) => match f.func { ExtFunc::Scalar(f) => { diff --git a/core/vector/mod.rs b/core/vector/mod.rs index 1742c0d96..0c3d427a3 100644 --- a/core/vector/mod.rs +++ b/core/vector/mod.rs @@ -108,7 +108,7 @@ pub fn vector_distance_l2(args: &[Register]) -> Result { pub fn vector_concat(args: &[Register]) -> Result { if args.len() != 2 { return Err(LimboError::ConversionError( - "distance_concat requires exactly two arguments".to_string(), + "concat requires exactly two arguments".to_string(), )); } @@ -127,3 +127,28 @@ pub fn vector_concat(args: &[Register]) -> Result { VectorType::Float64 => Ok(vector_serialize_f64(vector)), } } + +pub fn subvector(args: &[Register]) -> Result { + if args.len() != 3 { + return Err(LimboError::ConversionError( + "vector_sub requires exactly three arguments".to_string(), + )); + } + + let vector = parse_vector(&args[0], None)?; + let start_index = args[1].get_owned_value().as_int(); + let length = args[2].get_owned_value().as_int(); + + if start_index < 0 || length < 0 { + return Err(LimboError::InvalidArgument( + "start index or length can't be negative".into(), + )); + } + + let result = vector_types::subvector(&vector, start_index as usize, length as usize)?; + + match result.vector_type { + VectorType::Float32 => Ok(vector_serialize_f32(result)), + VectorType::Float64 => Ok(vector_serialize_f64(result)), + } +}