diff --git a/lib/api/src/grpc/conversions.rs b/lib/api/src/grpc/conversions.rs index 6fc7457fcf..16804e5c22 100644 --- a/lib/api/src/grpc/conversions.rs +++ b/lib/api/src/grpc/conversions.rs @@ -11,6 +11,9 @@ use segment::data_types::index::{ KeywordIndexType, TextIndexType, UuidIndexType, }; use segment::data_types::{facets as segment_facets, vectors as segment_vectors}; +use segment::index::query_optimization::rescore_formula::parsed_formula::{ + ParsedExpression, ParsedFormula, +}; use segment::types::{DateTimePayloadType, FloatPayloadType, default_quantization_ignore_value}; use segment::vector_storage::query as segment_query; use sparse::common::sparse_vector::validate_sparse_vector_impl; @@ -27,7 +30,8 @@ use super::qdrant::{ StrictModeMultivector, StrictModeMultivectorConfig, StrictModeSparse, StrictModeSparseConfig, UuidIndexParams, VectorsOutput, WithLookup, raw_query, start_from, }; -use crate::conversions::json; +use super::{Expression, Formula}; +use crate::conversions::json::{self, json_to_proto}; use crate::grpc::qdrant::condition::ConditionOneOf; use crate::grpc::qdrant::r#match::MatchValue; use crate::grpc::qdrant::payload_index_params::IndexParams; @@ -45,6 +49,7 @@ use crate::grpc::qdrant::{ TextIndexParams, TokenizerType, UpdateResult, UpdateResultInternal, ValuesCount, VectorsSelector, WithPayloadSelector, WithVectorsSelector, shard_key, with_vectors_selector, }; +use crate::grpc::{DivExpression, GeoDistance, MultExpression, PowExpression, SumExpression}; use crate::rest::models::{CollectionsResponse, VersionInfo}; use crate::rest::schema as rest; @@ -2587,3 +2592,100 @@ impl From for HardwareUsage { } } } + +impl Formula { + /// This implementation is only used to forward a request to remote shards. + /// + /// It is preferred to pay the cost of un-parsing->re-parsing the formula, and keep the parsed representation + /// out of the API surface, than to expose the implementation details to the interface and avoid the extra work. + /// Conversion should be cheap enough. + pub fn from_parsed(value: ParsedFormula) -> Self { + let ParsedFormula { + formula, + payload_vars: _, // they are already in the expression + conditions, + defaults, + } = value; + + let expression = unparse_expression(formula, &conditions); + + let defaults = defaults + .into_iter() + .map(|(key, value)| (key.unparse(), json_to_proto(value))) + .collect(); + + Formula { + expression: Some(expression), + defaults, + } + } +} + +fn unparse_expression( + formula: ParsedExpression, + conditions: &Vec, +) -> Expression { + use segment::index::query_optimization::rescore_formula::parsed_formula::VariableId; + + use super::expression::Variant; + + let variant = match formula { + ParsedExpression::Constant(c) => Variant::Constant(c), + ParsedExpression::Variable(variable_id) => match variable_id { + var_id @ VariableId::Score(_) => Variant::Variable(var_id.unparse()), + var_id @ VariableId::Payload(_) => Variant::Variable(var_id.unparse()), + VariableId::Condition(cond_idx) => { + Variant::Condition(Condition::from(conditions[cond_idx].clone())) + } + }, + ParsedExpression::Mult(exprs) => Variant::Mult(MultExpression { + mult: exprs + .into_iter() + .map(|expr| unparse_expression(expr, conditions)) + .collect(), + }), + ParsedExpression::Sum(exprs) => Variant::Sum(SumExpression { + sum: exprs + .into_iter() + .map(|expr| unparse_expression(expr, conditions)) + .collect(), + }), + ParsedExpression::Neg(expr) => { + Variant::Neg(Box::new(unparse_expression(*expr, conditions))) + } + ParsedExpression::Div { + left, + right, + by_zero_default, + } => Variant::Div(Box::new(DivExpression { + left: Some(Box::new(unparse_expression(*left, conditions))), + right: Some(Box::new(unparse_expression(*right, conditions))), + by_zero_default: Some(by_zero_default), + })), + ParsedExpression::Sqrt(expr) => { + Variant::Sqrt(Box::new(unparse_expression(*expr, conditions))) + } + ParsedExpression::Pow { base, exponent } => Variant::Pow(Box::new(PowExpression { + base: Some(Box::new(unparse_expression(*base, conditions))), + exponent: Some(Box::new(unparse_expression(*exponent, conditions))), + })), + ParsedExpression::Exp(expr) => { + Variant::Exp(Box::new(unparse_expression(*expr, conditions))) + } + ParsedExpression::Log10(expr) => { + Variant::Log10(Box::new(unparse_expression(*expr, conditions))) + } + ParsedExpression::Ln(expr) => Variant::Ln(Box::new(unparse_expression(*expr, conditions))), + ParsedExpression::Abs(expr) => { + Variant::Abs(Box::new(unparse_expression(*expr, conditions))) + } + ParsedExpression::GeoDistance { origin, key } => Variant::GeoDistance(GeoDistance { + origin: Some(GeoPoint::from(origin)), + to: key.to_string(), + }), + }; + + Expression { + variant: Some(variant), + } +} diff --git a/lib/api/src/grpc/mod.rs b/lib/api/src/grpc/mod.rs index 78ab4304c7..beba7751ed 100644 --- a/lib/api/src/grpc/mod.rs +++ b/lib/api/src/grpc/mod.rs @@ -10,6 +10,8 @@ pub mod grpc_health_v1; pub mod transport_channel_pool; pub mod validate; +pub use qdrant::*; + pub const fn api_crate_version() -> &'static str { env!("CARGO_PKG_VERSION") } diff --git a/lib/api/src/grpc/proto/points_internal_service.proto b/lib/api/src/grpc/proto/points_internal_service.proto index dc806282f9..f8f5bcd5fc 100644 --- a/lib/api/src/grpc/proto/points_internal_service.proto +++ b/lib/api/src/grpc/proto/points_internal_service.proto @@ -244,6 +244,7 @@ message QueryShardPoints { Fusion fusion = 2; // One of the fusion methods OrderBy order_by = 3; // Order by a field Sample sample = 4; // Sample points + Formula formula = 5; // Use an arbitrary formula to rescore points } } diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 6e183a0d05..37740e1bc5 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -9634,7 +9634,7 @@ pub mod query_shard_points { #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct Query { - #[prost(oneof = "query::Score", tags = "1, 2, 3, 4")] + #[prost(oneof = "query::Score", tags = "1, 2, 3, 4, 5")] pub score: ::core::option::Option, } /// Nested message and enum types in `Query`. @@ -9655,6 +9655,9 @@ pub mod query_shard_points { /// Sample points #[prost(enumeration = "super::super::Sample", tag = "4")] Sample(i32), + /// Use an arbitrary formula to rescore points + #[prost(message, tag = "5")] + Formula(super::super::Formula), } } #[derive(serde::Serialize)] diff --git a/lib/collection/src/operations/universal_query/formula.rs b/lib/collection/src/operations/universal_query/formula.rs index e0e4740c2c..4705014cab 100644 --- a/lib/collection/src/operations/universal_query/formula.rs +++ b/lib/collection/src/operations/universal_query/formula.rs @@ -1,12 +1,18 @@ -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use api::rest; use api::rest::GeoDistance; use common::types::ScoreType; +use itertools::Itertools; +use segment::index::query_optimization::rescore_formula::parsed_formula::{ + ParsedExpression, ParsedFormula, VariableId, +}; use segment::json_path::JsonPath; use segment::types::{Condition, GeoPoint}; use serde_json::Value; +use crate::operations::types::{CollectionError, CollectionResult}; + #[derive(Debug, Clone, PartialEq)] pub struct FormulaInternal { pub formula: ExpressionInternal, @@ -41,6 +47,106 @@ pub enum ExpressionInternal { }, } +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) + } + 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)?, + )), + }; + + 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 From for FormulaInternal { fn from(value: rest::FormulaQuery) -> Self { let rest::FormulaQuery { formula, defaults } = value; diff --git a/lib/collection/src/operations/universal_query/shard_query.rs b/lib/collection/src/operations/universal_query/shard_query.rs index 05c05228e7..3c09866e5a 100644 --- a/lib/collection/src/operations/universal_query/shard_query.rs +++ b/lib/collection/src/operations/universal_query/shard_query.rs @@ -1,5 +1,3 @@ -use std::collections::HashSet; - use api::conversions::json::proto_to_json; use api::grpc::conversions::grpc_condition_into_condition; use api::grpc::qdrant as grpc; @@ -9,22 +7,18 @@ use segment::data_types::order_by::OrderBy; use segment::data_types::vectors::{ DEFAULT_VECTOR_NAME, NamedQuery, NamedVectorStruct, VectorInternal, }; -use segment::index::query_optimization::rescore_formula::parsed_formula::{ - ParsedExpression, ParsedFormula, VariableId, -}; -use segment::json_path::JsonPath; +use segment::index::query_optimization::rescore_formula::parsed_formula::ParsedFormula; use segment::types::{ - Condition, Filter, Order, ScoredPoint, SearchParams, VectorName, VectorNameBuf, - WithPayloadInterface, WithVector, + 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::{CollectionError, CollectionResult}; +use crate::operations::types::CollectionResult; use crate::operations::universal_query::formula::{ExpressionInternal, FormulaInternal}; - /// Internal response type for a universal query request. /// /// Capable of returning multiple intermediate results if needed, like the case of RRF (Reciprocal Rank Fusion) @@ -199,105 +193,6 @@ 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) - } - 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)?, - )), - }; - - 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; @@ -621,6 +516,11 @@ impl ScoringQuery { grpc::query_shard_points::query::Score::Sample(sample) => { ScoringQuery::Sample(SampleInternal::try_from(sample)?) } + grpc::query_shard_points::query::Score::Formula(formula) => ScoringQuery::Formula( + ParsedFormula::try_from(FormulaInternal::try_from(formula)?).map_err(|e| { + Status::invalid_argument(format!("failed to parse formula: {e}")) + })?, + ), }; Ok(scoring_query) @@ -664,8 +564,9 @@ impl From for grpc::query_shard_points::Query { ScoringQuery::OrderBy(order_by) => Self { score: Some(Score::OrderBy(grpc::OrderBy::from(order_by))), }, - // TODO(score boosting): Implement conversion - ScoringQuery::Formula(_formula) => todo!(), + ScoringQuery::Formula(parsed_formula) => Self { + score: Some(Score::Formula(grpc::Formula::from_parsed(parsed_formula))), + }, ScoringQuery::Sample(sample) => Self { score: Some(Score::Sample(api::grpc::qdrant::Sample::from(sample) as i32)), }, 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 3ece3bef18..922d90ccf5 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 @@ -66,6 +66,16 @@ pub enum VariableId { Condition(ConditionId), } +impl VariableId { + pub fn unparse(self) -> String { + match self { + VariableId::Score(index) => format!("${SCORE_KEYWORD}[{index}]"), + VariableId::Payload(path) => path.to_string(), + VariableId::Condition(_) => unreachable!("there are no defaults for conditions"), + } + } +} + impl ParsedExpression { /// Default value for division by zero const fn by_zero_default() -> ScoreType {