use std::collections::{HashMap, HashSet};
use std::fmt;
use crate::common::types::ScoreType;
use itertools::Itertools;
use crate::segment::common::operation_error::{OperationError, OperationResult};
use crate::segment::index::query_optimization::rescore_formula::parsed_formula::*;
use crate::segment::json_path::JsonPath;
use crate::segment::types::{Condition, GeoPoint};
use serde::Serialize;
use serde_json::Value;
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct FormulaInternal {
pub formula: ExpressionInternal,
pub defaults: HashMap<String, Value>,
}
impl TryFrom<FormulaInternal> for ParsedFormula {
type Error = OperationError;
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()
.map_err(|msg| failed_to_parse("variable ID", &key, &msg))?;
OperationResult::Ok((key, value))
})
.try_collect()?;
Ok(ParsedFormula {
formula: parsed_expression,
payload_vars,
conditions,
defaults,
})
}
}
#[derive(Clone, Debug, PartialEq, Serialize)]
pub enum ExpressionInternal {
Constant(f32),
Variable(String),
Condition(Box<Condition>),
GeoDistance {
origin: GeoPoint,
to: JsonPath,
},
Datetime(String),
DatetimeKey(JsonPath),
Mult(Vec<ExpressionInternal>),
Sum(Vec<ExpressionInternal>),
Neg(Box<ExpressionInternal>),
Div {
left: Box<ExpressionInternal>,
right: Box<ExpressionInternal>,
by_zero_default: Option<ScoreType>,
},
Sqrt(Box<ExpressionInternal>),
Pow {
base: Box<ExpressionInternal>,
exponent: Box<ExpressionInternal>,
},
Exp(Box<ExpressionInternal>),
Log10(Box<ExpressionInternal>),
Ln(Box<ExpressionInternal>),
Abs(Box<ExpressionInternal>),
Decay {
kind: DecayKind,
x: Box<ExpressionInternal>,
target: Option<Box<ExpressionInternal>>,
midpoint: Option<f32>,
scale: Option<f32>,
},
}
impl ExpressionInternal {
fn parse_and_convert(
self,
payload_vars: &mut HashSet<JsonPath>,
conditions: &mut Vec<Condition>,
) -> OperationResult<ParsedExpression> {
let expr = match self {
ExpressionInternal::Constant(c) => {
ParsedExpression::Constant(PreciseScoreOrdered::from(PreciseScore::from(c)))
}
ExpressionInternal::Variable(var) => {
let var: VariableId = var
.parse()
.map_err(|msg| failed_to_parse("variable ID", &var, &msg))?;
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::GeoDistance { origin, to } => {
payload_vars.insert(to.clone());
ParsedExpression::new_geo_distance(origin, to)
}
ExpressionInternal::Datetime(dt_str) => {
ParsedExpression::Datetime(DatetimeExpression::Constant(
dt_str
.parse()
.map_err(|err| failed_to_parse("date-time", &dt_str, err))?,
))
}
ExpressionInternal::DatetimeKey(json_path) => {
payload_vars.insert(json_path.clone());
ParsedExpression::Datetime(DatetimeExpression::PayloadVariable(json_path))
}
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.map(PreciseScore::from),
),
ExpressionInternal::Sqrt(expression_internal) => ParsedExpression::Sqrt(Box::new(
expression_internal.parse_and_convert(payload_vars, conditions)?,
)),
ExpressionInternal::Pow { base, exponent } => ParsedExpression::Pow {
base: Box::new(base.parse_and_convert(payload_vars, conditions)?),
exponent: Box::new(exponent.parse_and_convert(payload_vars, conditions)?),
},
ExpressionInternal::Exp(expression_internal) => ParsedExpression::Exp(Box::new(
expression_internal.parse_and_convert(payload_vars, conditions)?,
)),
ExpressionInternal::Log10(expression_internal) => ParsedExpression::Log10(Box::new(
expression_internal.parse_and_convert(payload_vars, conditions)?,
)),
ExpressionInternal::Ln(expression_internal) => ParsedExpression::Ln(Box::new(
expression_internal.parse_and_convert(payload_vars, conditions)?,
)),
ExpressionInternal::Abs(expression_internal) => ParsedExpression::Abs(Box::new(
expression_internal.parse_and_convert(payload_vars, conditions)?,
)),
ExpressionInternal::Decay {
kind,
x,
target,
midpoint,
scale,
} => {
let lambda = ParsedExpression::decay_params_to_lambda(midpoint, scale, kind)?;
let x = x.parse_and_convert(payload_vars, conditions)?;
let target = target
.map(|t| t.parse_and_convert(payload_vars, conditions))
.transpose()?
.map(Box::new);
ParsedExpression::Decay {
kind,
x: Box::new(x),
target,
lambda: PreciseScoreOrdered::from(lambda),
}
}
};
Ok(expr)
}
}
fn failed_to_parse(what: &str, value: &str, message: impl fmt::Display) -> OperationError {
OperationError::validation_error(format!("failed to parse {what} {value}: {message}"))
}