[score boosting] create query structures and parse them (#5990)

* create REST, CollectionQuery, and ShardQuery structures. Parse all variables and conditions

* move stuff around, rename segment Expression to ParsedExpression

* rename to MultExpression

* add Div in REST, use unique geo_distance tag

* review fixes

---------

Co-authored-by: generall <andrey@vasnetsov.com>

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<String, Value>,
+}
+
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+#[serde(untagged)]
+pub enum Expression {
+    Constant(f32),
+    Variable(String),
+    Condition(Box<Condition>),
+    Mult(MultExpression),
+    Sum(SumExpression),
+    Neg(NegExpression),
+    Div(DivExpression),
+    GeoDistance(GeoDistance),
+}
+
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+pub struct MultExpression {
+    pub mult: Vec<Expression>,
+}
+
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+pub struct SumExpression {
+    pub sum: Vec<Expression>,
+}
+
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+pub struct NegExpression {
+    pub neg: Box<Expression>,
+}
+
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+pub struct DivExpression {
+    pub div: DivParams,
+}
+
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+pub struct DivParams {
+    pub left: Box<Expression>,
+    pub right: Box<Expression>,
+    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 {

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),

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<String, Value>,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum ExpressionInternal {
+    Constant(f32),
+    Variable(String),
+    Condition(Box<Condition>),
+    Mult(Vec<ExpressionInternal>),
+    Sum(Vec<ExpressionInternal>),
+    Neg(Box<ExpressionInternal>),
+    Div {
+        left: Box<ExpressionInternal>,
+        right: Box<ExpressionInternal>,
+        by_zero_default: ScoreType,
+    },
+    GeoDistance {
+        origin: GeoPoint,
+        to: JsonPath,
+    },
+}
+
+impl From<rest::FormulaInput> for FormulaInternal {
+    fn from(value: rest::FormulaInput) -> Self {
+        let rest::FormulaInput { formula, defaults } = value;
+
+        FormulaInternal {
+            formula: ExpressionInternal::from(formula),
+            defaults,
+        }
+    }
+}
+
+impl From<rest::Expression> 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 },
+        }
+    }
+}

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;

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<Vec<ScoredPoint>>;
 /// 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<ShardPrefetch>,
@@ -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<JsonPath>,
+        conditions: &mut Vec<Condition>,
+    ) -> CollectionResult<ParsedExpression> {
+        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<FormulaInternal> for ParsedFormula {
+    type Error = CollectionError;
+    fn try_from(value: FormulaInternal) -> Result<Self, Self::Error> {
+        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<grpc::QueryShardPoints> for ShardQueryRequest {
     type Error = Status;
 

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<PointOffsetType, ScoreType>],
     /// 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<ScoreType> {
         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)),
             (

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<JsonPath>,
+    pub payload_vars: HashSet<JsonPath>,
 
     /// Conditions used in the formula. Their index in the array is used as a variable id
-    pub(super) conditions: Vec<Condition>,
+    pub conditions: Vec<Condition>,
 
     /// Defaults to use when variable is not found
-    pub(super) defaults: HashMap<VariableId, Value>,
+    pub defaults: HashMap<VariableId, Value>,
 
     /// 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<Expression>),
-    Sum(Vec<Expression>),
+    Mult(Vec<ParsedExpression>),
+    Sum(Vec<ParsedExpression>),
     Div {
-        left: Box<Expression>,
-        right: Box<Expression>,
+        left: Box<ParsedExpression>,
+        right: Box<ParsedExpression>,
         by_zero_default: ScoreType,
     },
-    Neg(Box<Expression>),
+    Neg(Box<ParsedExpression>),
     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<Self, Self::Err> {
+        let var_id = match var_str.strip_prefix("$") {
+            Some(score) => {
+                // parse as reserved word
+                let json_path = score
+                    .parse::<JsonPath>()
+                    .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());
     }
 }