upstream-archive/qdrant
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant.git synced 2026-08-05 17:40:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

add formula to internal grpc api (#6057)

Browse Source
This commit is contained in:
Luis Cossío
2025-02-26 10:27:53 -03:00
committed by GitHub
parent 64ff5f0af5
commit 697c797fa4
7 changed files with 239 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
lib/api/src/grpc/conversions.rs
View File

@@ -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<HwMeasurementAcc> 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<segment::types::Condition>,
) -> 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),
}
}

2
lib/api/src/grpc/mod.rs
View File

@@ -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")
}

1
lib/api/src/grpc/proto/points_internal_service.proto
View File

@@ -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
}
}

5
lib/api/src/grpc/qdrant.rs
View File

@@ -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<query::Score>,
}
/// 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)]

108
lib/collection/src/operations/universal_query/formula.rs
View File

@@ -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<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)
}
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<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 From<rest::FormulaQuery> for FormulaInternal {
fn from(value: rest::FormulaQuery) -> Self {
let rest::FormulaQuery { formula, defaults } = value;

123
lib/collection/src/operations/universal_query/shard_query.rs
View File

@@ -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<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)
}
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<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;
@@ -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<ScoringQuery> 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)),
},

10
lib/segment/src/index/query_optimization/rescore_formula/parsed_formula.rs
View File

@@ -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 {