diff --git a/lib/api/src/rest/schema.rs b/lib/api/src/rest/schema.rs index cd119dedb2..a543a4de61 100644 --- a/lib/api/src/rest/schema.rs +++ b/lib/api/src/rest/schema.rs @@ -8,8 +8,8 @@ use segment::common::utils::MaybeOneOrMany; use segment::data_types::order_by::OrderBy; use segment::json_path::JsonPath; use segment::types::{ - Filter, IntPayloadType, Payload, PointIdType, SearchParams, ShardKey, VectorNameBuf, - WithPayloadInterface, WithVector, + Condition, Filter, GeoPoint, IntPayloadType, Payload, PointIdType, SearchParams, ShardKey, + VectorNameBuf, WithPayloadInterface, WithVector, }; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -505,6 +505,11 @@ pub struct FusionQuery { pub fusion: Fusion, } +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +pub struct FormulaQuery { + pub formula: FormulaInput, +} + #[derive(Debug, Serialize, Deserialize, JsonSchema)] #[serde(rename_all = "snake_case")] pub struct SampleQuery { @@ -623,6 +628,66 @@ impl ContextPair { } } +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +pub struct FormulaInput { + pub formula: Expression, + // TODO(score boosting): Validate defaults, particularly for score references + pub defaults: HashMap, +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +#[serde(untagged)] +pub enum Expression { + Constant(f32), + Variable(String), + Condition(Box), + Mult(MultExpression), + Sum(SumExpression), + Neg(NegExpression), + Div(DivExpression), + GeoDistance(GeoDistance), +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +pub struct MultExpression { + pub mult: Vec, +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +pub struct SumExpression { + pub sum: Vec, +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +pub struct NegExpression { + pub neg: Box, +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +pub struct DivExpression { + pub div: DivParams, +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +pub struct DivParams { + pub left: Box, + pub right: Box, + pub by_zero_default: ScoreType, +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +pub struct GeoDistance { + pub geo_distance: GeoDistanceParams, +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +pub struct GeoDistanceParams { + /// The origin geo point to measure from + pub origin: GeoPoint, + /// Payload field with the destination geo point + pub to: JsonPath, +} + #[derive(Debug, Serialize, Deserialize, JsonSchema)] #[serde(rename_all = "snake_case")] pub enum Sample { diff --git a/lib/collection/src/operations/universal_query/collection_query.rs b/lib/collection/src/operations/universal_query/collection_query.rs index dd253bd1a8..82cb406a69 100644 --- a/lib/collection/src/operations/universal_query/collection_query.rs +++ b/lib/collection/src/operations/universal_query/collection_query.rs @@ -91,6 +91,9 @@ pub enum Query { /// Order by a payload field OrderBy(OrderBy), + // TODO(score boosting): enable this + // /// Formula-based score fusion + // Formula(FormulaInternal), /// Sample points Sample(SampleInternal), } @@ -114,12 +117,15 @@ impl Query { } Query::Fusion(fusion) => ScoringQuery::Fusion(fusion), Query::OrderBy(order_by) => ScoringQuery::OrderBy(order_by), + // TODO(score boosting): enable this + // Query::Formula(formula) => ScoringQuery::Formula(ParsedFormula::try_from(formula)?), Query::Sample(sample) => ScoringQuery::Sample(sample), }; Ok(scoring_query) } } + #[derive(Clone, Debug, PartialEq)] pub enum VectorInputInternal { Id(PointIdType), diff --git a/lib/collection/src/operations/universal_query/formula.rs b/lib/collection/src/operations/universal_query/formula.rs new file mode 100644 index 0000000000..9c83124beb --- /dev/null +++ b/lib/collection/src/operations/universal_query/formula.rs @@ -0,0 +1,82 @@ +use std::collections::HashMap; + +use api::rest; +use api::rest::GeoDistance; +use common::types::ScoreType; +use segment::json_path::JsonPath; +use segment::types::{Condition, GeoPoint}; +use serde_json::Value; + +#[derive(Debug, Clone, PartialEq)] +pub struct FormulaInternal { + pub formula: ExpressionInternal, + pub defaults: HashMap, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ExpressionInternal { + Constant(f32), + Variable(String), + Condition(Box), + Mult(Vec), + Sum(Vec), + Neg(Box), + Div { + left: Box, + right: Box, + by_zero_default: ScoreType, + }, + GeoDistance { + origin: GeoPoint, + to: JsonPath, + }, +} + +impl From for FormulaInternal { + fn from(value: rest::FormulaInput) -> Self { + let rest::FormulaInput { formula, defaults } = value; + + FormulaInternal { + formula: ExpressionInternal::from(formula), + defaults, + } + } +} + +impl From for ExpressionInternal { + fn from(value: rest::Expression) -> Self { + match value { + rest::Expression::Constant(c) => ExpressionInternal::Constant(c), + rest::Expression::Variable(key) => ExpressionInternal::Variable(key), + rest::Expression::Condition(condition) => ExpressionInternal::Condition(condition), + rest::Expression::Mult(rest::MultExpression { mult: exprs }) => { + ExpressionInternal::Mult(exprs.into_iter().map(ExpressionInternal::from).collect()) + } + rest::Expression::Sum(rest::SumExpression { sum: exprs }) => { + ExpressionInternal::Sum(exprs.into_iter().map(ExpressionInternal::from).collect()) + } + rest::Expression::Neg(rest::NegExpression { neg: expr }) => { + ExpressionInternal::Neg(Box::new(ExpressionInternal::from(*expr))) + } + rest::Expression::Div(rest::DivExpression { + div: + rest::DivParams { + left, + right, + by_zero_default, + }, + }) => { + let left = Box::new((*left).into()); + let right = Box::new((*right).into()); + ExpressionInternal::Div { + left, + right, + by_zero_default, + } + } + rest::Expression::GeoDistance(GeoDistance { + geo_distance: rest::GeoDistanceParams { origin, to }, + }) => ExpressionInternal::GeoDistance { origin, to }, + } + } +} diff --git a/lib/collection/src/operations/universal_query/mod.rs b/lib/collection/src/operations/universal_query/mod.rs index de4a40f519..09e0d86f77 100644 --- a/lib/collection/src/operations/universal_query/mod.rs +++ b/lib/collection/src/operations/universal_query/mod.rs @@ -21,5 +21,6 @@ //! [`QueryShardPoints`]: api::grpc::qdrant::QueryShardPoints pub mod collection_query; +pub mod formula; pub mod planned_query; pub mod shard_query; diff --git a/lib/collection/src/operations/universal_query/shard_query.rs b/lib/collection/src/operations/universal_query/shard_query.rs index 298732fc46..1dc4f6a139 100644 --- a/lib/collection/src/operations/universal_query/shard_query.rs +++ b/lib/collection/src/operations/universal_query/shard_query.rs @@ -1,3 +1,5 @@ +use std::collections::HashSet; + use api::grpc::qdrant as grpc; use common::types::ScoreType; use itertools::Itertools; @@ -5,16 +7,21 @@ use segment::data_types::order_by::OrderBy; use segment::data_types::vectors::{ NamedQuery, NamedVectorStruct, VectorInternal, DEFAULT_VECTOR_NAME, }; +use segment::index::query_optimization::rescore_formula::parsed_formula::{ + ParsedExpression, ParsedFormula, VariableId, +}; +use segment::json_path::JsonPath; use segment::types::{ - Filter, Order, ScoredPoint, SearchParams, VectorName, VectorNameBuf, WithPayloadInterface, - WithVector, + Condition, Filter, Order, ScoredPoint, SearchParams, VectorName, VectorNameBuf, + WithPayloadInterface, WithVector, }; use segment::vector_storage::query::{ContextQuery, DiscoveryQuery, RecoQuery}; use tonic::Status; use crate::config::CollectionParams; use crate::operations::query_enum::QueryEnum; -use crate::operations::types::CollectionResult; +use crate::operations::types::{CollectionError, CollectionResult}; +use crate::operations::universal_query::formula::{ExpressionInternal, FormulaInternal}; /// Internal response type for a universal query request. /// @@ -24,6 +31,8 @@ pub type ShardQueryResponse = Vec>; /// Internal representation of a universal query request. /// /// Direct translation of the user-facing request, but with all point ids substituted with their corresponding vectors. +/// +/// For the case of formula queries, it collects conditions and variables too. #[derive(Clone, Debug)] pub struct ShardQueryRequest { pub prefetches: Vec, @@ -73,6 +82,8 @@ pub enum ScoringQuery { /// Order by a payload field OrderBy(OrderBy), + // TODO(score boosting): Enable this + // Formula(ParsedFormula), /// Sample points Sample(SampleInternal), } @@ -177,6 +188,86 @@ impl ShardPrefetch { } } +impl ExpressionInternal { + fn parse_and_convert( + self, + payload_vars: &mut HashSet, + conditions: &mut Vec, + ) -> CollectionResult { + let expr = match self { + ExpressionInternal::Constant(c) => ParsedExpression::Constant(c), + ExpressionInternal::Variable(var) => { + let var: VariableId = var.parse()?; + if let VariableId::Payload(payload_var) = var.clone() { + payload_vars.insert(payload_var); + } + ParsedExpression::Variable(var) + } + ExpressionInternal::Condition(condition) => { + let condition_id = conditions.len(); + conditions.push(*condition); + ParsedExpression::new_condition_id(condition_id) + } + ExpressionInternal::Mult(internal_expressions) => ParsedExpression::Mult( + internal_expressions + .into_iter() + .map(|expr| expr.parse_and_convert(payload_vars, conditions)) + .try_collect()?, + ), + ExpressionInternal::Sum(expression_internals) => ParsedExpression::Sum( + expression_internals + .into_iter() + .map(|expr| expr.parse_and_convert(payload_vars, conditions)) + .try_collect()?, + ), + ExpressionInternal::Neg(expression_internal) => ParsedExpression::new_neg( + expression_internal.parse_and_convert(payload_vars, conditions)?, + ), + ExpressionInternal::Div { + left, + right, + by_zero_default, + } => ParsedExpression::new_div( + left.parse_and_convert(payload_vars, conditions)?, + right.parse_and_convert(payload_vars, conditions)?, + by_zero_default, + ), + ExpressionInternal::GeoDistance { origin, to } => { + ParsedExpression::new_geo_distance(origin, to) + } + }; + + Ok(expr) + } +} + +impl TryFrom for ParsedFormula { + type Error = CollectionError; + fn try_from(value: FormulaInternal) -> Result { + let FormulaInternal { formula, defaults } = value; + + let mut payload_vars = HashSet::new(); + let mut conditions = Vec::new(); + + let parsed_expression = formula.parse_and_convert(&mut payload_vars, &mut conditions)?; + + let defaults = defaults + .into_iter() + .map(|(key, value)| { + let key = key.as_str().parse()?; + CollectionResult::Ok((key, value)) + }) + .try_collect()?; + + Ok(ParsedFormula { + formula: parsed_expression, + payload_vars, + conditions, + defaults, + }) + } +} + impl TryFrom for ShardQueryRequest { type Error = Status; diff --git a/lib/segment/src/index/query_optimization/rescore_formula/formula_scorer.rs b/lib/segment/src/index/query_optimization/rescore_formula/formula_scorer.rs index c66206c493..93132b6fb5 100644 --- a/lib/segment/src/index/query_optimization/rescore_formula/formula_scorer.rs +++ b/lib/segment/src/index/query_optimization/rescore_formula/formula_scorer.rs @@ -7,7 +7,7 @@ use common::types::{PointOffsetType, ScoreType}; use geo::{Distance, Haversine}; use serde_json::Value; -use super::parsed_formula::{Expression, ParsedFormula, VariableId}; +use super::parsed_formula::{ParsedExpression, ParsedFormula, VariableId}; use super::value_retriever::VariableRetrieverFn; use crate::common::operation_error::{OperationError, OperationResult}; use crate::index::query_optimization::optimized_filter::{check_condition, OptimizedCondition}; @@ -21,7 +21,7 @@ const DEFAULT_SCORE: ScoreType = 0.0; /// A scorer to evaluate the same formula for many points pub struct FormulaScorer<'a> { /// The formula to evaluate - formula: Expression, + formula: ParsedExpression, /// One hashmap for each prefetch results prefetches_scores: &'a [AHashMap], /// Payload key -> retriever function @@ -78,12 +78,12 @@ impl FormulaScorer<'_> { /// Evaluate the expression recursively fn eval_expression( &self, - expression: &Expression, + expression: &ParsedExpression, point_id: PointOffsetType, ) -> OperationResult { match expression { - Expression::Constant(c) => Ok(*c), - Expression::Variable(v) => match v { + ParsedExpression::Constant(c) => Ok(*c), + ParsedExpression::Variable(v) => match v { VariableId::Score(prefetch_idx) => Ok(self .prefetches_scores .get(*prefetch_idx) @@ -115,7 +115,7 @@ impl FormulaScorer<'_> { Ok(score) } }, - Expression::Mult(expressions) => { + ParsedExpression::Mult(expressions) => { let mut product = 1.0; for expr in expressions { let value = self.eval_expression(expr, point_id)?; @@ -127,11 +127,11 @@ impl FormulaScorer<'_> { } Ok(product) } - Expression::Sum(expressions) => expressions.iter().try_fold(0.0, |acc, expr| { + ParsedExpression::Sum(expressions) => expressions.iter().try_fold(0.0, |acc, expr| { let value = self.eval_expression(expr, point_id)?; Ok(acc + value) }), - Expression::Div { + ParsedExpression::Div { left, right, by_zero_default, @@ -149,11 +149,11 @@ impl FormulaScorer<'_> { Ok(left / right) } } - Expression::Neg(expr) => { + ParsedExpression::Neg(expr) => { let value = self.eval_expression(expr, point_id)?; Ok(value.neg()) } - Expression::GeoDistance { origin, key } => { + ParsedExpression::GeoDistance { origin, key } => { let value: GeoPoint = self .payload_retrievers .get(key) @@ -220,7 +220,7 @@ mod tests { ]; FormulaScorer { - formula: Expression::Constant(0.0), + formula: ParsedExpression::Constant(0.0), prefetches_scores, payload_retrievers, condition_checkers, @@ -231,38 +231,38 @@ mod tests { #[rstest] // Basic expressions, just variables - #[case(Expression::Constant(5.0), 5.0)] - #[case(Expression::new_score_id(0), 1.0)] - #[case(Expression::new_score_id(1), 2.0)] - #[case(Expression::new_payload_id(FIELD_NAME), 85.0)] - #[case(Expression::new_condition_id(0), 1.0)] - #[case(Expression::new_condition_id(1), 0.0)] + #[case(ParsedExpression::Constant(5.0), 5.0)] + #[case(ParsedExpression::new_score_id(0), 1.0)] + #[case(ParsedExpression::new_score_id(1), 2.0)] + #[case(ParsedExpression::new_payload_id(JsonPath::new(FIELD_NAME)), 85.0)] + #[case(ParsedExpression::new_condition_id(0), 1.0)] + #[case(ParsedExpression::new_condition_id(1), 0.0)] // Operations - #[case(Expression::Sum(vec![ - Expression::Constant(1.0), - Expression::new_score_id(0), - Expression::new_payload_id(FIELD_NAME), - Expression::new_condition_id(0), + #[case(ParsedExpression::Sum(vec![ + ParsedExpression::Constant(1.0), + ParsedExpression::new_score_id(0), + ParsedExpression::new_payload_id(JsonPath::new(FIELD_NAME)), + ParsedExpression::new_condition_id(0), ]), 1.0 + 1.0 + 85.0 + 1.0)] - #[case(Expression::Mult(vec![ - Expression::Constant(2.0), - Expression::new_score_id(0), - Expression::new_payload_id(FIELD_NAME), - Expression::new_condition_id(0), + #[case(ParsedExpression::Mult(vec![ + ParsedExpression::Constant(2.0), + ParsedExpression::new_score_id(0), + ParsedExpression::new_payload_id(JsonPath::new(FIELD_NAME)), + ParsedExpression::new_condition_id(0), ]), 2.0 * 1.0 * 85.0 * 1.0)] - #[case(Expression::Div { - left: Box::new(Expression::Constant(10.0)), - right: Box::new(Expression::new_score_id(0)), + #[case(ParsedExpression::Div { + left: Box::new(ParsedExpression::Constant(10.0)), + right: Box::new(ParsedExpression::new_score_id(0)), by_zero_default: f32::INFINITY, }, 10.0 / 1.0)] - #[case(Expression::new_neg(Expression::Constant(10.0)), -10.0)] - #[case(Expression::new_geo_distance(GeoPoint { lat: 25.717877679163667, lon: -100.43383200156751 }, JsonPath::new(GEO_FIELD_NAME)), 21926.494)] + #[case(ParsedExpression::new_neg(ParsedExpression::Constant(10.0)), -10.0)] + #[case(ParsedExpression::new_geo_distance(GeoPoint { lat: 25.717877679163667, lon: -100.43383200156751 }, JsonPath::new(GEO_FIELD_NAME)), 21926.494)] #[should_panic( expected = r#"called `Result::unwrap()` on an `Err` value: VariableTypeError { field_name: JsonPath { first_key: "number", rest: [] }, expected_type: "geo point" }"# )] - #[case(Expression::new_geo_distance(GeoPoint { lat: 25.717877679163667, lon: -100.43383200156751 }, JsonPath::new(FIELD_NAME)), 0.0)] + #[case(ParsedExpression::new_geo_distance(GeoPoint { lat: 25.717877679163667, lon: -100.43383200156751 }, JsonPath::new(FIELD_NAME)), 0.0)] #[test] - fn test_evaluation(#[case] expr: Expression, #[case] expected: ScoreType) { + fn test_evaluation(#[case] expr: ParsedExpression, #[case] expected: ScoreType) { let defaults = HashMap::new(); let scorer_fixture = make_formula_scorer(&defaults); @@ -274,17 +274,23 @@ mod tests { // Default values #[rstest] // Defined default score - #[case(Expression::new_score_id(3), 1.5)] + #[case(ParsedExpression::new_score_id(3), 1.5)] // score idx not defined - #[case(Expression::new_score_id(10), DEFAULT_SCORE)] + #[case(ParsedExpression::new_score_id(10), DEFAULT_SCORE)] // missing value in payload - #[case(Expression::new_payload_id(NO_VALUE_FIELD_NAME), 85.0)] + #[case( + ParsedExpression::new_payload_id(JsonPath::new(NO_VALUE_FIELD_NAME)), + 85.0 + )] // missing value and no default value provided - #[case(Expression::new_payload_id("missing_field"), DEFAULT_SCORE)] + #[case( + ParsedExpression::new_payload_id(JsonPath::new("missing_field")), + DEFAULT_SCORE + )] // geo distance with default value - #[case(Expression::new_geo_distance(GeoPoint { lat: 25.717877679163667, lon: -100.43383200156751 }, JsonPath::new(NO_VALUE_GEO_POINT)), 90951.3)] + #[case(ParsedExpression::new_geo_distance(GeoPoint { lat: 25.717877679163667, lon: -100.43383200156751 }, JsonPath::new(NO_VALUE_GEO_POINT)), 90951.3)] #[test] - fn test_default_values(#[case] expr: Expression, #[case] expected: ScoreType) { + fn test_default_values(#[case] expr: ParsedExpression, #[case] expected: ScoreType) { let defaults = [ (VariableId::Score(3), json!(1.5)), ( diff --git a/lib/segment/src/index/query_optimization/rescore_formula/parsed_formula.rs b/lib/segment/src/index/query_optimization/rescore_formula/parsed_formula.rs index 7d97a01c48..ac979d0164 100644 --- a/lib/segment/src/index/query_optimization/rescore_formula/parsed_formula.rs +++ b/lib/segment/src/index/query_optimization/rescore_formula/parsed_formula.rs @@ -1,49 +1,53 @@ use std::collections::{HashMap, HashSet}; +use std::str::FromStr; use common::types::ScoreType; use serde_json::Value; -use crate::json_path::JsonPath; +use crate::json_path::{JsonPath, JsonPathItem}; use crate::types::{Condition, GeoPoint}; +const SCORE_KEYWORD: &str = "score"; + pub type ConditionId = usize; +#[derive(Debug, Clone, PartialEq)] pub struct ParsedFormula { /// Variables used in the formula - pub(super) payload_vars: HashSet, + pub payload_vars: HashSet, /// Conditions used in the formula. Their index in the array is used as a variable id - pub(super) conditions: Vec, + pub conditions: Vec, /// Defaults to use when variable is not found - pub(super) defaults: HashMap, + pub defaults: HashMap, /// Root of the formula expression - pub(super) formula: Expression, + pub formula: ParsedExpression, } -#[derive(Clone)] -pub enum Expression { +#[derive(Debug, Clone, PartialEq)] +pub enum ParsedExpression { // Scalars Constant(ScoreType), Variable(VariableId), // Operations - Mult(Vec), - Sum(Vec), + Mult(Vec), + Sum(Vec), Div { - left: Box, - right: Box, + left: Box, + right: Box, by_zero_default: ScoreType, }, - Neg(Box), + Neg(Box), GeoDistance { origin: GeoPoint, key: JsonPath, }, } -#[derive(Clone, Hash, Eq, PartialEq)] +#[derive(Debug, Clone, Hash, Eq, PartialEq)] pub enum VariableId { /// Score index Score(usize), @@ -53,33 +57,118 @@ pub enum VariableId { Condition(ConditionId), } -impl Expression { - pub fn new_div(left: Expression, right: Expression, by_zero_default: ScoreType) -> Self { - Expression::Div { +impl ParsedExpression { + pub fn new_div( + left: ParsedExpression, + right: ParsedExpression, + by_zero_default: ScoreType, + ) -> Self { + ParsedExpression::Div { left: Box::new(left), right: Box::new(right), by_zero_default, } } - pub fn new_neg(expression: Expression) -> Self { - Expression::Neg(Box::new(expression)) + pub fn new_neg(expression: ParsedExpression) -> Self { + ParsedExpression::Neg(Box::new(expression)) } pub fn new_geo_distance(origin: GeoPoint, key: JsonPath) -> Self { - Expression::GeoDistance { origin, key } + ParsedExpression::GeoDistance { origin, key } } - #[cfg(feature = "testing")] - pub fn new_payload_id(path: &str) -> Self { - Expression::Variable(VariableId::Payload(JsonPath::new(path))) + pub fn new_payload_id(path: JsonPath) -> Self { + ParsedExpression::Variable(VariableId::Payload(path)) } pub fn new_score_id(index: usize) -> Self { - Expression::Variable(VariableId::Score(index)) + ParsedExpression::Variable(VariableId::Score(index)) } pub fn new_condition_id(index: ConditionId) -> Self { - Expression::Variable(VariableId::Condition(index)) + ParsedExpression::Variable(VariableId::Condition(index)) + } +} + +impl FromStr for VariableId { + type Err = String; + + fn from_str(var_str: &str) -> Result { + let var_id = match var_str.strip_prefix("$") { + Some(score) => { + // parse as reserved word + let json_path = score + .parse::() + .map_err(|_| format!("Invalid reserved variable: {var_str}"))?; + match json_path.first_key.as_str() { + SCORE_KEYWORD => match &json_path.rest[..] { + // Default prefetch index, like "$score" + [] => VariableId::Score(0), + // Specifies prefetch index, like "$score[2]" + [JsonPathItem::Index(idx)] => VariableId::Score(*idx), + _ => { + // Only direct index is supported + return Err(format!("Invalid reserved variable: {var_str}")); + } + }, + _ => { + // No other reserved words are supported + return Err(format!("Invalid reserved word: {var_str}")); + } + } + } + None => { + // parse as regular payload variable + let parsed = var_str + .parse() + .map_err(|_| format!("Invalid payload variable: {var_str}"))?; + VariableId::Payload(parsed) + } + }; + Ok(var_id) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_variable_id_from_str() { + // Test score variables + assert_eq!( + VariableId::from_str("$score").unwrap(), + VariableId::Score(0) + ); + assert_eq!( + VariableId::from_str("$score[0]").unwrap(), + VariableId::Score(0) + ); + assert_eq!( + VariableId::from_str("$score[1]").unwrap(), + VariableId::Score(1) + ); + assert!(VariableId::from_str("$score.invalid").is_err()); + assert!(VariableId::from_str("$score[1][2]").is_err()); + assert!(VariableId::from_str("$score[]").is_err()); + + // Test invalid reserved words + assert!(VariableId::from_str("$invalid").is_err()); + + // Test payload variables + assert_eq!( + VariableId::from_str("field").unwrap(), + VariableId::Payload("field".parse().unwrap()) + ); + assert_eq!( + VariableId::from_str("field.nested").unwrap(), + VariableId::Payload("field.nested".parse().unwrap()) + ); + assert_eq!( + VariableId::from_str("field[0]").unwrap(), + VariableId::Payload("field[0]".parse().unwrap()) + ); + assert!(VariableId::from_str("").is_err()); } }