diff --git a/qdrant_client/conversions/conversion.py b/qdrant_client/conversions/conversion.py index a4873fcc..9dbc83a1 100644 --- a/qdrant_client/conversions/conversion.py +++ b/qdrant_client/conversions/conversion.py @@ -1,6 +1,7 @@ from datetime import date, datetime, timezone from typing import Any, Mapping, Optional, Sequence, Union, get_args +from google.protobuf.internal.containers import MessageMap from google.protobuf.json_format import MessageToDict from google.protobuf.timestamp_pb2 import Timestamp @@ -90,7 +91,7 @@ def payload_to_grpc(payload: dict[str, Any]) -> dict[str, Value]: return dict((key, json_to_value(val)) for key, val in payload.items()) -def grpc_to_payload(grpc_: dict[str, Value]) -> dict[str, Any]: +def grpc_to_payload(grpc_: MessageMap[str, Value]) -> dict[str, Any]: return dict((key, value_to_json(val)) for key, val in grpc_.items()) @@ -288,17 +289,16 @@ class GrpcToRest: name = model.WhichOneof("variant") if name is None: raise ValueError(f"invalid MaxOptimizationThreads model: {model}") # pragma: no cover - val = getattr(model, name) if name == "setting": - if val == grpc.MaxOptimizationThreads.Setting.Auto: + if model.setting == grpc.MaxOptimizationThreads.Setting.Auto: return rest.MaxOptimizationThreadsSetting.AUTO else: raise ValueError( f"invalid MaxOptimizationThreads model: {model}" ) # pragma: no cover elif name == "value": - return val + return model.value else: raise ValueError(f"invalid MaxOptimizationThreads model: {model}") # pragma: no cover @@ -311,7 +311,7 @@ class GrpcToRest: max_optimization_threads = cls.convert_max_optimization_threads( model.max_optimization_threads ) - if not isinstance(max_optimization_threads, int): + if max_optimization_threads == rest.MaxOptimizationThreadsSetting.AUTO: max_optimization_threads = None return rest.OptimizersConfig( @@ -1182,6 +1182,79 @@ class GrpcToRest: raise ValueError(f"invalid Sample model: {model}") # pragma: no cover + @classmethod + def convert_formula_query(cls, model: grpc.Formula) -> rest.FormulaQuery: + defaults = grpc_to_payload(model.defaults) + return rest.FormulaQuery( + formula=cls.convert_expression(model.expression), defaults=defaults + ) + + @classmethod + def convert_expression(cls, model: grpc.Expression) -> rest.Expression: + name = model.WhichOneof("variant") + if name is None: + raise ValueError(f"invalid Query model: {model}") # pragma: no cover + + if name == "constant": + return model.constant + if name == "variable": + return model.variable + if name == "condition": + return cls.convert_condition(model.condition) + if name == "sum": + return cls.convert_sum_expression(model.sum) + if name == "mult": + return cls.convert_mult_expression(model.mult) + if name == "div": + return cls.convert_div_expression(model.div) + if name == "abs": + return rest.AbsExpression(abs=cls.convert_expression(model.abs)) + if name == "neg": + return rest.NegExpression(neg=cls.convert_expression(model.neg)) + if name == "log10": + return rest.Log10Expression(log10=cls.convert_expression(model.log10)) + if name == "ln": + return rest.LnExpression(ln=cls.convert_expression(model.ln)) + if name == "sqrt": + return rest.SqrtExpression(sqrt=cls.convert_expression(model.sqrt)) + if name == "exp": + return rest.ExpExpression(exp=cls.convert_expression(model.exp)) + if name == "pow": + return cls.convert_pow_expression(model.pow) + if name == "geo_distance": + return cls.convert_geo_distance(model.geo_distance) + + raise ValueError(f"Unknown function name: {name}") + + @classmethod + def convert_sum_expression(cls, model: grpc.SumExpression) -> rest.SumExpression: + return rest.SumExpression(sum=[cls.convert_expression(expr) for expr in model.sum]) + + @classmethod + def convert_mult_expression(cls, model: grpc.MultExpression) -> rest.MultExpression: + return rest.MultExpression(mult=[cls.convert_expression(expr) for expr in model.mult]) + + @classmethod + def convert_div_expression(cls, model: grpc.DivExpression) -> rest.DivExpression: + left = cls.convert_expression(model.left) + right = cls.convert_expression(model.right) + by_zero_default = model.by_zero_default if model.HasField("by_zero_default") else None + params = rest.DivParams(left=left, right=right, by_zero_default=by_zero_default) + return rest.DivExpression(div=params) + + @classmethod + def convert_pow_expression(cls, model: grpc.PowExpression) -> rest.PowExpression: + base = cls.convert_expression(model.base) + exponent = cls.convert_expression(model.exponent) + params = rest.PowParams(base=base, exponent=exponent) + return rest.PowExpression(pow=params) + + @classmethod + def convert_geo_distance(cls, model: grpc.GeoDistance) -> rest.GeoDistance: + origin = cls.convert_geo_point(model.origin) + params = rest.GeoDistanceParams(origin=origin, to=model.to) + return rest.GeoDistance(geo_distance=params) + @classmethod def convert_query(cls, model: grpc.Query) -> rest.Query: name = model.WhichOneof("variant") @@ -1210,6 +1283,9 @@ class GrpcToRest: if name == "sample": return rest.SampleQuery(sample=cls.convert_sample(val)) + if name == "formula": + return cls.convert_formula_query(val) + raise ValueError(f"invalid Query model: {model}") # pragma: no cover @classmethod @@ -3367,8 +3443,80 @@ class RestToGrpc: if isinstance(model, rest.SampleQuery): return grpc.Query(sample=cls.convert_sample(model.sample)) + if isinstance(model, rest.FormulaQuery): + return grpc.Query(formula=cls.convert_formula_query(model)) + raise ValueError(f"invalid Query model: {model}") # pragma: no cover + @classmethod + def convert_formula_query(cls, model: rest.FormulaQuery) -> grpc.Formula: + defaults = payload_to_grpc(model.defaults) if model.defaults is not None else None + expression = cls.convert_expression(model.formula) + return grpc.Formula(defaults=defaults, expression=expression) + + @classmethod + def convert_expression(cls, model: rest.Expression) -> grpc.Expression: + if isinstance(model, float): + return grpc.Expression(constant=model) + if isinstance(model, str): + return grpc.Expression(variable=model) + if isinstance(model, get_args_subscribed(rest.Condition)): + return grpc.Expression(condition=cls.convert_condition(model)) + if isinstance(model, rest.NegExpression): + return grpc.Expression(neg=cls.convert_expression(model.neg)) + if isinstance(model, rest.SumExpression): + return grpc.Expression(sum=cls.convert_sum_expression(model)) + if isinstance(model, rest.MultExpression): + return grpc.Expression(mult=cls.convert_mult_expression(model)) + if isinstance(model, rest.DivExpression): + return grpc.Expression(div=cls.convert_div_expression(model)) + if isinstance(model, rest.PowExpression): + return grpc.Expression(pow=cls.convert_pow_expression(model)) + if isinstance(model, rest.Log10Expression): + return grpc.Expression(log10=cls.convert_expression(model.log10)) + if isinstance(model, rest.LnExpression): + return grpc.Expression(ln=cls.convert_expression(model.ln)) + if isinstance(model, rest.AbsExpression): + return grpc.Expression(abs=cls.convert_expression(model.abs)) + if isinstance(model, rest.SqrtExpression): + return grpc.Expression(sqrt=cls.convert_expression(model.sqrt)) + if isinstance(model, rest.ExpExpression): + return grpc.Expression(exp=cls.convert_expression(model.exp)) + if isinstance(model, rest.GeoDistance): + return grpc.Expression(geo_distance=cls.convert_geo_distance(model)) + + raise ValueError(f"invalid Expression model: {model}") # pragma: no cover + + @classmethod + def convert_sum_expression(cls, model: rest.SumExpression) -> grpc.SumExpression: + return grpc.SumExpression(sum=[cls.convert_expression(expr) for expr in model.sum]) + + @classmethod + def convert_mult_expression(cls, model: rest.MultExpression) -> grpc.MultExpression: + return grpc.MultExpression(mult=[cls.convert_expression(expr) for expr in model.mult]) + + @classmethod + def convert_div_expression(cls, model: rest.DivExpression) -> grpc.DivExpression: + return grpc.DivExpression( + left=cls.convert_expression(model.div.left), + right=cls.convert_expression(model.div.right), + by_zero_default=model.div.by_zero_default, + ) + + @classmethod + def convert_pow_expression(cls, model: rest.PowExpression) -> grpc.PowExpression: + return grpc.PowExpression( + base=cls.convert_expression(model.pow.base), + exponent=cls.convert_expression(model.pow.exponent), + ) + + @classmethod + def convert_geo_distance(cls, model: rest.GeoDistance) -> grpc.GeoDistance: + return grpc.GeoDistance( + origin=cls.convert_geo_point(model.geo_distance.origin), + to=model.geo_distance.to, + ) + @classmethod def convert_query_interface(cls, model: rest.QueryInterface) -> grpc.Query: if isinstance(model, get_args_subscribed(rest.VectorInput)): diff --git a/qdrant_client/hybrid/formula.py b/qdrant_client/hybrid/formula.py new file mode 100644 index 00000000..26ae5be6 --- /dev/null +++ b/qdrant_client/hybrid/formula.py @@ -0,0 +1,303 @@ +import warnings +from typing import Union, Any + +import pytest +import numpy as np + +from qdrant_client._pydantic_compat import construct +from qdrant_client.conversions.common_types import get_args_subscribed +from qdrant_client.http import models +from qdrant_client.local.geo import geo_distance +from qdrant_client.local.payload_filters import check_condition +from qdrant_client.local.payload_value_extractor import value_by_key + +DEFAULT_SCORE = np.float32(0.0) + + +def evaluate_expression( + expression: models.Expression, + point_id: models.ExtendedPointId, + scores: list[dict[models.ExtendedPointId, float]], + payload: models.Payload, + has_vector: dict[str, bool], + defaults: dict[str, Any], +) -> np.float32: + if isinstance(expression, (float, int)): # Constant + return np.float32(expression) + + elif isinstance(expression, str): # Variable + return evaluate_variable(expression, point_id, scores, payload, defaults) + + elif isinstance(expression, get_args_subscribed(models.Condition)): + if check_condition(expression, payload, point_id, has_vector): # type: ignore + return np.float32(1.0) + return np.float32(0.0) + + elif isinstance(expression, models.MultExpression): + factors: list[np.float32] = [] + + for expr in expression.mult: + factor = evaluate_expression(expr, point_id, scores, payload, has_vector, defaults) + # Return early if any factor is zero + if factor == np.float32(0.0): + return factor + + factors.append(factor) + + return np.prod(factors, dtype=np.float32) + + elif isinstance(expression, models.SumExpression): + return np.sum( + [ + evaluate_expression(expr, point_id, scores, payload, has_vector, defaults) + for expr in expression.sum + ], + dtype=np.float32, + ) + + elif isinstance(expression, models.NegExpression): + return -evaluate_expression( + expression.neg, point_id, scores, payload, has_vector, defaults + ) + + elif isinstance(expression, models.AbsExpression): + return abs( + evaluate_expression(expression.abs, point_id, scores, payload, has_vector, defaults) + ) + + elif isinstance(expression, models.DivExpression): + left = evaluate_expression( + expression.div.left, point_id, scores, payload, has_vector, defaults + ) + + if left == np.float32(0.0): + return left + + right = evaluate_expression( + expression.div.right, point_id, scores, payload, has_vector, defaults + ) + + if right == np.float32(0.0): + if expression.div.by_zero_default is not None: + return np.float32(expression.div.by_zero_default) + raise_non_finite_error(f"{left}/{right}") + + with np.errstate(invalid="ignore"): + result = left / right + if np.isfinite(result): + return np.float32(result) + + raise_non_finite_error(f"{left}/{right}") + + elif isinstance(expression, models.SqrtExpression): + value = evaluate_expression( + expression.sqrt, point_id, scores, payload, has_vector, defaults + ) + + with np.errstate(invalid="ignore"): + sqrt_value = np.sqrt(value, dtype=np.float32) + if np.isfinite(sqrt_value): + return np.float32(sqrt_value) + + raise_non_finite_error(f"√{value}") + + elif isinstance(expression, models.PowExpression): + base = evaluate_expression( + expression.pow.base, point_id, scores, payload, has_vector, defaults + ) + exponent = evaluate_expression( + expression.pow.exponent, point_id, scores, payload, has_vector, defaults + ) + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore") + power = np.power(base, exponent, dtype=np.float32) + if np.isfinite(power): + return np.float32(power) + + raise_non_finite_error(f"{base}^{exponent}") + + elif isinstance(expression, models.ExpExpression): + value = evaluate_expression( + expression.exp, point_id, scores, payload, has_vector, defaults + ) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore") + exp_value = np.exp(value, dtype=np.float32) + if np.isfinite(exp_value): + return exp_value + + raise_non_finite_error(f"exp({value})") + + elif isinstance(expression, models.Log10Expression): + value = evaluate_expression( + expression.log10, point_id, scores, payload, has_vector, defaults + ) + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore") + log_value = np.log10(value, dtype=np.float32) + if np.isfinite(log_value): + return log_value + + raise_non_finite_error(f"log10({value})") + + elif isinstance(expression, models.LnExpression): + value = evaluate_expression(expression.ln, point_id, scores, payload, has_vector, defaults) + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore") + ln_value = np.log(value, dtype=np.float32) + if np.isfinite(ln_value): + return ln_value + + raise_non_finite_error(f"ln({value})") + + elif isinstance(expression, models.GeoDistance): + origin = expression.geo_distance.origin + to = expression.geo_distance.to + + # Get value from payload + geo_value = try_extract_payload_value(to, payload, defaults) + + if isinstance(geo_value, dict): + # let this fail if it is not a valid geo point + destination = construct(models.GeoPoint, **geo_value) + return np.float32( + geo_distance(origin.lon, origin.lat, destination.lon, destination.lat) + ) + + raise ValueError( + f"Expected geo point for {to} in the payload and/or in the formula defaults." + ) + + raise ValueError(f"Unsupported expression type: {type(expression)}") + + +def try_extract_payload_value(key: str, payload: models.Payload, defaults: dict[str, Any]) -> Any: + # Get value from payload + value = value_by_key(payload, key) + + if value is None or len(value) == 0: + # Or from defaults + value = defaults.get(key, None) + # Consider it None if it is an empty list + if isinstance(value, list) and len(value) == 0: + value = None + + # Consider it a single value if it's a list with one element + if isinstance(value, list) and len(value) == 1: + return value[0] + + if value is None: + raise ValueError(f"No value found for {key} in the payload nor the formula defaults") + + return value + + +def evaluate_variable( + variable: str, + point_id: models.ExtendedPointId, + scores: list[dict[models.ExtendedPointId, float]], + payload: models.Payload, + defaults: dict[str, Any], +) -> np.float32: + var = parse_variable(variable) + if isinstance(var, str): + value = try_extract_payload_value(var, payload, defaults) + + if is_number(value): + return value + + raise ValueError( + f"Expected number value for {var} in the payload and/or in the formula defaults. Error: Value is not a number" + ) + + elif isinstance(var, int): + # Get score from scores + score = None + if var < len(scores): + score = scores[var].get(point_id, None) + if score is not None: + return np.float32(score) + + defined_default = defaults.get(variable, None) + if defined_default is not None: + return defined_default + + return DEFAULT_SCORE + + raise ValueError(f"Invalid variable type: {type(var)}") + + +def parse_variable(var: str) -> Union[str, int]: + # Try to parse score pattern + if not var.startswith("$score"): + # Treat as payload path + return var + + remaining = var.replace("$score", "", 1) + if remaining == "": + # end of string, default idx is 0 + return 0 + + # it must proceed with brackets + if not remaining.startswith("["): + raise ValueError(f"Invalid score pattern: {var}") + + remaining = remaining.replace("[", "", 1) + bracket_end = remaining.find("]") + if bracket_end == -1: + raise ValueError(f"Invalid score pattern: {var}") + + # try parsing the content in between brackets as integer + try: + idx = int(remaining[:bracket_end]) + except ValueError: + raise ValueError(f"Invalid score pattern: {var}") + + # make sure the string ends after the closing bracket + if len(remaining) > bracket_end + 1: + raise ValueError(f"Invalid score pattern: {var}") + + return idx + + +def raise_non_finite_error(expression: str) -> None: + raise ValueError(f"The expression {expression} produced a non-finite number") + + +def is_number(value: Any) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) + + +def test_parsing_variable() -> None: + assert parse_variable("$score") == 0 + assert parse_variable("$score[0]") == 0 + assert parse_variable("$score[1]") == 1 + assert parse_variable("$score[2]") == 2 + + try: + parse_variable("$score[invalid]") + assert False + except ValueError as e: + assert str(e) == "Invalid score pattern: $score[invalid]" + + try: + parse_variable("$score[10].other") + assert False + except ValueError as e: + assert str(e) == "Invalid score pattern: $score[10].other" + + +@pytest.mark.parametrize( + "payload_value, expected", [(1.2, 1.2), ([1.2], 1.2), ([1.2, 2.3], [1.2, 2.3])] +) +def test_try_extract_payload_value(payload_value: Any, expected: Any) -> None: + empty_defaults: dict[str, Any] = {} + payload = {"key": payload_value} + assert try_extract_payload_value("key", payload, empty_defaults) == expected + + defaults = {"key": payload_value} + empty_payload: dict[str, Any] = {} + assert try_extract_payload_value("key", empty_payload, defaults) == expected diff --git a/qdrant_client/local/geo.py b/qdrant_client/local/geo.py index 62a655f8..6a36a3cf 100644 --- a/qdrant_client/local/geo.py +++ b/qdrant_client/local/geo.py @@ -1,5 +1,8 @@ from math import asin, cos, radians, sin, sqrt +# Radius of earth in meters, [as recommended by the IUGG](ftp://athena.fsv.cvut.cz/ZFG/grs80-Moritz.pdf) +MEAN_EARTH_RADIUS = 6371008.8 + def geo_distance(lon1: float, lat1: float, lon2: float, lat2: float) -> float: """ @@ -22,9 +25,8 @@ def geo_distance(lon1: float, lat1: float, lon2: float, lat2: float) -> float: dlat = lat2 - lat1 a = sin(dlat / 2) ** 2 + cos(lat1) * cos(lat2) * sin(dlon / 2) ** 2 c = 2 * asin(sqrt(a)) - # Radius of earth in kilometers is 6371 - km = 6371 * c - return km * 1000 + + return MEAN_EARTH_RADIUS * c def test_geo_distance() -> None: diff --git a/qdrant_client/local/local_collection.py b/qdrant_client/local/local_collection.py index 9b38b13f..b403a74d 100644 --- a/qdrant_client/local/local_collection.py +++ b/qdrant_client/local/local_collection.py @@ -14,7 +14,7 @@ from copy import deepcopy import numpy as np -from qdrant_client import grpc as grpc +from qdrant_client import grpc as grpc, hybrid from qdrant_client.common.client_warnings import show_warning_once from qdrant_client._pydantic_compat import construct, to_jsonable_python as _to_jsonable_python from qdrant_client.conversions import common_types as types @@ -23,6 +23,7 @@ from qdrant_client.conversions.conversion import GrpcToRest from qdrant_client.http import models from qdrant_client.http.models import ScoredPoint from qdrant_client.http.models.models import Distance, ExtendedPointId, SparseVector, OrderValue +from qdrant_client.hybrid.formula import evaluate_expression from qdrant_client.hybrid.fusion import reciprocal_rank_fusion, distribution_based_score_fusion from qdrant_client.local.distances import ( ContextPair, @@ -515,6 +516,14 @@ class LocalCollection: return mask + def _calculate_has_vector(self, internal_id: int) -> dict[str, bool]: + has_vector: dict[str, bool] = {} + for vector_name, deleted in self.deleted_per_vector.items(): + if not deleted[internal_id]: + has_vector[vector_name] = True + + return has_vector + def search( self, query_vector: Union[ @@ -786,8 +795,20 @@ class LocalCollection: scored.vector = fetched.vector return fused[offset:] + + elif isinstance(query, models.FormulaQuery): + # Re-score with formula + rescored = self._rescore_with_formula( + query=query, + prefetches_results=sources, + limit=limit + offset, + with_payload=with_payload, + with_vectors=with_vectors, + ) + + return rescored[offset:] else: - # Re-score + # Re-score with vector sources_ids = set() for source in sources: for point in source: @@ -903,6 +924,8 @@ class LocalCollection: raise ValueError(f"Unknown Sample variant: {query.sample}") elif isinstance(query, models.FusionQuery): raise AssertionError("Cannot perform fusion without prefetches") + elif isinstance(query, models.FormulaQuery): + raise AssertionError("Cannot perform formula without prefetches") else: # most likely a VectorInput, delegate to search return self.search( @@ -1971,6 +1994,54 @@ class LocalCollection: return result + def _rescore_with_formula( + self, + query: models.FormulaQuery, + prefetches_results: list[list[models.ScoredPoint]], + limit: int, + with_payload: Union[bool, Sequence[str], types.PayloadSelector], + with_vectors: Union[bool, Sequence[str]], + ) -> list[models.ScoredPoint]: + # collect prefetches in vec of dicts for faster lookup + prefetches_scores = [ + dict((point.id, point.score) for point in prefetch) for prefetch in prefetches_results + ] + + defaults = query.defaults or {} + + points_to_rescore: set[models.ExtendedPointId] = set() + for prefetch in prefetches_results: + for point in prefetch: + points_to_rescore.add(point.id) + + # Evaluate formula for each point + rescored: list[models.ScoredPoint] = [] + for point_id in points_to_rescore: + internal_id = self.ids[point_id] + payload = self._get_payload(internal_id, True) or {} + has_vector = self._calculate_has_vector(internal_id) + score = evaluate_expression( + expression=query.formula, + point_id=point_id, + scores=prefetches_scores, + payload=payload, + has_vector=has_vector, + defaults=defaults, + ) + point = construct( + models.ScoredPoint, + id=point_id, + score=float(score), + version=0, + payload=self._get_payload(internal_id, with_payload), + vector=self._get_vectors(internal_id, with_vectors), + ) + rescored.append(point) + + rescored.sort(key=lambda x: x.score, reverse=True) + + return rescored[:limit] + def _update_point(self, point: models.PointStruct) -> None: idx = self.ids[point.id] self.payload[idx] = deepcopy( diff --git a/qdrant_client/local/payload_filters.py b/qdrant_client/local/payload_filters.py index 4a8682e4..e41020e7 100644 --- a/qdrant_client/local/payload_filters.py +++ b/qdrant_client/local/payload_filters.py @@ -164,7 +164,7 @@ def check_nested_filter(nested_filter: models.Filter, values: list[Any]) -> bool def check_condition( condition: models.Condition, - payload: dict, + payload: dict[str, Any], point_id: models.ExtendedPointId, has_vector: Dict[str, bool], ) -> bool: diff --git a/qdrant_client/local/payload_value_extractor.py b/qdrant_client/local/payload_value_extractor.py index 6e650715..d343a709 100644 --- a/qdrant_client/local/payload_value_extractor.py +++ b/qdrant_client/local/payload_value_extractor.py @@ -8,7 +8,7 @@ from qdrant_client.local.json_path_parser import ( ) -def value_by_key(payload: dict, key: str, flat: bool = True) -> Optional[list[Any]]: +def value_by_key(payload: dict[str, Any], key: str, flat: bool = True) -> Optional[list[Any]]: """ Get value from payload by key. Args: diff --git a/tests/congruence_tests/test_query.py b/tests/congruence_tests/test_query.py index dbe1159d..18900718 100644 --- a/tests/congruence_tests/test_query.py +++ b/tests/congruence_tests/test_query.py @@ -1,5 +1,6 @@ -from typing import Tuple, Callable, Any +from typing import Tuple, Callable, Any, Union +from grpc import RpcError import numpy as np import pytest @@ -25,6 +26,7 @@ from tests.congruence_tests.test_common import ( generate_multivector_fixtures, multi_vector_config, ) +from tests.fixtures.expressions import one_random_expression_please from tests.fixtures.filters import one_random_filter_please from tests.fixtures.points import ( generate_random_sparse_vector, @@ -711,6 +713,41 @@ class TestSimpleSearcher: return result + def score_boosting( + self, client: QdrantBase, formula: models.FormulaQuery, point_id: int + ) -> Union[models.QueryResponse, str]: + def comparable_error(exception: Exception): + non_finite_message = "produced a non-finite number" + unexpected_type_message = "in the payload and/or in the formula defaults" + + if non_finite_message in str(exception): + return non_finite_message + elif unexpected_type_message in str(exception): + return unexpected_type_message + raise exception + + prefetch = models.Prefetch( + filter=models.Filter(must=[models.HasIdCondition(has_id=[point_id])]), + limit=100, + using="text", + ) + + try: + result = client.query_points( + collection_name=COLLECTION_NAME, + prefetch=prefetch, + query=formula, + limit=100, + ) + except ValueError as e: # local mode error + return comparable_error(e) + except UnexpectedResponse as e: # rest error + return comparable_error(e) + except RpcError as e: # grpc error + return comparable_error(e) + + return result + def group_by_keys(): return ["maybe", "rand_digit", "two_words", "city.name", "maybe_null", "id"] @@ -1493,3 +1530,42 @@ def test_random_sampling(): local_client, http_client, grpc_client = init_clients(fixture_points) compare_clients_results(local_client, http_client, grpc_client, searcher.random_query) + + +def test_formula_query(): + points_count = 100 + fixture_points = generate_fixtures(points_count) + + searcher = TestSimpleSearcher() + + local_client, http_client, grpc_client = init_clients(fixture_points) + + defaults = { + "rand_digit": 5, + "maybe_null": None, + "mixed_type": 0.3, + "city.geo": {"lon": 0.4, "lat": 0.5}, + } + + for _ in range(50): + formula = models.FormulaQuery( + formula=one_random_expression_please(max_depth=2), defaults=defaults + ) + + # We need to score point by point to make sure that the errors that come up correspond to the same point. + # + # Otherwise, we can have discrepancy where one point produced one error, + # and another caused a different error in the other client + for point_id in range(points_count): + try: + compare_clients_results( + local_client, + http_client, + grpc_client, + searcher.score_boosting, + formula=formula, + point_id=point_id, + ) + except Exception as e: + print(f"\nFailed with formula {formula} on point {fixture_points[point_id]}") + raise e diff --git a/tests/conversions/fixtures.py b/tests/conversions/fixtures.py index c5617c52..f5a1f95b 100644 --- a/tests/conversions/fixtures.py +++ b/tests/conversions/fixtures.py @@ -1346,6 +1346,54 @@ context_input_pair = grpc.ContextInputPair( context_input = grpc.ContextInput(pairs=[context_input_pair]) discover_input = grpc.DiscoverInput(target=vector_input_dense, context=context_input) +formula_defaults = payload_to_grpc( + { + "some_string": "some_value", + "some_number": 42.666, + "some_boolean": True, + "some_array": [1, 2, 3], + "some_object": {"key": "value"}, + "some_null": None, + } +) + +expression = grpc.Expression( + sum=grpc.SumExpression( + sum=[ + grpc.Expression(mult=grpc.MultExpression(mult=[grpc.Expression(constant=0.1)])), + grpc.Expression(variable="some_variable"), + grpc.Expression(condition=grpc.Condition(field=field_condition_match)), + grpc.Expression(geo_distance=grpc.GeoDistance(origin=geo_point, to="my_field")), + grpc.Expression(neg=grpc.Expression(constant=1.0)), + grpc.Expression(abs=grpc.Expression(variable="my_variable")), + grpc.Expression(sqrt=grpc.Expression(constant=4.0)), + grpc.Expression( + pow=grpc.PowExpression( + base=grpc.Expression(constant=2.0), exponent=grpc.Expression(constant=3.0) + ) + ), + grpc.Expression(exp=grpc.Expression(constant=1.0)), + grpc.Expression(log10=grpc.Expression(constant=100.0)), + grpc.Expression(ln=grpc.Expression(constant=2.718)), + grpc.Expression( + div=grpc.DivExpression( + left=grpc.Expression(constant=10.0), + right=grpc.Expression(constant=2.0), + ) + ), + grpc.Expression( + div=grpc.DivExpression( + left=grpc.Expression(constant=10.0), + right=grpc.Expression(constant=2.0), + by_zero_default=0.0, + ) + ), + ] + ) +) + +formula = grpc.Formula(defaults=formula_defaults, expression=expression) + query_nearest = grpc.Query(nearest=vector_input_sparse) query_recommend = grpc.Query(recommend=recommend_input) query_recommend_id = grpc.Query(recommend=recommend_input_strategy) @@ -1356,6 +1404,7 @@ query_order_by = grpc.Query(order_by=order_by) query_fusion = grpc.Query(fusion=grpc.Fusion.RRF) query_fusion_dbsf = grpc.Query(fusion=grpc.Fusion.DBSF) query_sample = grpc.Query(sample=grpc.Sample.Random) +query_formula = grpc.Query(formula=formula) deep_prefetch_query = grpc.PrefetchQuery(query=query_recommend) prefetch_query = grpc.PrefetchQuery( @@ -1567,6 +1616,7 @@ fixtures = { query_order_by, query_fusion, query_recommend_id, + query_formula, ], "FacetValueHit": [facet_string_hit, facet_integer_hit], "PrefetchQuery": [deep_prefetch_query, prefetch_query, prefetch_full_query, prefetch_many], diff --git a/tests/fixtures/expressions.py b/tests/fixtures/expressions.py new file mode 100644 index 00000000..937e96b1 --- /dev/null +++ b/tests/fixtures/expressions.py @@ -0,0 +1,96 @@ +import random + +import qdrant_client.models as models +from tests.fixtures.filters import one_random_condition_please + + +def one_random_expression_please(current_depth: int = 0, max_depth: int = 5) -> models.Expression: + """Generate a random expression for testing purposes. + + Args: + current_depth: Current depth of expression nesting + max_depth: Maximum allowed depth for nested expressions + + Returns: + A random Expression object + """ + + # Choose a random expression type with possible nesting + expression_choices = [ + # Terminal expressions (no nesting) with higher probability at deeper levels + lambda: round(random.uniform(-10.0, 10.0), 5), # Constant + lambda: random.choice( # Variable + [ + "rand_number", + "rand_digit", + "rand_signed_int", + "nested.rand_digit", + "$score", + "$score[0]", + "mixed_type", + ] + ), + lambda: one_random_condition_please(), # Condition + # Nested expressions + lambda: models.MultExpression( + mult=[ + one_random_expression_please(current_depth + 1, max_depth) + for _ in range(random.randint(2, 3)) + ] + ), + lambda: models.SumExpression( + sum=[ + one_random_expression_please(current_depth + 1, max_depth) + for _ in range(random.randint(2, 3)) + ] + ), + lambda: models.NegExpression( + neg=one_random_expression_please(current_depth + 1, max_depth) + ), + lambda: models.AbsExpression( + abs=one_random_expression_please(current_depth + 1, max_depth) + ), + lambda: models.DivExpression( + div=models.DivParams( + left=one_random_expression_please(current_depth + 1, max_depth), + right=one_random_expression_please(current_depth + 1, max_depth), + by_zero_default=round(random.uniform(0, 1), 5) if random.random() < 0.5 else None, + ) + ), + lambda: models.SqrtExpression( + sqrt=one_random_expression_please(current_depth + 1, max_depth) + ), + lambda: models.PowExpression( + pow=models.PowParams( + base=one_random_expression_please(current_depth + 1, max_depth), + exponent=one_random_expression_please(current_depth + 1, max_depth), + ) + ), + lambda: models.ExpExpression( + exp=one_random_expression_please(current_depth + 1, max_depth) + ), + lambda: models.Log10Expression( + log10=one_random_expression_please(current_depth + 1, max_depth) + ), + lambda: models.LnExpression(ln=one_random_expression_please(current_depth + 1, max_depth)), + # GeoDistance is a special case - needs specific structure + lambda: models.GeoDistance( + geo_distance=models.GeoDistanceParams( + origin=models.GeoPoint( + lon=round(random.uniform(-180, 180), 5), lat=round(random.uniform(-80, 80), 5) + ), + to="city.geo", # Using a field that would contain geo coordinates + ) + ), + ] + # If we've reached max depth, return a terminal expression (no nesting) + if current_depth >= max_depth: # Limit nesting depth + # Return a simple expression at max depth + return random.choice(expression_choices[:3])() + + # Give higher weight to terminal expressions at deeper levels + if current_depth > 2: + # Add more terminal expressions to increase their probability + expression_choices = expression_choices[:3] * 3 + expression_choices[3:] + + return random.choice(expression_choices)() diff --git a/tests/fixtures/payload.py b/tests/fixtures/payload.py index 23217ad6..c431504e 100644 --- a/tests/fixtures/payload.py +++ b/tests/fixtures/payload.py @@ -219,6 +219,18 @@ def one_random_payload_please(idx: int) -> dict[str, Any]: "city": random_city(), "rand_tuple": tuple(random.randint(0, 100) for _ in range(random.randint(1, 5))), "rand_bool": random.random() < 0.2, + "mixed_type": random.choice( + [ + None, + random_real_word(), + random.randint(0, 9), # int + [random.randint(0, 9) for _ in range(3)], # list of ints + round(random.random(), 5), # float + [round(random.random(), 5) for _ in range(3)], # list of floats + random.random() < 0.5, # bool + {"key": "value"}, + ] + ), } if random.random() < 0.5: