upstream-archive/qdrant-client
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant-client.git synced 2026-07-23 11:11:01 -05:00
Code Issues Packages Projects Releases Wiki Activity

Geo polygons support (#325)

Browse Source 
* regenerate client

* regen grpc client

* document build process

* work from Zein in https://github.com/qdrant/qdrant-client/pull/272

* update rest client with optional field

* code refactor and better naming

* remove shapely as a dependency

* fix typing

---------

Co-authored-by: zzzz-vincent <wenzishen.vincent@hotmail.com>
Co-authored-by: generall <andrey@vasnetsov.com>
This commit is contained in:
Arnaud Gourlay
2023-10-06 21:43:25 +02:00
committed by generall
parent d2a8b43a1b
commit 5b155f833f
13 changed files with 310 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@@ -11,3 +11,4 @@ htmlcov
dist
*.tar.gz
local_cache/*/*
.python-version

8
poetry.lock generated
View File

@@ -13,13 +13,13 @@ files = [
[[package]]
name = "annotated-types"
version = "0.5.0"
version = "0.6.0"
description = "Reusable constraint types to use with typing.Annotated"
optional = false
python-versions = ">=3.7"
python-versions = ">=3.8"
files = [
{file = "annotated_types-0.5.0-py3-none-any.whl", hash = "sha256:58da39888f92c276ad970249761ebea80ba544b77acddaa1a4d6cf78287d45fd"},
{file = "annotated_types-0.5.0.tar.gz", hash = "sha256:47cdc3490d9ac1506ce92c7aaa76c579dc3509ff11e098fc867e5130ab7be802"},
{file = "annotated_types-0.6.0-py3-none-any.whl", hash = "sha256:0641064de18ba7a25dee8f96403ebc39113d0cb953a01429249d5c7564666a43"},
{file = "annotated_types-0.6.0.tar.gz", hash = "sha256:563339e807e53ffd9c267e99fc6d9ea23eb8443c08f112651963e24e22f84a5d"},
]
[package.dependencies]

82
qdrant_client/grpc/points_pb2.py
View File

File diff suppressed because one or more lines are too long

1
qdrant_client/http/__init__.py
View File

@@ -4,6 +4,7 @@ from pydantic import BaseModel
from qdrant_client._pydantic_compat import update_forward_refs
from qdrant_client.http.api_client import ( # noqa F401
ApiClient as ApiClient,
AsyncApiClient as AsyncApiClient,
AsyncApis as AsyncApis,
SyncApis as SyncApis,
)

4
qdrant_client/http/api_client.py
View File

@@ -89,7 +89,7 @@ class ApiClient:
def send(self, request: Request, type_: Type[T]) -> T:
response = self.middleware(request, self.send_inner)
if response.status_code in [200, 201]:
if response.status_code in [200, 201, 202]:
try:
return parse_as_type(response.json(), type_)
except ValidationError as e:
@@ -161,7 +161,7 @@ class AsyncApiClient:
async def send(self, request: Request, type_: Type[T]) -> T:
response = await self.middleware(request, self.send_inner)
if response.status_code in [200, 201]:
if response.status_code in [200, 201, 202]:
try:
return parse_as_type(response.json(), type_)
except ValidationError as e:

28
qdrant_client/http/models/models.py
View File

@@ -432,6 +432,9 @@ class FieldCondition(BaseModel, extra="forbid"):
default=None, description="Check if points geo location lies in a given area"
)
geo_radius: Optional["GeoRadius"] = Field(default=None, description="Check if geo point is within a given radius")
geo_polygon: Optional["GeoPolygon"] = Field(
default=None, description="Check if geo point is within a given polygon"
)
values_count: Optional["ValuesCount"] = Field(default=None, description="Check number of values of the field")
@@ -462,6 +465,14 @@ class GeoBoundingBox(BaseModel, extra="forbid"):
)
class GeoLineString(BaseModel, extra="forbid"):
"""
Ordered sequence of GeoPoints representing the line
"""
points: List["GeoPoint"] = Field(..., description="Ordered sequence of GeoPoints representing the line")
class GeoPoint(BaseModel, extra="forbid"):
"""
Geo point payload schema
@@ -471,6 +482,21 @@ class GeoPoint(BaseModel, extra="forbid"):
lat: float = Field(..., description="Geo point payload schema")
class GeoPolygon(BaseModel, extra="forbid"):
"""
Geo filter request Matches coordinates inside the polygon, defined by `exterior` and `interiors`
"""
exterior: "GeoLineString" = Field(
...,
description="Geo filter request Matches coordinates inside the polygon, defined by `exterior` and `interiors`",
)
interiors: Optional[List["GeoLineString"]] = Field(
default=None,
description="Interior lines (if present) bound holes within the surface each GeoLineString must consist of a minimum of 4 points, and the first and last points must be the same.",
)
class GeoRadius(BaseModel, extra="forbid"):
"""
Geo filter request Matches coordinates inside the circle of `radius` and center with coordinates `center`
@@ -1618,7 +1644,7 @@ class UpdateCollection(BaseModel, extra="forbid"):
vectors: Optional["VectorsConfigDiff"] = Field(
default=None,
description="Vector data parameters to update. It is possible to provide one config for single vector mode and list of configs for multiple vectors mode.",
description="Map of vector data parameters to update for each named vector. To update parameters in a collection having a single unnamed vector, use an empty string as name.",
)
optimizers_config: Optional["OptimizersConfigDiff"] = Field(
default=None,

49
qdrant_client/local/geo.py
View File

@@ -1,4 +1,5 @@
from math import asin, cos, radians, sin, sqrt
from typing import List, Tuple
def geo_distance(lon1: float, lat1: float, lon2: float, lat2: float) -> float:
@@ -38,3 +39,51 @@ def test_geo_distance() -> None:
assert geo_distance(moscow["lon"], moscow["lat"], london["lon"], london["lat"]) < 2600 * 1000
assert geo_distance(moscow["lon"], moscow["lat"], berlin["lon"], berlin["lat"]) > 1600 * 1000
assert geo_distance(moscow["lon"], moscow["lat"], berlin["lon"], berlin["lat"]) < 1650 * 1000
def boolean_point_in_polygon(
point: Tuple[float, float],
exterior: List[Tuple[float, float]],
interiors: List[List[Tuple[float, float]]],
) -> bool:
inside_poly = False
if in_ring(point, exterior, True):
in_hole = False
k = 0
while k < len(interiors) and not in_hole:
if in_ring(point, interiors[k], False):
in_hole = True
k += 1
if not in_hole:
inside_poly = True
return inside_poly
def in_ring(
pt: Tuple[float, float], ring: List[Tuple[float, float]], ignore_boundary: bool
) -> bool:
is_inside = False
if ring[0][0] == ring[len(ring) - 1][0] and ring[0][1] == ring[len(ring) - 1][1]:
ring = ring[0 : len(ring) - 1]
j = len(ring) - 1
for i in range(0, len(ring)):
xi = ring[i][0]
yi = ring[i][1]
xj = ring[j][0]
yj = ring[j][1]
on_boundary = (
(pt[1] * (xi - xj) + yi * (xj - pt[0]) + yj * (pt[0] - xi) == 0)
and ((xi - pt[0]) * (xj - pt[0]) <= 0)
and ((yi - pt[1]) * (yj - pt[1]) <= 0)
)
if on_boundary:
return not ignore_boundary
intersect = ((yi > pt[1]) != (yj > pt[1])) and (
pt[0] < (xj - xi) * (pt[1] - yi) / (yj - yi) + xi
)
if intersect:
is_inside = not is_inside
j = i
return is_inside

22
qdrant_client/local/payload_filters.py
View File

@@ -3,7 +3,7 @@ from typing import Any, List, Optional
import numpy as np
from qdrant_client.http import models
from qdrant_client.local.geo import geo_distance
from qdrant_client.local.geo import boolean_point_in_polygon, geo_distance
from qdrant_client.local.payload_value_extractor import value_by_key
@@ -70,6 +70,22 @@ def check_geo_bounding_box(condition: models.GeoBoundingBox, values: Any) -> boo
return False
def check_geo_polygon(condition: models.GeoPolygon, values: Any) -> bool:
if isinstance(values, dict) and "lat" in values and "lon" in values:
lat = values["lat"]
lon = values["lon"]
exterior = [(point.lat, point.lon) for point in condition.exterior.points]
interiors = []
if condition.interiors is not None:
interiors = [
[(point.lat, point.lon) for point in interior.points]
for interior in condition.interiors
]
return boolean_point_in_polygon(point=(lat, lon), exterior=exterior, interiors=interiors)
return False
def check_range(condition: models.Range, value: Any) -> bool:
if not isinstance(value, (int, float)):
return False
@@ -141,6 +157,10 @@ def check_condition(
if condition.values_count is not None:
values = value_by_key(payload, condition.key, flat=False)
return check_values_count(condition.values_count, values)
if condition.geo_polygon is not None:
if values is None:
return False
return any(check_geo_polygon(condition.geo_polygon, v) for v in values)
elif isinstance(condition, models.NestedCondition):
values = value_by_key(payload, condition.nested.key)
if values is None:

47
qdrant_client/local/tests/test_payload_filters.py
View File

@@ -140,3 +140,50 @@ def test_nested_payload_filters():
res = check_filter(query, payload, 0)
assert res is False
def test_geo_polygon_filter_query():
payload = {
"location": [
{
"lon": 70.0,
"lat": 70.0,
},
]
}
query = models.Filter(
**{
"must": [
{
"key": "location",
"geo_polygon": {
"exterior": {
"points": [
{"lon": 55.455868, "lat": 55.495862},
{"lon": 86.455868, "lat": 55.495862},
{"lon": 86.455868, "lat": 86.495862},
{"lon": 55.455868, "lat": 86.495862},
{"lon": 55.455868, "lat": 55.495862},
]
},
},
}
]
}
)
res = check_filter(query, payload, 0)
assert res is True
payload = {
"location": [
{
"lon": 30.693738,
"lat": 30.502165,
},
]
}
res = check_filter(query, payload, 0)
assert res is False

10
qdrant_client/proto/collections.proto
View File

@@ -159,7 +159,7 @@ message OptimizersConfigDiff {
Do not create segments larger this size (in kilobytes).
Large segments might require disproportionately long indexation times,
therefore it makes sense to limit the size of segments.
If indexing speed is more important - make this parameter lower.
If search speed is more important - make this parameter higher.
Note: 1Kb = 1 vector of size 256
@@ -179,11 +179,11 @@ message OptimizersConfigDiff {
optional uint64 memmap_threshold = 5;
/*
Maximum size (in kilobytes) of vectors allowed for plain index, exceeding this threshold will enable vector indexing
Default value is 20,000, based on <https://github.com/google-research/google-research/blob/master/scann/docs/algorithms.md>.
To disable vector indexing, set to `0`.
Note: 1kB = 1 vector of size 256.
*/
optional uint64 indexing_threshold = 6;
@@ -412,7 +412,7 @@ message ShardTransferInfo {
}
message CollectionClusterInfoResponse {
uint64 peer_id = 1; // ID of this peer
uint64 peer_id = 1; // ID of this peer
uint64 shard_count = 2; // Total number of shards
repeated LocalShardInfo local_shards = 3; // Local shards
repeated RemoteShardInfo remote_shards = 4; // Remote shards

24
qdrant_client/proto/points.proto
View File

@@ -215,7 +215,7 @@ message SearchParams {
optional bool exact = 2;
/*
If set to true, search will ignore quantized vector data
If set to true, search will ignore quantized vector data
*/
optional QuantizationSearchParams quantization = 3;
/*
@@ -283,12 +283,12 @@ message ScrollPoints {
// How to use positive and negative vectors to find the results, default is `AverageVector`:
enum RecommendStrategy {
// Average positive and negative vectors and create a single query with the formula
// Average positive and negative vectors and create a single query with the formula
// `query = avg_pos + avg_pos - avg_neg`. Then performs normal search.
AverageVector = 0;
// Uses custom search objective. Each candidate is compared against all
// examples, its score is then chosen from the `max(max_pos_score, max_neg_score)`.
// Uses custom search objective. Each candidate is compared against all
// examples, its score is then chosen from the `max(max_pos_score, max_neg_score)`.
// If the `max_neg_score` is chosen then it is squared and negated.
BestScore = 1;
}
@@ -432,7 +432,7 @@ message GroupId {
message PointGroup {
GroupId id = 1; // Group id
repeated ScoredPoint hits = 2; // Points in the group
repeated ScoredPoint hits = 2; // Points in the group
RetrievedPoint lookup = 3; // Point(s) from the lookup collection that matches the group id
}
@@ -551,7 +551,7 @@ message FieldCondition {
GeoBoundingBox geo_bounding_box = 4; // Check if points geolocation lies in a given area
GeoRadius geo_radius = 5; // Check if geo point is within a given radius
ValuesCount values_count = 6; // Check number of values for a specific field
// GeoPolygon geo_polygon = 7; // Check if geo point is within a given polygon
GeoPolygon geo_polygon = 7; // Check if geo point is within a given polygon
}
message Match {
@@ -592,11 +592,15 @@ message GeoRadius {
float radius = 2; // In meters
}
message GeoLineString {
repeated GeoPoint points = 1; // Ordered sequence of GeoPoints representing the line
}
// For a valid GeoPolygon, both the exterior and interior GeoLineStrings must consist of a minimum of 4 points.
// Additionally, the first and last points of each GeoLineString must be the same.
message GeoPolygon {
// Ordered list of coordinates representing the vertices of a polygon.
// The minimum size is 4, and the first coordinate and the last coordinate
// should be the same to form a closed polygon.
repeated GeoPoint points = 1;
GeoLineString exterior = 1; // The exterior line bounds the surface
repeated GeoLineString interiors = 2; // Interior lines (if present) bound holes within the surface
}
message ValuesCount {

87
tests/congruence_tests/test_geo_polygon_filter_query.py Normal file
View File

@@ -0,0 +1,87 @@
import json
import random
from qdrant_client.http.models import models
from tests.congruence_tests.test_common import (
COLLECTION_NAME,
generate_fixtures,
init_client,
init_local,
init_remote,
)
def test_geo_polygon_filter_query():
# fix random seed
random.seed(42)
fixture_records = generate_fixtures(num=100)
local_client = init_local()
init_client(local_client, fixture_records)
remote_client = init_remote()
init_client(remote_client, fixture_records)
filter_ = models.Filter(
**{
"should": [
{
"key": "city.geo",
"geo_polygon": {
"exterior": {
"points": [
{"lon": -55.0, "lat": -55.0},
{"lon": 65.0, "lat": -55.0},
{"lon": 65.0, "lat": 65.0},
{"lon": 55.0, "lat": -65.0},
{"lon": -55.0, "lat": -55.0},
]
},
},
},
{
"key": "city.geo",
"geo_polygon": {
"exterior": {
"points": [
{"lon": 75.0, "lat": 75.0},
{"lon": 155.0, "lat": 75.0},
{"lon": 155.0, "lat": 85.0},
{"lon": 75.0, "lat": 85.0},
{"lon": 75.0, "lat": 75.0},
]
},
},
},
]
}
)
local_result, _next_page = local_client.scroll(
collection_name=COLLECTION_NAME,
scroll_filter=filter_,
limit=100,
with_payload=True,
)
remote_result, _next_page = remote_client.scroll(
collection_name=COLLECTION_NAME,
scroll_filter=filter_,
limit=100,
with_payload=True,
)
print("local_result:", len(local_result))
print("remote_result", len(remote_result))
assert len(local_result) == len(remote_result)
for local, remote in zip(local_result, remote_result):
if local.id != remote.id:
print(f"Local: {local.id}, Remote: {remote.id}")
print(f"Local:", json.dumps(local.payload["nested"]["array"], indent=2))
print(f"Remote:", json.dumps(remote.payload["nested"]["array"], indent=2))
assert False

6
tools/DEVELOPMENT.md
View File

@@ -12,6 +12,12 @@ For fixes:
---
* [ ] Create python virtual environment and install dependencies
* Install pyenv https://github.com/pyenv/pyenv#automatic-installer
* Install system pyenv dependencies https://github.com/pyenv/pyenv/wiki#suggested-build-environment
* `pyenv install 3.10.10` - install python 3.10.10
* `pyenv local 3.10.10` - set python version
* `pip install grpcio==1.48.2` - install grpcio
* `pip install grpcio-tools==1.48.2` - install grpcio-tools
* `pip install virtualenv` - install venv manager
* `virtualenv venv` - create virtual env
* `source venv/bin/activate` - enter venv